optimum-rbln 0.9.3rc0__py3-none-any.whl → 0.9.4a2__py3-none-any.whl

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

@@ -99,9 +99,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMix
         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.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("TextModel", "ForCausalLM")
@@ -135,7 +133,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMix
         return self.language_model.get_input_embeddings()
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return model.multi_modal_projector
 
     @classmethod
@@ -301,28 +299,60 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMix
         generate_idx: Optional[torch.Tensor] = None,
         padded_cache_lengths: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
         **lm_kwargs: Dict[str, Any],
     ) -> Union[Tuple, RBLNDecoderOnlyOutput]:
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.rbln_config.language_model.output_hidden_states
+        )
+        if output_hidden_states != self.rbln_config.language_model.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.language_model.output_hidden_states {self.rbln_config.language_model.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
+
         # prefill
         if cache_position is None:
             logits = []
             inputs_embeds = self._preprocess_prefill(input_ids, inputs_embeds, pixel_values)
             batch_size = inputs_embeds.shape[0]
 
+            all_hidden_states = (
+                tuple(
+                    torch.zeros(
+                        batch_size,
+                        inputs_embeds.shape[1],
+                        self.config.text_config.hidden_size,
+                        dtype=self.rbln_config.torch_dtype,
+                    )
+                    for _ in range(self.config.text_config.num_hidden_layers + 1)
+                )
+                if self.rbln_config.language_model.output_hidden_states
+                else None
+            )
+
             for b_idx in range(batch_size):
                 cache_position = torch.arange(0, generate_idx[b_idx].item(), dtype=torch.int32).unsqueeze(0)
                 token_type_id = token_type_ids[b_idx : b_idx + 1, attention_mask[b_idx].bool()]
                 cache_position = self.get_padded_cache_position(cache_position, token_type_id)
 
-                output = self.language_model.prefill_decoder(
+                outputs = self.language_model.prefill_decoder(
                     inputs_embeds=inputs_embeds[b_idx : b_idx + 1],
                     attention_mask=attention_mask[b_idx],
                     cache_position=cache_position,
                     batch_idx=b_idx,
                     token_type_ids=token_type_ids[b_idx : b_idx + 1],  # do not pass token_type_id
                 )
-                padded_cache_lengths[b_idx] += output.padded_cache_lengths
-                logits.append(output.logits)
+                padded_cache_lengths[b_idx] += outputs.padded_cache_lengths
+                logits.append(outputs.logits)
+                if self.rbln_config.language_model.output_hidden_states:
+                    for l_idx in range(self.config.text_config.num_hidden_layers + 1):
+                        mask_indices = torch.nonzero(attention_mask[b_idx], as_tuple=True)[0]
+                        all_hidden_states[l_idx][b_idx].index_copy_(
+                            dim=0, index=mask_indices, source=outputs.hidden_states[l_idx][0]
+                        )
 
             logits = torch.cat(logits, dim=0)
         # decoder
@@ -336,15 +366,20 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMix
                     f"Please run your model with one of these batch sizes or add support for batch size {batch_size}."
                 )
 
-            logits = self.language_model.decoders[batch_size](
+            outputs = self.language_model.decoders[batch_size](
                 input_ids=input_ids,
                 inputs_embeds=inputs_embeds,
                 cache_position=cache_position,
                 position_ids=position_ids if self.rbln_config.language_model.use_position_ids else None,
-            ).logits
+            )
+            logits = outputs.logits
+            all_hidden_states = outputs.hidden_states
 
         return RBLNDecoderOnlyOutput(
-            logits=logits, generate_idx=generate_idx, padded_cache_lengths=padded_cache_lengths
+            logits=logits,
+            generate_idx=generate_idx,
+            padded_cache_lengths=padded_cache_lengths,
+            hidden_states=all_hidden_states,
         )
 
 
@@ -405,26 +440,6 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             )
         return embed_tokens
 
