optimum-rbln 0.1.8__py3-none-any.whl → 0.1.11__py3-none-any.whl

optimum/rbln/transformers/models/dpt/modeling_dpt.py

@@ -22,13 +22,13 @@
 # from Rebellions Inc.
 
 import logging
-from typing import TYPE_CHECKING, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Union
 
 from transformers import AutoModelForDepthEstimation
 from transformers.modeling_outputs import DepthEstimatorOutput
 
 from ....modeling_base import RBLNModel
-from ....modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 logger = logging.getLogger(__name__)
@@ -47,9 +47,11 @@ class RBLNDPTForDepthEstimation(RBLNModel):
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_image_size: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
+        rbln_image_size = rbln_kwargs.get("image_size", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
@@ -79,10 +81,22 @@ class RBLNDPTForDepthEstimation(RBLNModel):
 
         input_info = [("pixel_values", [rbln_batch_size, 3, rbln_image_size[0], rbln_image_size[1]], "float32")]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
-        meta = {"rbln_image_size": rbln_image_size}
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update(
+            {
+                "image_size": rbln_image_size,
+                "batch_size": rbln_batch_size,
+            }
+        )
 
-        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        return rbln_config
 
     def forward(self, *args, **kwargs):
         predicted_depth = super().forward(*args, **kwargs)

optimum/rbln/transformers/models/gemma/gemma_architecture.py

@@ -39,9 +39,16 @@ from ...models.decoderonly import (
 class GemmaWrapper(DecoderOnlyWrapper):
     def get_forward_dict(self):
         forward_dict = {}
-        forward_dict.update({"wrapper": GemmaModel.forward, "model": DecoderOnlyDecoderLayer.forward, "decoder_layer": DecoderOnlyAttention.forward,})
+        forward_dict.update(
+            {
+                "wrapper": GemmaModel.forward,
+                "model": DecoderOnlyDecoderLayer.forward,
+                "decoder_layer": DecoderOnlyAttention.forward,
+            }
+        )
         return forward_dict
 
+
 class GemmaModel:
     def forward(
         self,
@@ -54,7 +61,7 @@ class GemmaModel:
         use_cache: Optional[bool] = True,
         output_attentions: Optional[bool] = False,
         output_hidden_states: Optional[bool] = False,
-        forward_dict : Optional[Dict[str, classmethod]] = None,
+        forward_dict: Optional[Dict[str, classmethod]] = None,
         rotary_pos_emb=None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         # embed positions
@@ -89,7 +96,7 @@ class GemmaModel:
                 batch_ids=batch_ids,
                 cos=cos,
                 sin=sin,
-                forward_dict=forward_dict
+                forward_dict=forward_dict,
             )
 
             hidden_states = layer_outputs[0]

optimum/rbln/transformers/models/gemma/modeling_gemma.py

@@ -23,14 +23,19 @@
 
 import inspect
 import logging
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
-from transformers import GemmaForCausalLM, PreTrainedModel
+from transformers import GemmaForCausalLM
 
 from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .gemma_architecture import GemmaWrapper
 
 
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+    from ....modeling_config import RBLNConfig
+
 logger = logging.getLogger(__name__)
 
 
@@ -46,7 +51,8 @@ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     @classmethod
-    def wrapping_torch_model(self, model: "PreTrainedModel", rbln_max_seq_len: int):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
         return GemmaWrapper(model, rbln_max_seq_len).eval()
 
     def __getattr__(self, __name: str) -> Any:

optimum/rbln/transformers/models/gpt2/gpt2_architecture.py

@@ -53,6 +53,7 @@ class GPT2LMHeadModelWrapper(torch.nn.Module):
         attention_mask,
         cache_position,
         batch_position,
+        query_idx,
         *past_key_values,
     ):
         if input_ids.shape[1] == 1:
@@ -79,11 +80,13 @@ class GPT2LMHeadModelWrapper(torch.nn.Module):
         )
 
         hidden_states = outputs[0]
+        if batch_position >= 0:
+            hidden_states = hidden_states[:, query_idx].unsqueeze(1)
         logits = self.lm_head(hidden_states)
 
         output = (logits,) + outputs[1:]
 
-        return output, batch_position
+        return output, batch_position + query_idx
 
 
 class _GPT2Model:

optimum/rbln/transformers/models/gpt2/modeling_gpt2.py

