optimum-rbln 0.9.3.post1__py3-none-any.whl

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

@@ -0,0 +1,144 @@
+# 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 math
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+import torch.nn as nn
+
+from ..decoderonly.decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyLayer,
+    DecoderOnlyModel,
+    DecoderOnlyWrapper,
+    apply_rotary_pos_emb_partial,
+    rotate_half,
+)
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel as MidmLMHeadModel
+
+
+def apply_rotary_to_tensor(tensor, cos, sin, rot_dim):
+    """Applies rotary position embedding to the specified dimension of the tensor."""
+    tensor_, tensor_pass = tensor[..., :rot_dim], tensor[..., rot_dim:]
+    tensor_embed = (tensor_ * cos) + (rotate_half(tensor_) * sin)
+    return torch.cat((tensor_embed, tensor_pass), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    """Applies Rotary Position Embedding to the query and key tensors."""
+    rot_dim = cos.shape[-1]
+    q_embed = apply_rotary_to_tensor(q, cos, sin, rot_dim)
+    k_embed = apply_rotary_to_tensor(k, cos, sin, rot_dim)
+    return q_embed, k_embed
+
+
+class MidmLMHeadModelWrapper(DecoderOnlyWrapper):
+    def get_rotary_emb(self, max_seq_len):
+        self.config.rope_theta = 10000
+        self.config.head_dim = self.config.n_embd // self.config.n_head
+        self.config.partial_rotary_factor = self.config.rotary_percentage
+        return super().get_rotary_emb(max_seq_len=max_seq_len)
+
+    def get_rbln_attn_class(self):
+        return MidmAttention
+
+    def get_rbln_layer_class(self):
+        return MidmLayer
+
+    def get_rbln_model_class(self):
+        return MidmModel
+
+    def get_model_layer(self, causal_lm: "MidmLMHeadModel"):
+        return causal_lm.transformer
+
+    def get_decoder_layers(self, causal_lm: "MidmLMHeadModel"):
+        return causal_lm.transformer.h
+
+
+class MidmModel(DecoderOnlyModel):
+    def get_layernorm1p(self, module: nn.LayerNorm):
+        def layernorm1p(input: torch.Tensor):
+            """Applies Layer Normalization with a slight modification on the weights."""
+            return torch.nn.functional.layer_norm(
+                input, module.normalized_shape, module.weight + 1, module.bias, module.eps
+            )
+
+        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
+
+    def get_embedding(self) -> nn.Embedding:
+        return self._original_mod.wte
+
+    def get_pos_embedding(self) -> nn.Embedding:
+        return self._original_mod.wpe
+
+
+class MidmLayer(DecoderOnlyLayer):
+    def get_layernorm1p(self, module: nn.LayerNorm):
+        def layernorm1p(input: torch.Tensor):
+            """Applies Layer Normalization with a slight modification on the weights."""
+            return torch.nn.functional.layer_norm(
+                input, module.normalized_shape, module.weight + 1, module.bias, module.eps
+            )
+
+        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
+
+    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
+
+
+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 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
+
+    def get_attn_scale(self):
+        scale = 1.0
+        if self._original_mod.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:
+            scale /= 1 + self.layer_idx
+
+        return scale
+
+    def apply_rotary_pos_embed(self, query_states, key_states, cos, sin):
+        return apply_rotary_pos_emb_partial(query_states, key_states, cos, sin, ndim=cos.shape[-1])

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

@@ -0,0 +1,144 @@
+# 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 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
+
+
+logger = logging.get_logger(__name__)
+
+
+class RBLNMidmLMHeadModel(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The MIDM Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based MidmForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers MidmForCausalLM model into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNMidmLMHeadModelConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNMidmLMHeadModelConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNMidmLMHeadModelConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNMidmLMHeadModel
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNMidmLMHeadModel.from_pretrained(
+            "KT-AI/midm-bitext-S-7B-inst-v1",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNMidmLMHeadModel.from_pretrained(
+            "KT-AI/midm-bitext-S-7B-inst-v1",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNMidmLMHeadModelConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNMidmLMHeadModelConfig
+
+        config = RBLNMidmLMHeadModelConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNMidmLMHeadModel.from_pretrained(
+            "KT-AI/midm-bitext-S-7B-inst-v1",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = MidmLMHeadModelWrapper
+    _hf_class = AutoModelForCausalLM
+    _supports_cache_class = True
+
+    @classmethod
+    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):
+            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
+
+        val = getattr(GenerationMixin, __name)
+
+        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
+            return redirect(val)
+        return val

optimum/rbln/transformers/models/mistral/__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_mistral import RBLNMistralForCausalLMConfig, RBLNMistralModelConfig
+from .modeling_mistral import RBLNMistralForCausalLM, RBLNMistralModel

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

@@ -0,0 +1,50 @@
+# 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 ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNMistralForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN Mistral models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNMistralForCausalLM, RBLNMistralForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNMistralForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNMistralForCausalLM.from_pretrained(
+        "mistralai/Mistral-7B-v0.1",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """
+
+
+class RBLNMistralModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Mistral models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

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

@@ -0,0 +1,19 @@
+# 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 ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper
+
+
+class MistralWrapper(DecoderOnlyWrapper):
+    pass

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

@@ -0,0 +1,115 @@
+# 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 transformers import PretrainedConfig
+
+from ....utils import logging
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModel,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
+from .mistral_architecture import MistralWrapper
+
+
+logger = logging.get_logger(__name__)
+
+
+class RBLNMistralForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Mistral Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based MistralForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers MistralForCausalLM model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNMistralForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNMistralForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNMistralForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNMistralForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNMistralForCausalLM.from_pretrained(
+            "mistralai/Mistral-7B-v0.1",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNMistralForCausalLM.from_pretrained(
+            "mistralai/Mistral-7B-v0.1",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+        # Using a RBLNMistralForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNMistralForCausalLMConfig
+
+        config = RBLNMistralForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNMistralForCausalLM.from_pretrained(
+            "mistralai/Mistral-7B-v0.1",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _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(
+        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/__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_opt import RBLNOPTForCausalLMConfig, RBLNOPTModelConfig
+from .modeling_opt import RBLNOPTForCausalLM, RBLNOPTModel

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

@@ -0,0 +1,29 @@
+# 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 ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
+
+
+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

@@ -0,0 +1,102 @@
+# 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 torch.nn as nn
+from transformers import PreTrainedModel
+
+from ....utils import logging
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from .opt_architecture import OPTWrapper
+
+
+logger = logging.get_logger(__name__)
+
+
+class MLP(nn.Module):
+    def __init__(self, fc1, fc2, activation_fn):
+        super(MLP, self).__init__()
+        self.fc1 = fc1
+        self.fc2 = fc2
+        self.activation_fn = activation_fn
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.activation_fn(x)
+        x = self.fc2(x)
+        return x
+
+
+class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The OPT Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based OPTForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers OPTForCausalLM model into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNOPTForCausalLM`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNOPTForCausalLM`] or a dictionary conforming to its structure.
+
+    See the [`RBLNOPTForCausalLM`] class for all available configuration options.
+    """
+
+    _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.model.decoder.layers)):
+            model.model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.model.decoder.layers[i])
+
+        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()