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

optimum/rbln/transformers/models/midm/midm_architecture.py

@@ -123,7 +123,10 @@ class MidmAttention(DecoderOnlyAttention):
         self.split_size = self._original_mod.split_size
         self.num_key_value_heads = self._original_mod.num_heads
 
-    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def projection(self, hidden_states, lora_int_id) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if lora_int_id is not None:
+            raise NotImplementedError("LoRA is not supported for MidmAttention")
+
         query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
         return query_states, key_states, value_states
 

optimum/rbln/transformers/models/midm/modeling_midm.py

@@ -13,11 +13,13 @@
 # limitations under the License.
 
 import inspect
-from typing import Any, Callable
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Union
 
 from transformers import AutoModelForCausalLM
 from transformers.generation.utils import GenerationMixin
 
+from ....configuration_utils import RBLNModelConfig
 from ....utils import logging
 from ..decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .midm_architecture import MidmLMHeadModelWrapper
@@ -91,9 +93,45 @@ class RBLNMidmLMHeadModel(RBLNDecoderOnlyModelForCausalLM):
     _supports_cache_class = True
 
     @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        kwargs.setdefault("trust_remote_code", True)
-        return super().from_pretrained(*args, **kwargs)
+    def from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        *,
+        export: Optional[bool] = None,
+        rbln_config: Optional[Union[Dict, RBLNModelConfig]] = None,
+        trust_remote_code: Optional[bool] = None,
+        **kwargs: Any,
+    ):
+        """
+        The `from_pretrained()` function is utilized in its standard form as in the HuggingFace transformers library.
+        User can use this function to load a pre-trained model from the HuggingFace library and convert it to a RBLN model to be run on RBLN NPUs.
+
+        Args:
+            model_id (Union[str, Path]): The model id of the pre-trained model to be loaded.
+                It can be downloaded from the HuggingFace model hub or a local path, or a model id of a compiled model using the RBLN Compiler.
+            export (Optional[bool]): A boolean flag to indicate whether the model should be compiled.
+                If None, it will be determined based on the existence of the compiled model files in the model_id.
+            rbln_config (Optional[Union[Dict, RBLNModelConfig]]): Configuration for RBLN model compilation and runtime.
+                This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNMidmLMHeadModelConfig` for Mi:dm models).
+                For detailed configuration options, see the specific model's configuration class documentation.
+            trust_remote_code (bool): Whether or not to trust the remote code when loading a model from the Hub.
+            kwargs: Additional keyword arguments. Arguments with the prefix `rbln_` are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+
+        Returns:
+            (RBLNModel): A RBLN model instance ready for inference on RBLN NPU devices.
+        """
+
+        if trust_remote_code is not None:
+            kwargs["trust_remote_code"] = trust_remote_code
+        elif "trust_remote_code" not in kwargs:
+            kwargs["trust_remote_code"] = True
+
+        return super().from_pretrained(
+            model_id=model_id,
+            export=export,
+            rbln_config=rbln_config,
+            **kwargs,
+        )
 
     def __getattr__(self, __name: str) -> Any:
         def redirect(func):

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_mistral import RBLNMistralForCausalLMConfig
-from .modeling_mistral import RBLNMistralForCausalLM
+from .configuration_mistral import RBLNMistralForCausalLMConfig, RBLNMistralModelConfig
+from .modeling_mistral import RBLNMistralForCausalLM, RBLNMistralModel

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

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

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

@@ -15,5 +15,5 @@
 from ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper
 
 
-class MistralForCausalLMWrapper(DecoderOnlyWrapper):
+class MistralWrapper(DecoderOnlyWrapper):
     pass

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

@@ -15,8 +15,12 @@
 from transformers import PretrainedConfig
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
-from .mistral_architecture import MistralForCausalLMWrapper
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModel,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
+from .mistral_architecture import MistralWrapper
 
 
 logger = logging.get_logger(__name__)
