optimum-rbln 0.9.4a2__py3-none-any.whl → 0.9.5a4__py3-none-any.whl

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

@@ -16,7 +16,6 @@ import copy
 from typing import Optional, Tuple, Union
 
 import torch
-from transformers.models.gemma3.modeling_gemma3 import Gemma3RMSNorm
 
 from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
@@ -95,16 +94,18 @@ class Gemma3TextModel(DecoderOnlyModel):
         else:
             seq_positions = cache_position[:, :1]
 
-        sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position)
+        cache_seq_len, cache_offset, swa_attn_mask = self.get_swa_custom_op_args(position_ids, query_position)
+        sliding_cache_pos = (cache_seq_len, cache_offset)
 
         all_hidden_states = () if output_hidden_states else None
         for layer_idx, layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
             is_sliding = True if layer_idx in self.sliding_window_layers else False
+            is_sliding_decode = is_sliding and self.phase == "decode"
             hidden_states = layer(
                 hidden_states=hidden_states,
-                attention_mask=attention_mask,
+                attention_mask=swa_attn_mask if is_sliding_decode else attention_mask,
                 seq_positions=sliding_cache_pos if is_sliding else seq_positions,
                 past_key_values=past_key_values,
                 cos=cos_local if is_sliding else cos_global,
@@ -120,11 +121,8 @@ class Gemma3TextModel(DecoderOnlyModel):
 
 
 class Gemma3DecoderLayer(DecoderOnlyLayer):
-    def get_pre_feedforward_layernorm(self) -> Gemma3RMSNorm:
-        return self._original_mod.pre_feedforward_layernorm
-
-    def get_post_feedforward_layernorm(self) -> Gemma3RMSNorm:
-        return self._original_mod.post_feedforward_layernorm
+    _PRE_FF_LAYERNORM_ATTRS = ["pre_feedforward_layernorm"]
+    _POST_FF_LAYERNORM_ATTRS = ["post_feedforward_layernorm"]
 
     def forward(
         self,
@@ -164,13 +162,13 @@ class Gemma3DecoderLayer(DecoderOnlyLayer):
 
 
 class Gemma3Attention(DecoderOnlyAttention):
-    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.o_proj = self._original_mod.o_proj
-        self.q_norm = self._original_mod.q_norm
-        self.k_norm = self._original_mod.k_norm
-
-    def get_attn_scale(self):
-        return self._original_mod.config.query_pre_attn_scalar**-0.5
+    def __post_init__(self, self_attn):
+        self.q_proj = self_attn.q_proj
+        self.k_proj = self_attn.k_proj
+        self.v_proj = self_attn.v_proj
+        self.o_proj = self_attn.o_proj
+        self.q_norm = self_attn.q_norm
+        self.k_norm = self_attn.k_norm
+
+    def get_attn_scale(self, self_attn):
+        return self_attn.config.query_pre_attn_scalar**-0.5

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

@@ -325,7 +325,7 @@ class RBLNGemma3ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMix
                         batch_size,
                         inputs_embeds.shape[1],
                         self.config.text_config.hidden_size,
-                        dtype=self.rbln_config.torch_dtype,
+                        dtype=self.rbln_config.dtype,
                     )
                     for _ in range(self.config.text_config.num_hidden_layers + 1)
                 )

optimum/rbln/transformers/models/gpt2/gpt2_architecture.py

@@ -20,8 +20,6 @@ import torch.nn as nn
 
 from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
-    DecoderOnlyLayer,
-    DecoderOnlyModel,
     DecoderOnlyWrapper,
 )
 
@@ -34,12 +32,6 @@ class GPT2Wrapper(DecoderOnlyWrapper):
     def get_rbln_attn_class(self):
         return GPT2Attention
 
-    def get_rbln_layer_class(self):
-        return GPT2Layer
-
-    def get_rbln_model_class(self):
-        return GPT2Model
-
     def get_attn_layer(self, layer: nn.Module):
         return layer.attn
 
@@ -50,30 +42,12 @@ class GPT2Wrapper(DecoderOnlyWrapper):
         return model.transformer.h if self.is_causal_lm else model.h
 
 
