optimum-rbln 0.9.3rc0__py3-none-any.whl → 0.9.5a4__py3-none-any.whl

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):
@@ -109,7 +116,6 @@ class RBLNLlavaForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixi
     RBLNLlavaForConditionalGeneration is a multi-modal model that combines vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
     This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
-
     Important Note:
         This model includes a Large Language Model (LLM) as a submodule. For optimal performance, it is highly recommended to use
         tensor parallelism for the language model. This can be achieved by using the `rbln_config` parameter in the
@@ -175,9 +181,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 +198,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 +212,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 +225,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 +336,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 +355,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/midm/midm_architecture.py

@@ -71,6 +71,12 @@ class MidmLMHeadModelWrapper(DecoderOnlyWrapper):
 
 
 class MidmModel(DecoderOnlyModel):
+    def __init__(self, model, layers, rbln_config, use_learned_pos_emb=None, use_rotary_emb=True):
+        super().__init__(
+            model, layers, rbln_config, use_learned_pos_emb=use_learned_pos_emb, use_rotary_emb=use_rotary_emb
+        )
+        self.use_layernorm1p = getattr(model, "use_layernorm1p", False)
+
     def get_layernorm1p(self, module: nn.LayerNorm):
         def layernorm1p(input: torch.Tensor):
             """Applies Layer Normalization with a slight modification on the weights."""
@@ -81,19 +87,22 @@ class MidmModel(DecoderOnlyModel):
         return layernorm1p
 
     def get_last_layernorm(self) -> nn.LayerNorm:
-        if self._original_mod.use_layernorm1p:
-            return self.get_layernorm1p(self._original_mod.ln_f)
-        else:
-            return self._original_mod.ln_f
+        if self.use_layernorm1p:
+            return self.get_layernorm1p(self.norm)
+        return self.norm
 
     def get_embedding(self) -> nn.Embedding:
-        return self._original_mod.wte
+        return self.embed_tokens
 
     def get_pos_embedding(self) -> nn.Embedding:
-        return self._original_mod.wpe
+        return self.embed_positions
 
 
 class MidmLayer(DecoderOnlyLayer):
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config=None):
+        super().__init__(layer, self_attn, lora_config)
+        self.use_layernorm1p = getattr(layer, "use_layernorm1p", False)
+
     def get_layernorm1p(self, module: nn.LayerNorm):
         def layernorm1p(input: torch.Tensor):
             """Applies Layer Normalization with a slight modification on the weights."""
@@ -104,24 +113,22 @@ class MidmLayer(DecoderOnlyLayer):
         return layernorm1p
 
     def get_pre_attention_layernorm(self) -> nn.LayerNorm:
-        if self._original_mod.use_layernorm1p:
-            return self.get_layernorm1p(self._original_mod.ln_1)
-        else:
-            return self._original_mod.ln_1
+        if self.use_layernorm1p:
+            return self.get_layernorm1p(self.pre_attention_layernorm)
+        return self.pre_attention_layernorm
 
     def get_post_attention_layernorm(self) -> nn.LayerNorm:
-        if self._original_mod.use_layernorm1p:
-            return self.get_layernorm1p(self._original_mod.ln_2)
-        else:
-            return self._original_mod.ln_2
+        if self.use_layernorm1p:
+            return self.get_layernorm1p(self.post_attention_layernorm)
+        return self.post_attention_layernorm
 
 
 class MidmAttention(DecoderOnlyAttention):
-    def __post_init__(self):
-        self.c_attn = self._original_mod.c_attn
-        self.o_proj = self._original_mod.c_proj
-        self.split_size = self._original_mod.split_size
-        self.num_key_value_heads = self._original_mod.num_heads
+    def __post_init__(self, self_attn):
+        self.c_attn = self_attn.c_attn
+        self.o_proj = self_attn.c_proj
+        self.split_size = self_attn.split_size
+        self.num_key_value_heads = self_attn.num_heads
 
     def projection(self, hidden_states, lora_int_id) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         if lora_int_id is not None:
@@ -130,12 +137,12 @@ class MidmAttention(DecoderOnlyAttention):
         query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
         return query_states, key_states, value_states
 
-    def get_attn_scale(self):
+    def get_attn_scale(self, self_attn):
         scale = 1.0
-        if self._original_mod.scale_attn_weights:
+        if self_attn.scale_attn_weights:
             scale /= math.sqrt(self.head_dim)
 
