optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3__py3-none-any.whl

optimum/rbln/transformers/models/idefics3/configuration_idefics3.py

@@ -12,13 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNIdefics3VisionTransformerConfig(RBLNModelConfig):
-    pass
+    """
+    Configuration class for RBLNIdefics3VisionTransformer.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Idefics3 vision transformer.
+    """
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
 
 class RBLNIdefics3ForConditionalGenerationConfig(RBLNModelConfig):
@@ -39,17 +58,21 @@ class RBLNIdefics3ForConditionalGenerationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         vision_model: Optional[RBLNModelConfig] = None,
         text_model: Optional[RBLNModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
             vision_model (Optional[RBLNModelConfig]): Configuration for the vision transformer component.
+                This can include settings specific to the vision encoder, such as input resolution or other vision-related parameters.
+                If not provided, default settings will be used.
             text_model (Optional[RBLNModelConfig]): Configuration for the text model component.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+                This can include settings specific to the language model, such as tensor parallelism or other text-related parameters.
+                If not provided, default settings will be used.
+            kwargs: Additional arguments passed to the parent `RBLNModelConfig`.
 
         Raises:
-            ValueError: If batch_size is not a positive integer.
+            ValueError: If `batch_size` is not a positive integer.
         """
 
         super().__init__(**kwargs)
@@ -57,5 +80,10 @@ class RBLNIdefics3ForConditionalGenerationConfig(RBLNModelConfig):
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
-        self.vision_model = vision_model
-        self.text_model = text_model
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Idefics3 vision transformer. It will be set to 1.")
+
+        self.vision_model = self.initialize_submodule_config(
+            submodule_config=vision_model, batch_size=1, force_kwargs=True
+        )
+        self.text_model = self.initialize_submodule_config(submodule_config=text_model)

optimum/rbln/transformers/models/idefics3/modeling_idefics3.py

@@ -34,17 +34,12 @@ from transformers.models.idefics3.modeling_idefics3 import Idefics3CausalLMOutpu
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.runtime_utils import RBLNPytorchRuntime
-from ..decoderonly.modeling_decoderonly import (
-    RBLNDecoderOnlyOutput,
-)
+from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
 class RBLNRuntimeVisionModel(RBLNPytorchRuntime):
@@ -81,10 +76,12 @@ class RBLNRuntimeVisionModel(RBLNPytorchRuntime):
 
         hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)
 
-        return super().forward(hidden_states.contiguous())
+        return super().forward(hidden_states.contiguous(), **kwargs)
 
 
 class RBLNIdefics3VisionTransformer(RBLNModel):
+    _tp_support = False
+
     def __post_init__(self, **kwargs):
         artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
         with no_init_weights():
@@ -113,7 +110,7 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
         return self.embeddings
 
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
         class Idefics3VisionTransformerWrapper(torch.nn.Module):
             def __init__(self, model: "Idefics3VisionTransformer"):
                 super().__init__()
@@ -124,9 +121,6 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
                 encoder_outputs = self.encoder(
                     inputs_embeds=hidden_states,
                     attention_mask=patch_attention_mask,
-                    output_attentions=None,
-                    output_hidden_states=None,
-                    return_dict=False,
                 )
                 last_hidden_state = encoder_outputs[0]
                 last_hidden_state = self.post_layernorm(last_hidden_state)
@@ -146,8 +140,7 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
             (
                 "hidden_states",
                 [
-                    # batch_size * num_patches (dependent on image size) -> compile with 1 and use for loop
-                    1,
+                    rbln_config.batch_size,
                     (model_config.image_size // model_config.patch_size) ** 2,
                     model_config.hidden_size,
                 ],
@@ -166,29 +159,31 @@ class RBLNIdefics3VisionTransformer(RBLNModel):
         return_dict: Optional[bool] = None,
         **kwargs,
     ) -> Union[Tuple, BaseModelOutput]:
-        batch_size = pixel_values.shape[0]
-        last_hidden_state = []
-        for i in range(batch_size):
+        last_hidden_state_size = [
+            pixel_values.shape[0],
+            (self.config.image_size // self.config.patch_size) ** 2,
+            self.config.hidden_size,
+        ]
+        last_hidden_state = torch.empty(size=last_hidden_state_size, dtype=torch.float32, device="cpu")
+        for i in range(pixel_values.shape[0]):
             if patch_attention_mask is not None:
                 batch_attention_mask = patch_attention_mask[i : i + 1,]
             else:
                 batch_attention_mask = None
 
-            batch_hidden_state = self.model(
+            self.model(
                 pixel_values[i : i + 1,],
                 batch_attention_mask,
+                out=last_hidden_state[i : i + 1,],
                 return_dict=False,
             )
-            last_hidden_state.append(batch_hidden_state)
-        last_hidden_state = torch.cat(last_hidden_state, dim=0)
-
         if not return_dict:
             return (last_hidden_state,)
         else:
             return BaseModelOutput(last_hidden_state=last_hidden_state)
 
 
-class RBLNIdefics3ForConditionalGeneration(RBLNModel):
+class RBLNIdefics3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNIdefics3ForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
@@ -245,9 +240,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
         return True
 
     @classmethod
-    def get_pytorch_model(cls, *args, **kwargs):
-        model = super().get_pytorch_model(*args, **kwargs)
-
+    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
         with no_init_weights():
             model_cls_name = model.model.text_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
@@ -276,7 +269,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
         return self.text_model.get_input_embeddings()
 
     @classmethod
-    def wrap_model_if_needed(cls, model, rbln_config):
+    def _wrap_model_if_needed(cls, model, rbln_config):
         return model.model.connector
 
     @classmethod
@@ -291,8 +284,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
             (
                 "image_hidden_states",
                 [
-                    # batch_size * num_patches (dependent on image size) -> compile with 1 and use for loop
-                    1,
+                    rbln_config.vision_model.batch_size,
                     (model_config.vision_config.image_size // model_config.vision_config.patch_size) ** 2,
                     model_config.vision_config.hidden_size,
                 ],
@@ -431,10 +423,15 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
                 pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True
             ).last_hidden_state
 
-            connector_outputs = []
+            connector_output_size = [
+                image_hidden_states.shape[0],
+                image_hidden_states.shape[1] // self.config.scale_factor**2,
+                self.config.text_config.hidden_size,
+            ]
+            connector_outputs = torch.empty(size=connector_output_size, dtype=torch.float32, device="cpu")
             for i in range(image_hidden_states.shape[0]):
-                connector_outputs.append(self.connector(image_hidden_states[i : i + 1,]))
-            image_hidden_states = torch.cat(connector_outputs, dim=0)
+                self.connector(image_hidden_states[i : i + 1,], out=connector_outputs[i : i + 1,])
+            image_hidden_states = connector_outputs
 
         elif image_hidden_states is not None:
             image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device)

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_llama import RBLNLlamaForCausalLMConfig
-from .modeling_llama import RBLNLlamaForCausalLM
+from .configuration_llama import RBLNLlamaForCausalLMConfig, RBLNLlamaModelConfig
+from .modeling_llama import RBLNLlamaForCausalLM, RBLNLlamaModel

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNLlamaModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Llama models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

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

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .llama_architecture import LlamaWrapper
 
 
@@ -81,3 +81,24 @@ class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = LlamaWrapper
+
+
+class RBLNLlamaModel(RBLNDecoderOnlyModel):
+    """
+    The Llama Model transformer outputting raw hidden-states without any specific head on top.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based LlamaModel on RBLN devices.
+    It implements the methods to convert a pre-trained transformers LlamaModel into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNLlamaModelConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNLlamaModelConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNLlamaModelConfig`] class for all available configuration options.
+    """
+
+    _decoder_wrapper_cls = LlamaWrapper

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

@@ -0,0 +1,16 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# 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.
+
+from .configuration_llava import RBLNLlavaForConditionalGenerationConfig
+from .modeling_llava import RBLNLlavaForConditionalGeneration

optimum/rbln/transformers/models/llava/configuration_llava.py

@@ -0,0 +1,72 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# 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.
+
+from typing import Any, Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class RBLNLlavaForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNLlavaForConditionalGenerationConfig.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized LLaVA models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
+    submodules = ["vision_tower", "language_model"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+                This can include settings specific to the vision encoder, such as input resolution or other vision-related parameters.
+                If not provided, default settings will be used.
+            language_model (Optional[RBLNModelConfig]): Configuration for the language model component.
+                This can include settings specific to the language model, such as tensor parallelism or other text-related parameters.
+                If not provided, default settings will be used.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If `batch_size` is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Llava vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower,
+            batch_size=1,  # vision_tower batch_size is always 1 in Llava
+            force_kwargs=True,
+        )
+
+        self.language_model = self.initialize_submodule_config(
+            submodule_config=language_model,
+        )