optimum-rbln 0.7.4a1__py3-none-any.whl → 0.7.4a2__py3-none-any.whl

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.4a1'
+__version__ = version = '0.7.4a2'
 __version_tuple__ = version_tuple = (0, 7, 4)

optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -578,11 +578,41 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         nbits_per_param: int,
         n_model_params: int,
     ) -> int:
+        """
+        We are finding max_n_blocks(x) that satisfies the following equation:
+
+        available_dram - kernel_size - buffer
+            - num_layers * 2 * tensor_parallel_size
+            * align_2MB(
+                x
+                * block_size
+                * align_64(head_dim)
+                * math.ceil(num_key_value_heads / tensor_parallel_size)
+                * 2
+            ) > 0
+
+        This inequality can be rewritten as follows:
+
+        a - c * align_2MB(b * x) > 0
+        where
+           a = available_dram - kernel_size - buffer
+           b = block_size * align_64(head_dim) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+           c = num_layers * 2 * tensor_parallel_size
+
+        We can rewrite the inequality as follows:
+        k > align_2MB(b*x)
+        where
+           k = a / c
+
+        After that, we can derive the following equation:
+        x = floor(2**21 / b * floor((k - 1) / 2**21))
+        """
+
         def align(x: int, nbytes: int) -> int:
             return int(math.ceil(x / nbytes) * nbytes)
 
         def align_2MB(x: int) -> int:
-            return align(x, 2 * 1024 * 1024)
+            return align(x, 2**21)
 
         num_attention_heads = getattr(config, "n_head", None) or getattr(config, "num_attention_heads")
         num_layers = getattr(config, "n_layer", None) or getattr(config, "num_hidden_layers")
@@ -612,27 +642,16 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         available_dram -= kernel_size
 
         # TODO: Accurate buffer estimation
-        buffer = 2**30  # 1GB Buffer
-        if tensor_parallel_size <= 4:
-            buffer /= 4
-
+        buffer_per_core = 2**29  # 500MB per npu
+        buffer = buffer_per_core * tensor_parallel_size
         available_dram -= buffer
 
-        # Estimate nbytes per a single kvcache block
-        nbytes_per_block = (
-            align_2MB(
-                kvcache_block_size
-                * head_dim
-                * math.ceil(num_key_value_heads / tensor_parallel_size)  # Shard
-                * 2  # (fp16)
-            )
-            * num_layers
-            * 2  # (k, v)
-            * tensor_parallel_size
-        )
-        n_blocks = available_dram // nbytes_per_block
+        b = kvcache_block_size * align(head_dim, 64) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+        c = num_layers * 2 * tensor_parallel_size
+        k = available_dram / c
+        max_n_blocks = math.floor(2**21 / b * math.floor((k - 1) / 2**21))
 