-        if self._original_mod.scale_attn_by_inverse_layer_idx and not self._original_mod.scale_qk_by_inverse_layer_idx:
+        if self_attn.scale_attn_by_inverse_layer_idx:
             scale /= 1 + self.layer_idx
 
         return scale

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/opt/opt_architecture.py

@@ -14,14 +14,7 @@
 
 from typing import TYPE_CHECKING
 
-import torch.nn as nn
-
-from ...models.decoderonly.decoderonly_architecture import (
-    DecoderOnlyAttention,
-    DecoderOnlyLayer,
-    DecoderOnlyModel,
-    DecoderOnlyWrapper,
-)
+from ...models.decoderonly.decoderonly_architecture import DecoderOnlyWrapper
 
 
 if TYPE_CHECKING:
@@ -31,44 +24,8 @@ if TYPE_CHECKING:
 class OPTWrapper(DecoderOnlyWrapper):
     _use_learned_pos_emb = True
 
-    def get_rbln_attn_class(self):
-        return OPTAttention
-
-    def get_rbln_layer_class(self):
-        return OPTDecoderLayer
-
-    def get_rbln_model_class(self):
-        return OPTModel
-
     def get_model_layer(self, model: "OPTForCausalLM"):
         return model.model.decoder if self.is_causal_lm else model.decoder
 
     def get_decoder_layers(self, model: "OPTForCausalLM"):
         return model.model.decoder.layers if self.is_causal_lm else model.decoder.layers
-
-
-class OPTAttention(DecoderOnlyAttention):
-    def __post_init__(self):
-        self.k_proj = self._original_mod.k_proj
-        self.v_proj = self._original_mod.v_proj
-        self.q_proj = self._original_mod.q_proj
-        self.o_proj = self._original_mod.out_proj
-
-
-class OPTModel(DecoderOnlyModel):
-    def get_embedding(self) -> nn.Embedding:
-        return self._original_mod.embed_tokens
-
-    def get_pos_embedding(self):
-        return self._original_mod.embed_positions
-
-    def get_last_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.final_layer_norm
-
-
-class OPTDecoderLayer(DecoderOnlyLayer):
-    def get_pre_attention_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.self_attn_layer_norm
-
-    def get_post_attention_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.final_layer_norm

optimum/rbln/transformers/models/paligemma/__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_paligemma import RBLNPaliGemmaForConditionalGenerationConfig, RBLNPaliGemmaModelConfig
+from .modeling_paligemma import RBLNPaliGemmaForConditionalGeneration, RBLNPaliGemmaModel

optimum/rbln/transformers/models/paligemma/configuration_paligemma.py

@@ -0,0 +1,129 @@
+# 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 RBLNPaliGemmaForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNPaliGemmaForConditionalGenerationConfig.
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized PaliGemma models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
+    submodules = ["vision_tower", "language_model"]
+    _allow_no_compile_cfgs = True
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        output_hidden_states: Optional[bool] = 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.
+            output_hidden_states (Optional[bool]): Whether to output the hidden states of the decoder. Defaults to False.
+            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 PaliGemma vision tower. It will be set to 1.")
+
+        self.output_hidden_states = output_hidden_states or False
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower,
+            batch_size=1,  # vision_tower batch_size is always 1 in PaliGemma
+            force_kwargs=True,
+        )
+        self.language_model = self.initialize_submodule_config(
+            submodule_config=language_model,
+            batch_size=batch_size,
+            use_position_ids=True,
+            use_attention_mask=True,
+            use_inputs_embeds=True,
+        )
+
+
+class RBLNPaliGemmaModelConfig(RBLNModelConfig):
+    submodules = ["vision_tower", "language_model"]
+    _allow_no_compile_cfgs = True
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        output_hidden_states: Optional[bool] = 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.
+            output_hidden_states (Optional[bool]): Whether to output the hidden states of the decoder. Defaults to False.
+            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 PaliGemma vision tower. It will be set to 1.")
+
+        self.output_hidden_states = output_hidden_states or False
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower,
+            batch_size=1,  # vision_tower batch_size is always 1 in PaliGemma
+            force_kwargs=True,
+        )
+
+        self.language_model = self.initialize_submodule_config(
+            submodule_config=language_model,
+            batch_size=batch_size,
+            use_position_ids=True,
+            use_attention_mask=True,
+            use_inputs_embeds=True,
+            output_hidden_states=output_hidden_states,
+        )