@@ -79,7 +83,26 @@ class RBLNMistralForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         ```
     """
 
-    _decoder_wrapper_cls = MistralForCausalLMWrapper
+    _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):
+    """
+    The Mistral Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    """
+
+    _decoder_wrapper_cls = MistralWrapper
 
     @classmethod
     def _update_sliding_window_config(

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_opt import RBLNOPTForCausalLMConfig
-from .modeling_opt import RBLNOPTForCausalLM
+from .configuration_opt import RBLNOPTForCausalLMConfig, RBLNOPTModelConfig
+from .modeling_opt import RBLNOPTForCausalLM, RBLNOPTModel

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNOPTForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -20,3 +20,10 @@ class RBLNOPTForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     Configuration class for OPT causal language model.
     Inherits from RBLNDecoderOnlyModelForCausalLMConfig with no additional parameters.
     """
+
+
+class RBLNOPTModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for OPT model.
+    Inherits from RBLNDecoderOnlyModelConfig with no additional parameters.
+    """

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

@@ -16,7 +16,7 @@ import torch.nn as nn
 from transformers import PreTrainedModel
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from ...models.decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 from .opt_architecture import OPTWrapper
 
@@ -69,22 +69,34 @@ class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         return layer
 
     @classmethod
-    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
-        wrapper_cfg = {
-            "max_seq_len": rbln_config.max_seq_len,
-            "attn_impl": rbln_config.attn_impl,
-            "kvcache_partition_len": rbln_config.kvcache_partition_len,
-            "kvcache_block_size": rbln_config.kvcache_block_size,
-            "use_rotary_emb": cls._use_rotary_emb,
-            "use_attention_mask": rbln_config.use_attention_mask,
-            "use_position_ids": rbln_config.use_position_ids,
-            "use_inputs_embeds": rbln_config.use_inputs_embeds,
-            "cache_impl": rbln_config.cache_impl,
-            "sliding_window": rbln_config.sliding_window,
-            "sliding_window_layers": rbln_config.sliding_window_layers,
-        }
-
+    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])
 
-        return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()
+        return cls._decoder_wrapper_cls(model, rbln_config=rbln_config, use_rotary_emb=cls._use_rotary_emb).eval()
+
+
+class RBLNOPTModel(RBLNDecoderOnlyModel):
+    """
+    The OPT Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    """
+
+    _decoder_wrapper_cls = OPTWrapper
+    _use_rotary_emb = False
+
+    def modify_opt_decoder_layer(layer):
+        mlp = MLP(layer.fc1, layer.fc2, layer.activation_fn)
+        layer.mlp = mlp
+        del layer.fc1
+        del layer.fc2
+        del layer.activation_fn
+
+        return layer
+
+    @classmethod
+    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])
+
+        return cls._decoder_wrapper_cls(model, rbln_config=rbln_config, use_rotary_emb=cls._use_rotary_emb).eval()

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

@@ -40,11 +40,11 @@ class OPTWrapper(DecoderOnlyWrapper):
     def get_rbln_model_class(self):
         return OPTModel
 
-    def get_model_layer(self, causal_lm: "OPTForCausalLM"):
-        return causal_lm.model.decoder
+    def get_model_layer(self, model: "OPTForCausalLM"):
+        return model.model.decoder if self.is_causal_lm else model.decoder
 
-    def get_decoder_layers(self, causal_lm: "OPTForCausalLM"):
-        return causal_lm.model.decoder.layers
+    def get_decoder_layers(self, model: "OPTForCausalLM"):
+        return model.model.decoder.layers if self.is_causal_lm else model.decoder.layers
 
 
 class OPTAttention(DecoderOnlyAttention):

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

@@ -0,0 +1,17 @@
+# 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 ....ops import paged_attn_decode, paged_causal_attn_decode
+from .configuration_pegasus import RBLNPegasusForConditionalGenerationConfig, RBLNPegasusModelConfig
+from .modeling_pegasus import RBLNPegasusForConditionalGeneration, RBLNPegasusModel

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

@@ -0,0 +1,38 @@
+# 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_generic import RBLNTransformerEncoderForFeatureExtractionConfig
+from ..seq2seq import RBLNModelForSeq2SeqLMConfig
+
+
+class RBLNPegasusModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
+    """
+    Configuration class for RBLNPegasusModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized PEGASUS models for feature extraction tasks.
+    """
+
+    rbln_model_input_names = ["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask"]
+
+
+class RBLNPegasusForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
+    """
+    Configuration class for RBLNPegasusForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized PEGASUS models for conditional text generation tasks.
+    """
+
+    support_paged_attention = True

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

@@ -0,0 +1,71 @@
+# 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.
+
+import inspect
+from typing import TYPE_CHECKING, Any, Callable
+
+from transformers import PegasusForConditionalGeneration, PreTrainedModel
+
+from ....utils.logging import get_logger
+from ...modeling_generic import RBLNTransformerEncoderForFeatureExtraction
+from ...models.seq2seq import RBLNModelForSeq2SeqLM
+from .configuration_pegasus import RBLNPegasusForConditionalGenerationConfig
+from .pegasus_architecture import PegasusWrapper
+
+
+logger = get_logger()
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+
+class RBLNPegasusModel(RBLNTransformerEncoderForFeatureExtraction):
+    """
+    RBLN optimized PEGASUS model for feature extraction tasks.
+
+    This class provides hardware-accelerated inference for PEGASUS encoder models
+    on RBLN devices, optimized for feature extraction use cases.
+    """
+
+    rbln_model_input_names = ["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask"]
+
+
+class RBLNPegasusForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    """
+    RBLN optimized PEGASUS model for conditional text generation tasks.
+
+    This class provides hardware-accelerated inference for PEGASUS models
+    on RBLN devices, supporting sequence-to-sequence generation tasks
+    such as summarization, translation, and text generation.
+    """
+
+    support_causal_attn = True
+
+    @classmethod
+    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
+        )
+
+    def __getattr__(self, __name: str) -> Any:
+        def redirect(func):
+            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
+
+        val = getattr(PegasusForConditionalGeneration, __name)
+
+        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
+            return redirect(val)
+
+        return val

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

@@ -0,0 +1,161 @@
+# 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 Tuple
+
+import torch
+from torch import nn
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.utils import logging
+
+from ..seq2seq.seq2seq_architecture import (
+    Seq2SeqCrossAttention,
+    Seq2SeqDecoder,
+    Seq2SeqDecoderLayer,
+    Seq2SeqDecoderWrapper,
+    Seq2SeqEncoderWrapper,
+    Seq2SeqForConditionalGeneration,
+    Seq2SeqSelfAttention,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+class PegasusWrapper:
+    def __init__(self, model: nn.Module, enc_max_seq_len: int, use_attention_mask: bool):
+        self.encoder = Seq2SeqEncoderWrapper(model, enc_max_seq_len)
+        self.decoder = PegasusDecoderWrapper(model, use_attention_mask=use_attention_mask)
+
+
+class PegasusDecoderWrapper(Seq2SeqDecoderWrapper):
+    def convert_to_rbln_conditional_generation(self, model: nn.Module):
+        new_layers = []
+        for layer in model.get_decoder().layers:
+            self_attn = PegasusSelfAttention(layer.self_attn, use_attention_mask=self.use_attention_mask)
+            cross_attn = PegasusCrossAttention(layer.encoder_attn)
+            new_layers.append(PegasusDecoderLayer(layer, self_attn, cross_attn))
+
+        decoder_model = PegasusDecoder(model.get_decoder(), new_layers)
+        new_model = PegasusForConditionalGeneration(model, decoder_model)
+
+        return new_model
+
+
+class PegasusForConditionalGeneration(Seq2SeqForConditionalGeneration):
+    pass
+
+
+class PegasusDecoder(Seq2SeqDecoder):
+    has_pos_emb = True
+
+    def __post_init__(self):
+        self.embed_positions = self._original_mod.embed_positions
+        self.embed_scale = getattr(self._original_mod, "embed_scale", None)
+        self.final_layer_norm = getattr(self._original_mod, "layer_norm", None)
+
+    def prepare_attn_mask(self, attention_mask, encoder_attention_mask, **kwargs):
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, None, None, :]
+        encoder_attention_mask = _prepare_4d_attention_mask(encoder_attention_mask, torch.float32, tgt_len=1)
+
+        return attention_mask, encoder_attention_mask
+
+    def apply_position_embedding(self, inputs_embeds, cache_position):
+        hidden_all = []
+        for i in range(inputs_embeds.shape[0]):
+            positions_idx = cache_position[i]
+            position_weight = self.embed_positions.weight
+            position = position_weight[positions_idx]
+            batch_hidden = position + inputs_embeds[i]
+            hidden_all.append(batch_hidden)
+        hidden_states = torch.stack(hidden_all, dim=0)
+
+        return hidden_states
+
+    def get_embedding(self):
+        if self.embed_scale is not None:
+            return lambda x: self.embed_tokens(x) * self.embed_scale
+        else:
+            return self.embed_tokens
+
+
+class PegasusLayerFF(nn.Module):
+    def __init__(self, decoder_layer):
+        super().__init__()
+        self.fc1 = decoder_layer.fc1
+        self.fc2 = decoder_layer.fc2
+        self.activation_fn = decoder_layer.activation_fn
+        self.layer_norm = decoder_layer.final_layer_norm
+
+    def forward(self, hidden_states):
+        # Residual Connection
+        residual = hidden_states
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class PegasusDecoderLayer(Seq2SeqDecoderLayer):
+    def __post_init__(self):
+        self.self_attn_layer_norm = self._original_mod.self_attn_layer_norm
+        self.encoder_attn = self._original_mod.encoder_attn
+        self.encoder_attn_layer_norm = self._original_mod.encoder_attn_layer_norm
+        self.ff_layer = PegasusLayerFF(self._original_mod)
+
+    def pre_self_attn_layer_norm(self, hidden_states):
+        return self.self_attn_layer_norm(hidden_states)
+
+    def post_self_attn_layer_norm(self, hidden_states):
+        return hidden_states
+
+    def pre_cross_attn_layer_norm(self, hidden_states):
+        return self.encoder_attn_layer_norm(hidden_states)
+
+    def post_cross_attn_layer_norm(self, hidden_states):
+        return hidden_states
+
+
+class PegasusSelfAttention(Seq2SeqSelfAttention):
+    def __post_init__(self, use_attention_mask: bool = True):
+        self.q_proj = self._original_mod.q_proj
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.out_proj = self._original_mod.out_proj
+        self.num_heads = self._original_mod.num_heads
+        self.head_dim = self._original_mod.embed_dim // self._original_mod.num_heads
+        self.scaling = self.head_dim**-0.5
+        if use_attention_mask:
+            self.attn_decode = torch.ops.rbln_custom_ops.paged_attn_decode
+        else:
+            self.attn_decode = torch.ops.rbln_custom_ops.paged_causal_attn_decode
+
+    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        query_states = self.q_proj(hidden_states) * self.scaling
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        return query_states, key_states, value_states
+
+
+class PegasusCrossAttention(Seq2SeqCrossAttention):
+    def __post_init__(self):
+        self.q_proj = self._original_mod.q_proj
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.out_proj = self._original_mod.out_proj
+        self.num_heads = self._original_mod.num_heads
+        self.head_dim = self._original_mod.embed_dim // self._original_mod.num_heads
+        self.embed_dim = self._original_mod.embed_dim

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_phi import RBLNPhiForCausalLMConfig
-from .modeling_phi import RBLNPhiForCausalLM
+from .configuration_phi import RBLNPhiForCausalLMConfig, RBLNPhiModelConfig
+from .modeling_phi import RBLNPhiForCausalLM, RBLNPhiModel

optimum/rbln/transformers/models/phi/configuration_phi.py

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

optimum/rbln/transformers/models/phi/modeling_phi.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .phi_architecture import PhiWrapper
 
 
@@ -81,3 +81,12 @@ class RBLNPhiForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = PhiWrapper
+
+
+class RBLNPhiModel(RBLNDecoderOnlyModel):
+    """
+    The Phi Model transformer without a language modeling head.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    """
+
+    _decoder_wrapper_cls = PhiWrapper

optimum/rbln/transformers/models/phi/phi_architecture.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 from transformers import PhiForCausalLM
@@ -27,7 +27,7 @@ from ..decoderonly.decoderonly_architecture import (
 
 
 if TYPE_CHECKING:
-    from transformers import PhiForCausalLM
+    from transformers import PhiForCausalLM, PhiModel
 
 
 class PhiWrapper(DecoderOnlyWrapper):
@@ -40,11 +40,11 @@ class PhiWrapper(DecoderOnlyWrapper):
     def get_rbln_model_class(self):
         return PhiModel
 
-    def get_model_layer(self, causal_lm: "PhiForCausalLM"):
-        return causal_lm.model
+    def get_model_layer(self, model: Union["PhiForCausalLM", "PhiModel"]):
+        return model.model if self.is_causal_lm else model
 
-    def get_decoder_layers(self, causal_lm: "PhiForCausalLM"):
-        return causal_lm.model.layers
+    def get_decoder_layers(self, model: Union["PhiForCausalLM", "PhiModel"]):
+        return model.model.layers if self.is_causal_lm else model.layers
 
 
 class PhiAttention(DecoderOnlyAttention):
@@ -56,7 +56,10 @@ class PhiAttention(DecoderOnlyAttention):
         self.qk_layernorm = self._original_mod.qk_layernorm
         self.rotary_ndims = self._original_mod.rotary_ndims
 
-    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def projection(self, hidden_states, lora_int_id) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if lora_int_id is not None:
+            raise NotImplementedError("LoRA is not supported for PhiAttention")
+
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
@@ -84,6 +87,7 @@ class PhiLayer(DecoderOnlyLayer):
         cos: Optional[torch.Tensor] = None,
         sin: Optional[torch.Tensor] = None,
         block_tables: Optional[torch.Tensor] = None,
+        lora_int_id: Optional[torch.Tensor] = None,
     ):
         residual = hidden_states