-    @classmethod
-    def _update_sliding_window_config(cls, model_config: PretrainedConfig, rbln_config: RBLNGemma3ForCausalLMConfig):
-        sliding_window = getattr(model_config, "sliding_window", None)
-        sliding_window_pattern = getattr(model_config, "sliding_window_pattern", None)
-        if sliding_window_pattern is None:
-            if hasattr(model_config, "layer_types"):
-                first_full_attention_index = model_config.layer_types.index("full_attention")
-                sliding_window_pattern = first_full_attention_index + 1
-            else:
-                raise ValueError("Cannot determine sliding_window_pattern from model_config")
-
-        if sliding_window_pattern <= model_config.num_hidden_layers:
-            rbln_config.cache_impl = "hybrid"
-            rbln_config.sliding_window = sliding_window
-            rbln_config.sliding_window_layers = [
-                i for i in range(model_config.num_hidden_layers) if (i + 1) % sliding_window_pattern > 0
-            ]
-
-        return rbln_config
-
     @classmethod
     def _update_submodule_config(
         cls,
@@ -482,7 +497,7 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNGemma3ForCausalLMConfig):
-        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
+        wrapped_model = cls._wrap_model_if_needed(model, rbln_config)
 
         rbln_compile_configs = rbln_config.compile_cfgs
         prefill_compile_config = rbln_compile_configs[0]

optimum/rbln/transformers/models/grounding_dino/grounding_dino_architecture.py

@@ -150,7 +150,7 @@ class _GroundingDinoEncoder(torch.nn.Module):
         all_attn_fused_vision = () if output_attentions else None
         all_attn_enhanced_text = () if output_attentions else None
         all_attn_deformable = () if output_attentions else None
-        for i, encoder_layer in enumerate(self.layers):
+        for _, encoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 encoder_vision_states += (vision_features,)
                 encoder_text_states += (text_features,)

optimum/rbln/transformers/models/grounding_dino/modeling_grounding_dino.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor, nn
@@ -206,8 +206,7 @@ class RBLNGroundingDinoForObjectDetection(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def get_pytorch_model(cls, *args, **kwargs):
-        model = super().get_pytorch_model(*args, **kwargs)
+    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
         model.encoder = model.model.encoder
         model.decoder = model.model.decoder
         model.text_backbone = model.model.text_backbone
@@ -217,7 +216,7 @@ class RBLNGroundingDinoForObjectDetection(RBLNModel):
         return model
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNGroundingDinoForObjectDetectionConfig
     ) -> torch.nn.Module:
         return model.model.text_projection
@@ -305,7 +304,6 @@ class RBLNGroundingDinoForObjectDetection(RBLNModel):
         for feature_map, mask in vision_features:
             # position encoding
             position_embeddings_list.append(self.backbone_position_embedding(feature_map, mask).to(feature_map.dtype))
-        vision_features, position_embeddings_list
 
         # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
         feature_maps = []
@@ -530,9 +528,26 @@ class RBLNGroundingDinoForObjectDetection(RBLNModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        labels: List[Dict[str, Union[torch.LongTensor, torch.FloatTensor]]] = None,
         **kwargs,
-    ):
+    ) -> Union[GroundingDinoObjectDetectionOutput, Tuple]:
+        """
+        Forward pass for the RBLN-optimized GroundingDinoForObjectDetection model.
+
+        Args:
+            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)): The tensors corresponding to the input images.
+            input_ids (torch.LongTensor of shape (batch_size, text_sequence_length)): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it.
+            token_type_ids (torch.LongTensor of shape (batch_size, text_sequence_length), optional): Segment token indices to indicate first and second portions of the inputs.
+            attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+            pixel_mask (torch.Tensor of shape (batch_size, height, width), optional): Mask to avoid performing attention on padding pixel values.
+            encoder_outputs (Tuple consists of last_hidden_state of shape(batch_size, sequence_length, hidden_size), optional): A sequence of hidden-states at the output of the last layer of the encoder.
+            output_attentions (bool, optional): Whether or not to return the attentions tensors of all attention layers.
+            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers.
+            return_dict (bool, optional): Whether or not to return a ModelOutput instead of a plain tuple.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a GroundingDinoObjectDetectionOutput object.
+        """
+
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # Pad image to rbln_config.image_height and rbln_config.image_width
@@ -663,7 +678,7 @@ class RBLNGroundingDinoEncoder(RBLNModel):
         self.encoder_runtime = RBLNPytorchRuntime(self.model[0])
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNGroundingDinoForObjectDetectionConfig
     ) -> torch.nn.Module:
         model = _GroundingDinoEncoder(model, rbln_config).eval()
@@ -861,7 +876,7 @@ class RBLNGroundingDinoDecoder(RBLNModel):
         self.decoder_runtime = RBLNPytorchRuntime(self.model[0])
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNGroundingDinoForObjectDetectionConfig
     ) -> torch.nn.Module:
         return _GroundingDinoDecoder(model, rbln_config).eval()

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