-class GPT2Model(DecoderOnlyModel):
-    def get_last_layernorm(self) -> nn.LayerNorm:
-        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 GPT2Layer(DecoderOnlyLayer):
-    def get_pre_attention_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.ln_1
-
-    def get_post_attention_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.ln_2
-
-
 class GPT2Attention(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
+    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:
@@ -82,12 +56,12 @@ class GPT2Attention(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:
+        if self_attn.scale_attn_by_inverse_layer_idx:
             scale /= 1 + self.layer_idx
 
         return scale

optimum/rbln/transformers/models/gpt_oss/__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_gpt_oss import RBLNGptOssForCausalLMConfig
+from .modeling_gpt_oss import RBLNGptOssForCausalLM

optimum/rbln/transformers/models/gpt_oss/configuration_gpt_oss.py

@@ -0,0 +1,41 @@
+# 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 RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNGptOssForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN GPT-OSS models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNGptOssForCausalLM, RBLNGptOssForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNGptOssForCausalLMConfig(
+        batch_size=1,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNGptOssForCausalLM.from_pretrained(
+        "openai/gpt-oss-20b",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/gpt_oss/gpt_oss_architecture.py

@@ -0,0 +1,122 @@
+# 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 Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ..decoderonly.configuration_decoderonly import RBLNLoRAConfig
+from ..decoderonly.decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyLayer,
+    DecoderOnlyWrapper,
+)
+
+
+class RBLNGptOssWrapper(DecoderOnlyWrapper):
+    def get_rbln_layer_class(self):
+        return RBLNGptOssLayer
+
+
+class RBLNGptOssLayer(DecoderOnlyLayer):
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config: Optional[RBLNLoRAConfig] = None):
+        super().__init__(layer, self_attn, lora_config)
+        self.mlp = RBLNGptOssMLP(layer.mlp)
+
+    def get_mlp(self) -> nn.Module:
+        return self.mlp
+
+
+class RBLNGptOssTopKRouter(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.weight = model.weight
+        self.bias = model.bias
+
+    def forward(self, hidden_states):
+        return F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
+
+
+class RBLNGptOssExperts(nn.Module):
+    def __init__(self, model, top_k: Optional[int] = None):
+        super().__init__()
+        self.intermediate_size = model.intermediate_size
+        self.num_experts = model.num_experts
+        self.hidden_size = model.hidden_size
+
+        self.register_buffer(
+            "gate_proj_blocks",
+            model.gate_up_proj_blocks.data[:, ::2, :, :].reshape(self.num_experts, self.intermediate_size, -1),
+        )
+        self.register_buffer("gate_proj_scales", model.gate_up_proj_scales.data[:, ::2, :])
+        self.register_buffer(
+            "gate_proj_bias",
+            model.gate_up_proj_bias.data[:, ::2].reshape(self.num_experts, self.intermediate_size),
+        )
+
+        self.register_buffer(
+            "up_proj_blocks",
+            model.gate_up_proj_blocks.data[:, 1::2, :, :].reshape(self.num_experts, self.intermediate_size, -1),
+        )
+        self.register_buffer("up_proj_scales", model.gate_up_proj_scales.data[:, 1::2, :])
+        self.register_buffer(
+            "up_proj_bias", model.gate_up_proj_bias.data[:, 1::2].reshape(self.num_experts, self.intermediate_size)
+        )
+
+        self.register_buffer(
+            "down_proj_blocks", model.down_proj_blocks.data.reshape(self.num_experts, self.hidden_size, -1)
+        )
+        self.register_buffer("down_proj_scales", model.down_proj_scales.data)
+        self.register_buffer("down_proj_bias", model.down_proj_bias.data)
+
+        self.alpha = model.alpha  # 1.702
+        self.limit = model.limit  # 7.0
+        self.top_k = top_k
+
+    def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor:
+        return torch.ops.rbln_custom_ops.custom_moe_glu_mxfp4(
+            hidden_states,
+            self.gate_proj_blocks,
+            self.gate_proj_scales,
+            self.gate_proj_bias,
+            self.up_proj_blocks,
+            self.up_proj_scales,
+            self.up_proj_bias,
+            self.down_proj_blocks,
+            self.down_proj_scales,
+            self.down_proj_bias,
+            router_logits,
+            torch.tensor(self.alpha, dtype=hidden_states.dtype),
+            torch.tensor(self.limit, dtype=hidden_states.dtype),
+            k=self.top_k,
+            post_norm=True,
+        )
+
+
+class RBLNGptOssMLP(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.router = RBLNGptOssTopKRouter(model.router)
+        self.experts = RBLNGptOssExperts(model.experts, top_k=model.router.top_k)
+
+    def forward(self, hidden_states):
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        router_logits = self.router(hidden_states)
+        routed_out = self.experts(hidden_states, router_logits=router_logits)
+        routed_out = routed_out.reshape(batch_size, sequence_length, hidden_dim)
+        return routed_out

optimum/rbln/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -0,0 +1,165 @@
+# 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 TYPE_CHECKING, Optional, Union
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig
+from transformers.integrations.mxfp4 import Mxfp4GptOssExperts
+from transformers.modeling_utils import PreTrainedModel, no_init_weights
+
+from ....utils.logging import get_logger
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModelConfig,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
+from ...utils.rbln_quantization import load_weight_files
+from .gpt_oss_architecture import RBLNGptOssWrapper
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+logger = get_logger(__name__)
+
+
+class RBLNGptOssForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The GPT-OSS 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 GPT-OSSForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GPT-OSSForCausalLM 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 [`RBLNGptOssForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNGptOssForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNGptOssForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNGptOssForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNGptOssForCausalLM.from_pretrained(
+            "openai/gpt-oss-20b",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNGptOssForCausalLM.from_pretrained(
+            "openai/gpt-oss-20b",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNGptOssForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNGptOssForCausalLMConfig
+
+        config = RBLNGptOssForCausalLMConfig(
+            batch_size=1,
+            tensor_parallel_size=4
+        )
+        model = RBLNGptOssForCausalLM.from_pretrained(
+            "openai/gpt-oss-20b",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = RBLNGptOssWrapper
+
+    @staticmethod
+    def _get_dtype(dtype: Union[str, torch.dtype] = None, torch_dtype: Union[str, torch.dtype] = None):
+        # For BC on torch_dtype argument
+        if torch_dtype is not None:
+            logger.warning_once("`torch_dtype` is deprecated! Use `dtype` instead!")
+            # If both kwargs are provided, use `dtype`
+            dtype = dtype if dtype is not None else torch_dtype
+
+        # As mxfp4_quantizer's default dtype
+        if dtype is None or dtype == "auto":
+            dtype = torch.bfloat16
+
+        return dtype
+
+    @classmethod
+    def get_pytorch_model(
+        cls,
+        model_id: str,
+        *args,
+        rbln_config: Optional[RBLNDecoderOnlyModelConfig] = None,
+        dtype: Union[str, torch.dtype] = None,
+        torch_dtype: Union[str, torch.dtype] = None,
+        config: Optional[PretrainedConfig] = None,
+        **kwargs,
+    ) -> PreTrainedModel:
+        safetensor_files = load_weight_files(model_id, exception_keywords=["original"])
+        state_dict = {k: v for f in safetensor_files for k, v in load_file(f).items()}
+
+        if config is None:
+            config, kwargs = AutoConfig.from_pretrained(model_id, return_unused_kwargs=True)
+
+        dtype = cls._get_dtype(dtype, torch_dtype)
+
+        with no_init_weights():
+            model = AutoModelForCausalLM.from_config(config, dtype=dtype, **kwargs)
+
+        _replace_with_mxfp4_linear(model, config)
+        model.load_state_dict(state_dict, strict=False)
+
+        return model
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None,
+    ) -> RBLNDecoderOnlyModelForCausalLMConfig:
+        rbln_config = super()._update_rbln_config(preprocessors, model, model_config, rbln_config)
+
+        if rbln_config.use_attention_mask:
+            raise ValueError(
+                "use_attention_mask is not supported for GPT-OSS because custom attention does not support attention sink for masked attention"
+            )
+
+        return rbln_config
+
+
+def _replace_with_mxfp4_linear(
+    model,
+    config,
+):
+    for name, module in model.named_children():
+        if module.__class__.__name__ == "GptOssExperts":
+            model._modules[name] = Mxfp4GptOssExperts(config)
+        if len(list(module.children())) > 0:
+            _replace_with_mxfp4_linear(module, config)

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

@@ -50,11 +50,14 @@ class RBLNGroundingDinoForObjectDetectionConfig(RBLNImageModelConfig):
         Raises:
             ValueError: If batch_size is not a positive integer.
         """
-        super().__init__(**kwargs)
-        self.encoder = encoder
-        self.decoder = decoder
-        self.text_backbone = text_backbone
-        self.backbone = backbone
+
+        super().__init__(batch_size=batch_size, **kwargs)
+        self.encoder = self.initialize_submodule_config(submodule_config=encoder, batch_size=self.batch_size)
+        self.decoder = self.initialize_submodule_config(submodule_config=decoder, batch_size=self.batch_size)
+        self.text_backbone = self.initialize_submodule_config(
+            submodule_config=text_backbone, batch_size=self.batch_size
+        )
+        self.backbone = self.initialize_submodule_config(submodule_config=backbone, batch_size=self.batch_size)
         self.output_attentions = output_attentions if output_attentions is not None else False
         self.output_hidden_states = output_hidden_states if output_hidden_states is not None else False
 

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

@@ -509,10 +509,12 @@ class _GroundingDinoBiMultiHeadAttention(torch.nn.Module):
 
         # mask vision for language
         if vision_attention_mask is not None:
-            # RBLN FIX: bool tensor to float tensor
-            mask = vision_attention_mask * torch.finfo(torch.float16).min
-            text_attn_weights = text_attn_weights.transpose(1, 2) + mask
-            text_attn_weights = text_attn_weights.transpose(1, 2)
+            # RBLN FIX: bool tensor to float tensor, broadcast across heads and src_len
+            mask = vision_attention_mask
+            if mask.dim() == 3:
+                mask = mask[..., 0]
+            mask = mask[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1)
+            text_attn_weights = text_attn_weights + mask * torch.finfo(text_attn_weights.dtype).min
 
         text_attn_weights = text_attn_weights.softmax(dim=-1)
 

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

@@ -116,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

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