-        return n_blocks, nbytes_per_block
+        return max_n_blocks
 
     @classmethod
     def _get_rbln_config(
@@ -689,7 +708,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         rbln_kvcache_num_blocks = (rbln_max_seq_len // rbln_kvcache_block_size) * rbln_batch_size
         if rbln_attn_impl == "flash_attn":
-            max_num_blocks, _ = cls.get_maximum_num_blocks(
+            max_num_blocks = cls.get_maximum_num_blocks(
                 config=model_config,
                 tensor_parallel_size=rbln_kwargs.get("tensor_parallel_size", 1),
                 kvcache_block_size=rbln_kvcache_block_size,

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -36,19 +36,50 @@ from .t5_architecture import T5Wrapper
 logger = get_logger()
 
 if TYPE_CHECKING:
+    from rebel import Runtime
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
 
 
 class RBLNRuntimeModel(RBLNPytorchRuntime):
+    def __init__(
+        self,
+        runtime: "Runtime",
+        max_seq_len: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(runtime, **kwargs)
+        self.max_seq_len = max_seq_len
+
+    def _prepare_inputs(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.LongTensor,
+    ):
+        input_len = input_ids.shape[-1]
+        pad_len = None
+        if input_len > self.max_seq_len:
+            raise ValueError(f"Error input_len({input_len}) exceed max_seq_len({self.max_seq_len}).")
+        elif input_len < self.max_seq_len and input_len > 0:
+            pad_len = self.max_seq_len - input_len
+            logger.warning(
+                f"Warning: The input was padded with {pad_len} tokens to meet the compiled model's requirements. "
+                "For optimal performance, consider recompiling with a shorter 'rbln_max_seq_len'."
+            )
+            input_ids = torch.nn.functional.pad(input_ids, (0, pad_len))
+            attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_len), value=0)
+
+        return input_ids, attention_mask, pad_len
+
     def forward(
         self,
         input_ids: torch.LongTensor,
-        attention_mask: torch.FloatTensor,
+        attention_mask: torch.LongTensor,
         head_mask: torch.FloatTensor,
         inputs_embeds: torch.FloatTensor,
         **kwargs,
     ):
-        return super().forward(
+        input_ids, attention_mask, pad_len = self._prepare_inputs(input_ids, attention_mask)
+        logits = super().forward(
             input_ids,
             attention_mask,
             head_mask,
@@ -56,6 +87,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             **kwargs,
         )
 
+        return logits[:, :-pad_len, :] if pad_len is not None else logits
+
 
 class T5EncoderWrapper(torch.nn.Module):
     def __init__(self, model: "T5EncoderModel") -> None:
@@ -72,7 +105,8 @@ class RBLNT5EncoderModel(RBLNModel):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
     def __post_init__(self, **kwargs):
-        self.model = RBLNRuntimeModel(runtime=self.model[0])
+        max_seq_len = self.rbln_config.model_cfg["max_seq_len"]
+        self.model = RBLNRuntimeModel(runtime=self.model[0], max_seq_len=max_seq_len)
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):

optimum/rbln/transformers/models/whisper/whisper_architecture.py

@@ -190,11 +190,11 @@ class WhisperDecoder(nn.Module):
         all_hiddens = []
         for i in range(inputs_embeds.shape[0]):
             position_id = cache_position[i]
-            position = self.embed_positions(input_ids, position_ids=position_id)
+            position = self.embed_positions.weight[position_id]
             batch_hidden = position + inputs_embeds[i]
             all_hiddens.append(batch_hidden)
 
-        hidden_states = torch.stack(all_hiddens, dim=0)
+        hidden_states = torch.cat(all_hiddens, dim=0).unsqueeze(1)
 
         # prepare attn mask (normal attention - masked)
         if attention_mask is not None:

{optimum_rbln-0.7.4a1.dist-info → optimum_rbln-0.7.4a2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.4a1
+Version: 0.7.4a2
 Summary: Optimum RBLN is the interface between the Hugging Face Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.7.4a1.dist-info → optimum_rbln-0.7.4a2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 optimum/rbln/__init__.py,sha256=qW45z47BiNLTDtRFEhVEzr4THNFX0ygqCbdNKqI0biI,6992
-optimum/rbln/__version__.py,sha256=KifVR95YmJmHh5f74wGiEAzd-c6ElHQ3XFHbY8VRp14,513
+optimum/rbln/__version__.py,sha256=7nOn__ePBcQxiENj-KnOIjfCYYWSP4QlkMW59HLjtRs,513
 optimum/rbln/modeling.py,sha256=nJsAs5zs--VVOYGFjYNpqfxYIemJIK4Lr0WEzlDLdP0,8390
 optimum/rbln/modeling_base.py,sha256=dNCL-BhrWCpuOVkZaj8-MW567Tf4lLo3p3Z3ldjWJfU,21779
 optimum/rbln/modeling_config.py,sha256=7104bxmrvKW4Q6XTruQayiIGl8GHDFmPkJ3cknMIInE,11335
@@ -63,7 +63,7 @@ optimum/rbln/transformers/models/clip/__init__.py,sha256=H9vuBwrmFO0-CqZhXUrKF-u
 optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=NiSm7bHs4SReHDUr53BBWSX0Y8bkKOeUSpsBDrp8YDw,6628
 optimum/rbln/transformers/models/decoderonly/__init__.py,sha256=pDogsdpJKKB5rqnVFrRjwfhUvOSV-jZ3oARMsqSvOOQ,665
 optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=m93-qKN7NMw3i0XDmFmttmRIRK4np_fWtLFlBb2RFgU,41351
-optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=qeZWdfLU0gCssxBODJsjQWMjfQWxK9vgC2Xt9eA5j4I,39147
+optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=qNnWSe3p2LDwNJ6utrilsqid-rQ8YLloqYkSOZamvhs,39918
 optimum/rbln/transformers/models/dpt/__init__.py,sha256=gP1tkR3XMNlHq1GT87ugIVvb2o_1eAUg1JaniXjy1Lw,651
 optimum/rbln/transformers/models/dpt/modeling_dpt.py,sha256=ZsS2SOiqcA4azULB-WFEMQZbgIoOyVUKqVKqrw_tWzA,3430
 optimum/rbln/transformers/models/exaone/__init__.py,sha256=zYH_5tVa8-juEdsOIky7I33WSC3Zuhoq1upI0OHYeVw,859
@@ -96,7 +96,7 @@ optimum/rbln/transformers/models/seq2seq/__init__.py,sha256=EmEMV4rOYqKyruX85d0f
 optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py,sha256=XcZb57v42wju1qOJ1AKqmtJXcmz6MEWaJZ8jyzaEiTw,17701
 optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py,sha256=tvzacIZam1sIr_1BvvZ_fDr8u5dXAiYiynFdX9tArtY,18877
 optimum/rbln/transformers/models/t5/__init__.py,sha256=1skR1RmnG62WTAP3-F5P1x-V_ReFhMyirH3u56vWwvc,675
-optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=-fG-h0wwsfjZ3par0QHbXKA7hbvw_lPJOIf8iXQDOfM,8082
+optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=Gyq5aAfkl4hBbLiR0114nDxLBs5P6YTw7hCnyuDyRrM,9494
 optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=Ups6drBbYe4wEAiBLcBIyO9wqrIQbvOPFR_ybbAgR8c,9722
 optimum/rbln/transformers/models/time_series_transformers/__init__.py,sha256=RL4SO8tKEd4wQrzyU4Nv4-hhITKPhblUsBd3anXNkA8,1079
 optimum/rbln/transformers/models/time_series_transformers/modeling_time_series_transformers.py,sha256=1Ippt0Rmt2TxJ5X4-4tlALQOkKmOfMaTrbOLWIUIKWw,16614
@@ -106,7 +106,7 @@ optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py,sha256=JYJmV52j6c
 optimum/rbln/transformers/models/whisper/__init__.py,sha256=ktnNe5ri3ycCWZ_W_voFB9y9-vgGgxS1X9s8LBRZmWc,665
 optimum/rbln/transformers/models/whisper/generation_whisper.py,sha256=GIHTca3b1VtW81kp7BzKQ7f77c2t9OsEsbZetripgDo,4582
 optimum/rbln/transformers/models/whisper/modeling_whisper.py,sha256=GegyAi3a8fF0psdYsffTQ1pC4KAUqE7WYLj4ZqObWXI,18184
-optimum/rbln/transformers/models/whisper/whisper_architecture.py,sha256=k_aDk2B58IxQimf6yW36Wgc0uw5PqB85Or8ie_6ZZ70,14205
+optimum/rbln/transformers/models/whisper/whisper_architecture.py,sha256=DS9AQYhNkaR7sUz_loee-fFtCCYy1BUsx7_dX_o1Le8,14199
 optimum/rbln/transformers/models/xlm_roberta/__init__.py,sha256=fC7iNcdxBZ_6eOF2snStmf8r2M3c8O_-XcXnQEaHQCE,653
 optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py,sha256=8YNLz0bc5ze-QuU8rN-QhUfGzlSUs3iMJiWTxO3o6AM,4366
 optimum/rbln/transformers/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -120,7 +120,7 @@ optimum/rbln/utils/model_utils.py,sha256=DfD_Z2qvZHqcddXqnzTM1AN8khanj3-DXK2lJvV
 optimum/rbln/utils/runtime_utils.py,sha256=5-DYniyP59nx-mrrbi7AqA77L85b4Cm5oLpaxidSyss,3699
 optimum/rbln/utils/save_utils.py,sha256=hG5uOtYmecSXZuGTvCXsTM-SiyZpr5q3InUGCCq_jzQ,3619
 optimum/rbln/utils/submodule.py,sha256=oZoGrItB8WqY4i-K9WJPlLlcLohc1YGB9OHB8_XZw3A,4071
-optimum_rbln-0.7.4a1.dist-info/METADATA,sha256=dMl4yloIz6iqjC2SN8CE1rVP9Kftw50Z01zocntnguE,5300
-optimum_rbln-0.7.4a1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.7.4a1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.7.4a1.dist-info/RECORD,,
+optimum_rbln-0.7.4a2.dist-info/METADATA,sha256=Pl6SOVN73gxS7Po-R5hMdTI_2izOrjWqyb9FhaDnr-A,5300
+optimum_rbln-0.7.4a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.7.4a2.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.7.4a2.dist-info/RECORD,,