@@ -23,23 +23,18 @@
 
 import inspect
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable
 
-from transformers import GPT2LMHeadModel, PretrainedConfig, PreTrainedModel
+from transformers import GPT2LMHeadModel
 
-from ....modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ....modeling_config import RBLNConfig
 from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .gpt2_architecture import GPT2LMHeadModelWrapper
 
 
 logger = logging.getLogger(__name__)
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-        PretrainedConfig,
-    )
+    from transformers import PreTrainedModel
 
 
 class RBLNGPT2LMHeadModel(RBLNDecoderOnlyModelForCausalLM):
@@ -57,7 +52,8 @@ class RBLNGPT2LMHeadModel(RBLNDecoderOnlyModelForCausalLM):
     """
 
     @classmethod
-    def wrapping_torch_model(self, model: "PreTrainedModel", rbln_max_seq_len: int):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
         return GPT2LMHeadModelWrapper(model, rbln_max_seq_len).eval()
 
     def __getattr__(self, __name: str) -> Any:
@@ -74,82 +70,3 @@ class RBLNGPT2LMHeadModel(RBLNDecoderOnlyModelForCausalLM):
         if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
             return redirect(val)
         return val
-
-    @classmethod
-    def _get_rbln_config(
-        cls,
-        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "PretrainedConfig",
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        **kwargs,
-    ) -> RBLNConfig:
-        meta = {}
-
-        prefill_chunk_size = 128
-        if rbln_max_seq_len is None:  # differenct from llama
-            rbln_max_seq_len = getattr(model_config, "n_positions", None)
-        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-
-        meta["rbln_max_seq_len"] = rbln_max_seq_len
-        meta["rbln_batch_size"] = rbln_batch_size
-        meta["rbln_prefill_chunk_size"] = prefill_chunk_size
-
-        def get_input_info(
-            batch_size,
-            query_length,
-        ):
-            head_dim = (
-                model_config.head_dim
-                if hasattr(model_config, "head_dim")
-                else model_config.hidden_size // model_config.n_head
-            )
-            input_info = [
-                ("input_ids", [batch_size, query_length], "int64"),
-                ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "int64"),
-                (
-                    "cache_position",
-                    [batch_size, query_length],
-                    "int32",
-                ),
-                ("batch_position", [], "int16"),
-            ]
-
-            input_info.extend(
-                [
-                    (
-                        f"past_key_values_{i}",
-                        [
-                            rbln_batch_size,
-                            model_config.n_head,  # differenct from llama
-                            rbln_max_seq_len,
-                            head_dim,
-                        ],
-                        "float32",
-                    )
-                    for i in range(model_config.n_layer * 2)  # differenct from llama
-                ]
-            )
-
-            return input_info
-
-        prefill_input_info = get_input_info(
-            batch_size=1,
-            query_length=prefill_chunk_size,
-        )
-        dec_input_info = get_input_info(
-            batch_size=rbln_batch_size,
-            query_length=1,
-        )
-
-        prefill_rbln_runtime_config = RBLNRuntimeConfig(input_info=prefill_input_info)
-        dec_rbln_runtime_config = RBLNRuntimeConfig(input_info=dec_input_info)
-
-        dec_rbln_runtime_config.batch_size = rbln_batch_size
-
-        rbln_config = RBLNConfig.from_rbln_runtime_configs(
-            [prefill_rbln_runtime_config, dec_rbln_runtime_config],
-            _rbln_meta=meta,
-        )
-
-        return rbln_config

optimum/rbln/transformers/models/llama/modeling_llama.py

@@ -23,14 +23,19 @@
 
 import inspect
 import logging
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
-from transformers import LlamaForCausalLM, PreTrainedModel
+from transformers import LlamaForCausalLM
 
 from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .llama_architecture import LlamaWrapper
 
 
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+    from ....modeling_config import RBLNConfig
+
 logger = logging.getLogger(__name__)
 
 
@@ -46,7 +51,8 @@ class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     @classmethod
-    def wrapping_torch_model(self, model: "PreTrainedModel", rbln_max_seq_len: int):
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
         return LlamaWrapper(model, rbln_max_seq_len).eval()
 
     def __getattr__(self, __name: str) -> Any:

optimum/rbln/transformers/models/llava_next/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .modeling_llava_next import RBLNLlavaNextForConditionalGeneration