@@ -110,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__()
@@ -240,9 +240,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationM
         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")
@@ -271,7 +269,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationM
         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

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

@@ -88,15 +88,22 @@ class LoopVisionTower(LoopProcessor):
 
 
 class LoopProjector(LoopProcessor):
-    def __init__(self, multi_modal_projector: "RBLNModel"):
+    def __init__(self, multi_modal_projector: "RBLNModel", rbln_config=None):
         super().__init__(model=multi_modal_projector)
+        self.rbln_config = rbln_config
 
     def _get_batch_size(self, image_feature, **kwargs):
         return image_feature.shape[0]
 
     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"]]
+        if hasattr(self.rbln_config.vision_tower, "max_image_size"):
+            out_buffer = [
+                tensor[:, index * image_feature.shape[1] : (index + 1) * image_feature.shape[1], :]
+                for tensor in kwargs["out"]
+            ]
+        else:
+            out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
         return ([image_feature_item], {"out": out_buffer})
 
     def _process_outputs(self, outputs: list, **kwargs):
@@ -175,9 +182,7 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
         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.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
@@ -194,7 +199,7 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
     def __post_init__(self, **kwargs):
         self.vision_tower = LoopVisionTower(self.rbln_submodules[0])
         self.language_model = self.rbln_submodules[1]
-        self.multi_modal_projector = LoopProjector(self.model[0])
+        self.multi_modal_projector = LoopProjector(self.model[0], rbln_config=self.rbln_config)
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         return super().__post_init__(**kwargs)
 
@@ -208,7 +213,7 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
         return self.language_model.get_input_embeddings()
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return model.multi_modal_projector
 
     @classmethod
@@ -221,10 +226,8 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
     ) -> RBLNModelConfig:
         # support for pixtral that needs padding
         if hasattr(rbln_config.vision_tower, "max_image_size"):
