optimum-rbln 0.9.1__py3-none-any.whl → 0.9.2a0__py3-none-any.whl

optimum/rbln/transformers/models/colpali/configuration_colpali.py

@@ -14,6 +14,10 @@
 from typing import Any, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
@@ -47,6 +51,7 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
 
     def __init__(
         self,
+        batch_size: Optional[int] = None,
         max_seq_lens: Union[int, List[int]] = None,
         output_hidden_states: Optional[bool] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
@@ -54,6 +59,8 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
     ):
         """
         Args:
+            batch_size (Optional[int]): The batch size for the model.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
             max_seq_lens (Union[int, List[int]]): The maximum sequence lengths for the language model.
                 This can be multiple values, and the model will be compiled for each max_seq_len, allowing selection of the most appropriate max_seq_len at inference time.
             output_hidden_states (Optional[bool]): Whether to output the hidden states of the language model.
@@ -63,6 +70,15 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
             ValueError: If batch_size is not a positive integer.
         """
         super().__init__(**kwargs)
-        self.vision_tower = vision_tower
+        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 ColPali vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower, batch_size=1, force_kwargs=True
+        )
         self.max_seq_lens = max_seq_lens
         self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/colpali/modeling_colpali.py

@@ -26,6 +26,7 @@ from transformers.models.paligemma.modeling_paligemma import PaliGemmaMultiModal
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
+from ...utils.rbln_runtime_wrapper import LoopProcessor
 from .colpali_architecture import RBLNColPaliForRetrievalWrapper
 
 
@@ -33,93 +34,64 @@ if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
-class LoopVisionTower:
-    def __init__(self, vision_tower: RBLNModel) -> None:
-        self.vision_tower = vision_tower
+class LoopVisionTower(LoopProcessor):
+    def __init__(self, vision_tower: "RBLNModel"):
+        super().__init__(model=vision_tower.model[0])
 
-    def forward(self, pixel_values, **kwargs):
-        batch_size = pixel_values.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.vision_tower(pixel_values[i : i + 1]))
+    def _get_batch_size(self, pixel_values, **kwargs):
+        return pixel_values.shape[0]
 
-        last_hidden_states = [output.last_hidden_state for output in outputs]
-        last_hidden_states = torch.cat(last_hidden_states, dim=0)
+    def _prepare_inputs_for_iteration(self, index, common_inputs, pixel_values, **kwargs):
+        pixel_values_item = pixel_values[index : index + 1]
+        out_buffer = kwargs["out"][index : index + 1]
+        return ([pixel_values_item], {"out": out_buffer})
 
+    def _process_outputs(self, outputs: list, **kwargs) -> "BaseModelOutputWithPooling":
         return BaseModelOutputWithPooling(
-            last_hidden_state=last_hidden_states,
+            last_hidden_state=kwargs["out"],
         )
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.vision_tower)
-
 
-class LoopLanguageModel:
-    def __init__(self, language_model: RBLNModel, rbln_config: RBLNModelConfig) -> None:
-        self.language_model = language_model
+class LoopLanguageModel(LoopProcessor):
+    def __init__(self, language_model: RBLNModel, rbln_config: RBLNModelConfig):
+        super().__init__(model=language_model)
         self.rbln_config = rbln_config
 
-    def prepare_inputs(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor):
+    def _get_batch_size(self, inputs_embeds, **kwargs):
+        return inputs_embeds.shape[0]
+
+    def _prepare_inputs_before_loop(self, *, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, **kwargs):
         input_len = inputs_embeds.shape[1]
         idx = bisect.bisect_left(self.rbln_config.max_seq_lens, input_len)
         if idx == len(self.rbln_config.max_seq_lens):
             raise ValueError(
                 f"Required seq_len({input_len}) is larger than available max_seq_lens({self.rbln_config.max_seq_lens})."
             )
-        else:
-            max_seq_len = self.rbln_config.max_seq_lens[idx]
-
-        inputs_embed = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, max_seq_len - input_len))
-        attn_mask = torch.nn.functional.pad(attention_mask, (0, max_seq_len - input_len)).to(torch.float32)
-        position_ids = torch.arange(max_seq_len, dtype=torch.int32).view(1, -1)
-
-        return inputs_embed, attn_mask, position_ids
-
-    def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, **kwargs):
-        padded_inputs_embed, padded_attn_mask, padded_position_ids = self.prepare_inputs(inputs_embeds, attention_mask)
-        input_batch_size = inputs_embeds.shape[0]
-        input_seq_len = inputs_embeds.shape[1]
-
-        all_embeddings = []
-        all_hidden_states = []
-        for i in range(input_batch_size):
-            outputs = self.language_model(
-                inputs_embeds=padded_inputs_embed[i : i + 1],
-                attention_mask=padded_attn_mask[i : i + 1],
-                position_ids=padded_position_ids,
-            )
-
-            if self.rbln_config.output_hidden_states:
-                embedding = outputs[0]
-                hidden_states = outputs[1:]
-            else:
-                embedding = outputs
-                hidden_states = None
+        max_seq_len = self.rbln_config.max_seq_lens[idx]
+        padded_inputs_embed = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, max_seq_len - input_len))
+        padded_attn_mask = torch.nn.functional.pad(attention_mask, (0, max_seq_len - input_len)).to(torch.float32)
+        padded_position_ids = torch.arange(max_seq_len, dtype=torch.int32).view(1, -1)
+
+        return {
+            "padded_inputs_embed": padded_inputs_embed,
+            "padded_attn_mask": padded_attn_mask,
+            "padded_position_ids": padded_position_ids,
+        }
 
-            all_embeddings.append(embedding)
-            all_hidden_states.append(hidden_states)
+    def _prepare_inputs_for_iteration(self, index: int, common_inputs, *args, **kwargs):
+        item_kwargs = {
+            "inputs_embeds": common_inputs["padded_inputs_embed"][index : index + 1],
+            "attention_mask": common_inputs["padded_attn_mask"][index : index + 1],
+            "position_ids": common_inputs["padded_position_ids"],
+            "out": [tensor[index : index + 1] for tensor in kwargs["out"]],
+        }
+        return ([], item_kwargs)
 
-        embeddings = torch.cat(all_embeddings, dim=0)[:, :input_seq_len]
+    def _process_outputs(self, outputs: list, **kwargs):
         if self.rbln_config.output_hidden_states:
-            hidden_states = [
-                torch.cat(
-                    [batch_hidden_states[layer_idx][:, :input_seq_len] for batch_hidden_states in all_hidden_states],
-                    dim=0,
-                )
-                for layer_idx in range(len(all_hidden_states[0]))
-            ]
-            return embeddings, tuple(hidden_states)
+            return kwargs["out"][0], tuple(kwargs["out"][1:])
         else:
-            return embeddings
-
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.language_model)
+            return kwargs["out"]
 
 
 class RBLNColPaliForRetrieval(RBLNModel):
@@ -212,7 +184,7 @@ class RBLNColPaliForRetrieval(RBLNModel):
     @classmethod
     def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return RBLNColPaliForRetrievalWrapper(
-            causal_lm=model.vlm.language_model,
+            causal_lm=model.vlm,
             embedding_proj_layer=model.embedding_proj_layer,
             max_seq_len=max(rbln_config.max_seq_lens),
             output_hidden_states=rbln_config.output_hidden_states,
@@ -252,9 +224,9 @@ class RBLNColPaliForRetrieval(RBLNModel):
         input_infos = []
         for max_seq_len in rbln_config.max_seq_lens:
             input_info = [
-                ("inputs_embeds", [1, max_seq_len, hidden_size], "float32"),
-                ("attention_mask", [1, max_seq_len], "float32"),
-                ("position_ids", [1, max_seq_len], "int32"),
+                ("inputs_embeds", [rbln_config.vision_tower.batch_size, max_seq_len, hidden_size], "float32"),
+                ("attention_mask", [rbln_config.vision_tower.batch_size, max_seq_len], "float32"),
+                ("position_ids", [rbln_config.vision_tower.batch_size, max_seq_len], "int32"),
             ]
             input_infos.append(input_info)
 
@@ -306,8 +278,7 @@ class RBLNColPaliForRetrieval(RBLNModel):
     def get_pytorch_model(cls, *args, **kwargs):
         model = super().get_pytorch_model(*args, **kwargs)
         model.vision_tower = model.vlm.vision_tower
-        del model.vlm.vision_tower
-
+        del model.vlm.model.vision_tower
         return model
 
     def get_image_features(self, pixel_values: torch.Tensor):
@@ -318,8 +289,14 @@ class RBLNColPaliForRetrieval(RBLNModel):
         # Returns:
         #     image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
 
-        vision_outputs = self.vision_tower(pixel_values).last_hidden_state
-        image_features = self.multi_modal_projector(vision_outputs)
+        vision_output_size = [
+            pixel_values.shape[0],
+            self.config.vlm_config.vision_config.num_image_tokens,
+            self.config.vlm_config.vision_config.hidden_size,
+        ]
+        vision_output = torch.empty(size=vision_output_size, dtype=torch.float32, device="cpu")
+        self.vision_tower(pixel_values, out=vision_output)
+        image_features = self.multi_modal_projector(vision_output)
         image_features = image_features / (self.config.text_config.hidden_size**0.5)
         return image_features
 
@@ -385,11 +362,27 @@ class RBLNColPaliForRetrieval(RBLNModel):
             input_ids=input_ids, inputs_embeds=inputs_embeds, pixel_values=pixel_values
         )
 