-            num_positions = (
-                rbln_config.batch_size
-                * (rbln_config.vision_tower.max_image_size[0] // model_config.vision_config.patch_size)
-                * (rbln_config.vision_tower.max_image_size[1] // model_config.vision_config.patch_size)
+            num_positions = (rbln_config.vision_tower.max_image_size[0] // model_config.vision_config.patch_size) * (
+                rbln_config.vision_tower.max_image_size[1] // model_config.vision_config.patch_size
             )
             selected_image_feature_dim = num_positions
 
@@ -334,7 +337,7 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
             pooler_out_size = [pixel_values.shape[0], self.config.vision_config.hidden_size]
 
         vision_out_buffer = []
-        for i in range(self.config.vision_config.num_hidden_layers + 2):
+        for _ in range(self.config.vision_config.num_hidden_layers + 2):
             vision_out_buffer.append(torch.empty(size=vision_out_size, dtype=torch.float32, device="cpu"))
         if pooler_out_size is not None:
             vision_out_buffer.insert(1, torch.empty(size=pooler_out_size, dtype=torch.float32, device="cpu"))
@@ -353,23 +356,32 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
 
         if hasattr(self.rbln_config.vision_tower, "max_image_size"):
             num_real_patches = selected_image_feature.shape[1]
-            max_patches = (
-                (self.rbln_config.vision_tower.max_image_size[0] // self.config.vision_config.patch_size)
-                * (self.rbln_config.vision_tower.max_image_size[1] // self.config.vision_config.patch_size)
-                * pixel_values.shape[0]
+            max_patches = (self.rbln_config.vision_tower.max_image_size[0] // self.config.vision_config.patch_size) * (
+                self.rbln_config.vision_tower.max_image_size[1] // self.config.vision_config.patch_size
             )
-            num_padding_patches = max_patches - num_real_patches
 
-            projector_out_size = [1, max_patches, self.config.text_config.hidden_size]
+            chunks = []
+            for i in range(0, num_real_patches, max_patches):
+                chunk = selected_image_feature[:, i : i + max_patches, :]
+                chunk_size = chunk.shape[1]
+
+                if chunk_size < max_patches:
+                    padding_tensor = torch.zeros(
+                        (selected_image_feature.shape[0], max_patches - chunk_size, selected_image_feature.shape[2]),
+                        dtype=selected_image_feature.dtype,
+                    )
+                    chunk = torch.cat([chunk, padding_tensor], dim=1)
+                chunks.append(chunk)
+
+            split_features = torch.cat(chunks, dim=0)
+            num_chunks = len(chunks)
+            projector_out_size = [1, max_patches * num_chunks, self.config.text_config.hidden_size]
             projector_out_buffer = [torch.empty(size=projector_out_size, dtype=torch.float32, device="cpu")]
-
-            padding_tensor = torch.zeros(
-                (selected_image_feature.shape[0], num_padding_patches, selected_image_feature.shape[2]),
-                dtype=selected_image_feature.dtype,
+            projected_features = self.multi_modal_projector(split_features, out=projector_out_buffer)
+            projected_features = projected_features.view(
+                selected_image_feature.shape[0], num_chunks * max_patches, self.config.text_config.hidden_size
             )
-            padded_feature = torch.cat([selected_image_feature, padding_tensor], dim=1)
-            padded_projected_feature = self.multi_modal_projector(padded_feature, out=projector_out_buffer)
-            image_features = padded_projected_feature[:, :num_real_patches, :]
+            image_features = projected_features[:, :num_real_patches, :]
         else:
             projector_out_size = [
                 pixel_values.shape[0] * pixel_values.shape[1],

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

@@ -139,9 +139,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGeneration
         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.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
@@ -192,7 +190,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGeneration
         return self.language_model.get_input_embeddings()
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return model.multi_modal_projector
 
     @classmethod
@@ -302,7 +300,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGeneration
         ]
         pooler_out_size = [pixel_values.shape[0] * pixel_values.shape[1], self.config.vision_config.hidden_size]
         vision_out_buffer = []
-        for i in range(self.config.vision_config.num_hidden_layers + 2):
+        for _ in range(self.config.vision_config.num_hidden_layers + 2):
             vision_out_buffer.append(torch.empty(size=vision_out_size, dtype=torch.float32, device="cpu"))
         vision_out_buffer.insert(1, torch.empty(size=pooler_out_size, dtype=torch.float32, device="cpu"))
 

optimum/rbln/transformers/models/mistral/modeling_mistral.py

@@ -12,13 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from transformers import PretrainedConfig
 
 from ....utils import logging
 from ...models.decoderonly import (
     RBLNDecoderOnlyModel,
     RBLNDecoderOnlyModelForCausalLM,
-    RBLNDecoderOnlyModelForCausalLMConfig,
 )
 from .mistral_architecture import MistralWrapper
 
@@ -85,16 +83,6 @@ class RBLNMistralForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     _decoder_wrapper_cls = MistralWrapper
 
-    @classmethod
-    def _update_sliding_window_config(
-        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ):
-        rbln_config.cache_impl = "sliding_window"
-        rbln_config.sliding_window = model_config.sliding_window
-        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
-
-        return rbln_config
-
 
 class RBLNMistralModel(RBLNDecoderOnlyModel):
     """
@@ -103,13 +91,3 @@ class RBLNMistralModel(RBLNDecoderOnlyModel):
     """
 
     _decoder_wrapper_cls = MistralWrapper
-
-    @classmethod
-    def _update_sliding_window_config(
-        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ):
-        rbln_config.cache_impl = "sliding_window"
-        rbln_config.sliding_window = model_config.sliding_window
-        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
-
-        return rbln_config

optimum/rbln/transformers/models/opt/modeling_opt.py

@@ -69,7 +69,7 @@ class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         return layer
 
     @classmethod
-    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+    def _wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         for i in range(len(model.model.decoder.layers)):
             model.model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.model.decoder.layers[i])
 
@@ -95,7 +95,7 @@ class RBLNOPTModel(RBLNDecoderOnlyModel):
         return layer
 
     @classmethod
-    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+    def _wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         for i in range(len(model.decoder.layers)):
             model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.decoder.layers[i])
 

optimum/rbln/transformers/models/pegasus/modeling_pegasus.py

@@ -54,7 +54,7 @@ class RBLNPegasusForConditionalGeneration(RBLNModelForSeq2SeqLM):
     support_causal_attn = True
 
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNPegasusForConditionalGenerationConfig):
+    def _wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNPegasusForConditionalGenerationConfig):
         return PegasusWrapper(
             model, enc_max_seq_len=rbln_config.enc_max_seq_len, use_attention_mask=rbln_config.use_attention_mask
         )

optimum/rbln/transformers/models/pixtral/modeling_pixtral.py

@@ -229,7 +229,7 @@ class RBLNPixtralVisionModel(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNPixtralVisionModelConfig
     ) -> torch.nn.Module:
         wrapper_cfg = {
@@ -293,6 +293,18 @@ class RBLNPixtralVisionModel(RBLNModel):
         return_dict: bool = True,
         **kwargs,
     ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Forward pass for the RBLN-optimized Pixtral vision model.
+
+        Args:
+            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)) — The tensors corresponding to the input images. Pixel values can be obtained using PixtralImageProcessor. See PixtralImageProcessor.call() for details (PixtralProcessor uses PixtralImageProcessor for processing images).
+            image_sizes (torch.Tensor of shape (batch_size, 2), optional) — The sizes of the images in the batch, being (height, width) for each image.
+            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
+            return_dict (bool, optional) — Whether or not to return a ModelOutput instead of a plain tuple.
+
+        Returns:
+            BaseModelOutput or tuple(torch.FloatTensor)
+        """
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
         )

optimum/rbln/transformers/models/pixtral/pixtral_architecture.py

@@ -24,8 +24,8 @@ class PixtralAttention(nn.Module):
     def __init__(self, self_attention):
         super().__init__()
         self.original_model = self_attention
-        self.num_heads = getattr(self.original_model, "num_heads", None) or getattr(
-            self.original_model.config, "num_attention_heads"
+        self.num_heads = (
+            getattr(self.original_model, "num_heads", None) or self.original_model.config.num_attention_heads
         )
         self.head_dim = self.original_model.head_dim
         self.scaling = self.head_dim**-0.5

optimum/rbln/transformers/models/qwen2/modeling_qwen2.py

@@ -12,13 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from transformers import PretrainedConfig
 
 from ....utils import logging
 from ...models.decoderonly import (
     RBLNDecoderOnlyModel,
     RBLNDecoderOnlyModelForCausalLM,
-    RBLNDecoderOnlyModelForCausalLMConfig,
 )
 from .qwen2_architecture import QWEN2Wrapper
 
@@ -87,19 +85,6 @@ class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     _decoder_wrapper_cls = QWEN2Wrapper
 
-    @classmethod
-    def _update_sliding_window_config(
-        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ):
-        # https://github.com/huggingface/transformers/issues/35896
-        # There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
-        # we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
-
-        rbln_config.cache_impl = "sliding_window"
-        rbln_config.sliding_window = model_config.sliding_window
-        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
-        return rbln_config
-
 
 class RBLNQwen2Model(RBLNDecoderOnlyModel):
     """
@@ -108,16 +93,3 @@ class RBLNQwen2Model(RBLNDecoderOnlyModel):
     """
 
     _decoder_wrapper_cls = QWEN2Wrapper
-
-    @classmethod
-    def _update_sliding_window_config(
-        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ):
-        # https://github.com/huggingface/transformers/issues/35896
-        # There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
-        # we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
-
-        rbln_config.cache_impl = "sliding_window"
-        rbln_config.sliding_window = model_config.sliding_window
-        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
-        return rbln_config

optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -88,7 +88,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: "PreTrainedModel", rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig
     ):
         return Qwen2_5_VisionTransformerWrapper(model).eval()
@@ -111,10 +111,10 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
         model_config: "PretrainedConfig" = None,
         rbln_config: Optional[RBLNQwen2_5_VisionTransformerPretrainedModelConfig] = None,
     ) -> RBLNQwen2_5_VisionTransformerPretrainedModelConfig:
-        window_size = getattr(model_config, "window_size")
-        patch_size = getattr(model_config, "patch_size")
-        hidden_size = getattr(model_config, "hidden_size")
-        num_heads = getattr(model_config, "num_heads")
+        window_size = model_config.window_size
+        patch_size = model_config.patch_size
+        hidden_size = model_config.hidden_size
+        num_heads = model_config.num_heads
         head_dim = hidden_size // num_heads
         window_seq_len = (window_size // patch_size) ** 2
 
@@ -294,10 +294,10 @@ class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
             try:
                 ws_index = torch.searchsorted(self.max_seq_lens, window_padded_len).item()
                 max_seq_len = self.max_seq_lens[ws_index]
-            except Exception:
+            except Exception as e:
                 raise ValueError(
                     f"Required seq_len({window_padded_len}) is larger than available max_seq_lens({self.max_seq_lens.tolist()})."
-                )
+                ) from e
 
             # Padding for Window Attention Layers
             hidden_state_padded, cos_padded, sin_padded, window_attn_masks, window_valid_lengths = (
@@ -393,8 +393,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         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"):
         model.model.lm_head = model.lm_head
         model.lm_head = None
         del model.lm_head