+        outputs = []
+        language_model_out_size = [inputs_embeds.shape[0], self.rbln_config.max_seq_lens[0], self.config.embedding_dim]
+        language_model_hidden_states_size = [
+            inputs_embeds.shape[0],
+            self.rbln_config.max_seq_lens[0],
+            self.rbln_config.max_seq_lens[0],
+        ]
+        outputs.append(torch.empty(size=language_model_out_size, dtype=torch.float32, device="cpu"))
+        if self.rbln_config.output_hidden_states:
+            for i in range(self.config.vlm_config.text_config.num_hidden_layers + 1):
+                outputs.append(torch.empty(size=language_model_hidden_states_size, dtype=torch.float32, device="cpu"))
+
         # Embedding_proj_layer is fused on the bottom of the language model.
-        outputs = self.language_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask)
+        self.language_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, out=outputs)
 
-        embeddings = outputs if not self.rbln_config.output_hidden_states else outputs[0]
-        hidden_states = None if not self.rbln_config.output_hidden_states else outputs[1]
+        embeddings = outputs[0][:, : inputs_embeds.shape[1]]
+        hidden_states = (
+            None
+            if not self.rbln_config.output_hidden_states
+            else [tensor[0][:, : inputs_embeds.shape[1]] for tensor in outputs[1:]]
+        )
 
         # L2 normalization
         embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)  # (batch_size, sequence_length, dim)

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

@@ -57,7 +57,6 @@ class RBLNDecoderOnlyModel(RBLNModel, RBLNDecoderOnlyFlashAttentionMixin):
     1. Converting pre-trained transformer models to RBLN-optimized format
     2. Handling the compilation process for RBLN devices
     3. Managing inference operations for decoder-only architectures
-
     This class inherits from RBLNModel and implements specific methods required for
     decoder-only architectures.
 
@@ -68,6 +67,8 @@ class RBLNDecoderOnlyModel(RBLNModel, RBLNDecoderOnlyFlashAttentionMixin):
         - The class handles RBLN-specific optimizations automatically during compilation
     """
 
+    _tp_support = True
+
     main_input_name = "input_ids"
     auto_model_class = AutoModel
     _decoder_wrapper_cls = DecoderOnlyWrapper
@@ -642,7 +643,6 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNDecoderOnlyModel, RBLNDecoderOnlyGener
     1. Converting pre-trained transformer models to RBLN-optimized format
     2. Handling the compilation process for RBLN devices
     3. Managing inference operations for causal language modeling
-
     This class inherits from RBLNModel and implements specific methods required for
     decoder-only architectures and causal language modeling tasks.
 

optimum/rbln/transformers/models/gemma3/configuration_gemma3.py

@@ -14,8 +14,11 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
 from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
-from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
+
+
+logger = get_logger(__name__)
 
 
 class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -89,8 +92,13 @@ class RBLNGemma3ForConditionalGenerationConfig(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_tower = self.init_submodule_config(RBLNSiglipVisionModelConfig, vision_tower)
-        self.language_model = self.init_submodule_config(RBLNGemma3ForCausalLMConfig, language_model)
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Gemma3 vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower, batch_size=1, force_kwargs=True
+        )
+        self.language_model = self.initialize_submodule_config(submodule_config=language_model)
 
     @property
     def image_prefill_chunk_size(self):

optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -11,6 +11,7 @@
 # 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.
+import importlib
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -25,6 +26,7 @@ from transformers.models.gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbed
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ...modeling_outputs import RBLNDecoderOnlyOutput
+from ...utils.rbln_runtime_wrapper import LoopProcessor
 from ..decoderonly.decoderonly_runtime_utils import RBLNPageTableManager
 from ..decoderonly.modeling_decoderonly import (
     RBLNDecoderOnlyModelForCausalLM,
@@ -38,58 +40,41 @@ if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, Gemma3ForConditionalGeneration
 
 
-class LoopVisionTower:
-    def __init__(self, vision_tower: RBLNModel) -> None:
-        self.vision_tower = vision_tower
+class LoopVisionTower(LoopProcessor):
+    def __init__(self, vision_tower: "RBLNModel"):
+        super().__init__(model=vision_tower)
 
-    def forward(self, *args, **kwargs):
-        # Loop instead of batch
-        # shape of pixel_values : [batch, num_channel, height, width]
-        pixel_values = args[0]
+    def _get_batch_size(self, pixel_values, **kwargs):
+        return pixel_values.shape[0]
 
-        batch_size = pixel_values.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.vision_tower(pixel_values=pixel_values[i : i + 1], return_dict=True))
+    def _prepare_inputs_for_iteration(self, index, common_inputs, pixel_values, **kwargs):
+        pixel_values_item = pixel_values[index : index + 1]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        return ([pixel_values_item], {"out": out_buffer})
 
-        last_hidden_states = [output.last_hidden_state for output in outputs]
-
-        # FIXME:: This can be optimized using out= API of rbln runtime.
-        last_hidden_states = torch.cat(last_hidden_states, dim=0)
+    def _process_outputs(self, outputs: list, **kwargs) -> "BaseModelOutputWithPooling":
+        output = kwargs["out"]
 
         return BaseModelOutputWithPooling(
-            last_hidden_state=last_hidden_states,
+            last_hidden_state=output[0],
         )
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.vision_tower)
-
 
-class LoopProjector:
-    def __init__(self, multi_modal_projector) -> None:
-        self.multi_modal_projector = multi_modal_projector
+class LoopProjector(LoopProcessor):
+    def __init__(self, multi_modal_projector: "RBLNModel"):
+        super().__init__(model=multi_modal_projector)
 
-    def forward(self, *args, **kwargs):
-        # Loop instead of batch
-        image_feature = args[0]
+    def _get_batch_size(self, image_feature, **kwargs):
+        return image_feature.shape[0]
 
-        batch_size = image_feature.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.multi_modal_projector(image_feature[i : i + 1]))
+    def _prepare_inputs_for_iteration(self, index, common_inputs, image_feature, **kwargs):
+        image_feature_item = image_feature[index : index + 1]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        return ([image_feature_item], {"out": out_buffer})
 
-        # FIXME:: This can be optimized using out= API of rbln runtime.
-        outputs = torch.cat(outputs, dim=0)
-        return outputs
-
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.multi_modal_projector)
+    def _process_outputs(self, outputs: list, **kwargs):
+        output = kwargs["out"]
+        return output[0]
 
 
 class RBLNGemma3ForConditionalGeneration(RBLNModel):
@@ -112,6 +97,23 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
     def can_generate(self):
         return True
 
+    @classmethod
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
+        with no_init_weights():
+            model_cls_name = model.model.language_model.__class__.__name__
+            causal_model_cls_name = model_cls_name.replace("TextModel", "ForCausalLM")
+            causal_model_cls = getattr(importlib.import_module("transformers"), causal_model_cls_name)
+            new_language_model = causal_model_cls(model.model.language_model.config)
+
+        new_language_model.lm_head = model.lm_head
+        new_language_model.model = model.model.language_model
+        model.model.language_model = new_language_model
+        model.lm_head = None
+        del model.lm_head
+        return model
+
     def __post_init__(self, **kwargs):
         self.vision_tower = LoopVisionTower(self.rbln_submodules[0])
         self.language_model = self.rbln_submodules[1]
@@ -210,8 +212,21 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel):
         # Returns:
         #     Image feature tensor of shape `(num_images, image_length, embed_dim)`.
 
-        vision_outputs = self.vision_tower(pixel_values).last_hidden_state
-        image_features = self.multi_modal_projector(vision_outputs)
+        vision_out_buffer = []
+        vision_out_size = [
+            pixel_values.shape[0],
+            (self.config.vision_config.image_size // self.config.vision_config.patch_size) ** 2,
+            self.config.vision_config.hidden_size,
+        ]
+        projector_out_size = [
+            pixel_values.shape[0],
+            self.config.mm_tokens_per_image,
+            self.config.text_config.hidden_size,
+        ]
+        vision_out_buffer.append(torch.empty(size=vision_out_size, dtype=torch.float32, device="cpu"))
+        projector_out_buffer = [torch.empty(size=projector_out_size, dtype=torch.float32, device="cpu")]
+        vision_outputs = self.vision_tower(pixel_values, out=vision_out_buffer).last_hidden_state
+        image_features = self.multi_modal_projector(vision_outputs, out=projector_out_buffer)
         return image_features
 
     def _preprocess_prefill(

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

@@ -15,10 +15,29 @@
 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):
@@ -61,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

@@ -75,10 +75,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():
@@ -140,8 +142,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,
                 ],
@@ -160,22 +161,24 @@ 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:
@@ -285,8 +288,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,
                 ],
@@ -425,10 +427,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/llava/configuration_llava.py

@@ -15,6 +15,10 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNLlavaForConditionalGenerationConfig(RBLNModelConfig):
@@ -54,5 +58,15 @@ class RBLNLlavaForConditionalGenerationConfig(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_tower = vision_tower
-        self.language_model = language_model
+        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,
+        )