optimum-rbln 0.9.4a2__py3-none-any.whl → 0.10.0.post1__py3-none-any.whl

optimum/rbln/transformers/models/qwen3_moe/__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_qwen3_moe import RBLNQwen3MoeForCausalLMConfig
+from .modeling_qwen3_moe import RBLNQwen3MoeForCausalLM

optimum/rbln/transformers/models/qwen3_moe/configuration_qwen3_moe.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 ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNQwen3MoeForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN Qwen3 Moe models.
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNQwen3MoeForCausalLM, RBLNQwen3MoeForCausalLMConfig
+    # Create a configuration object
+    config = RBLNQwen3MoeForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=262144,
+        tensor_parallel_size=4
+    )
+    # Use the configuration with from_pretrained
+    model = RBLNQwen3MoeForCausalLM.from_pretrained(
+        "Qwen/Qwen3-30B-A3B-Thinking-2507",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/qwen3_moe/modeling_qwen3_moe.py

@@ -0,0 +1,68 @@
+# 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 ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from .qwen3_moe_architecture import Qwen3MoeWrapper
+
+
+class RBLNQwen3MoeForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Qwen3 Moe is a Mixture-of-Experts (MoE) variant of Qwen3, available as a base model and an aligned chat model.
+    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 Qwen3MoeForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen3MoeForCausalLM 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 [`RBLNQwen3MoeForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen3MoeForCausalLMConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNQwen3MoeForCausalLMConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen3MoeForCausalLM
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen3MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen3-30B-A3B-Thinking-2507",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 262144,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNQwen3MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen3-30B-A3B-Thinking-2507",
+            export=True,
+            rbln_config=rbln_config
+        )
+        # Using a RBLNQwen3ForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNQwen3MoeForCausalLMConfig
+        config = RBLNQwen3MoeForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=262144,
+            tensor_parallel_size=4
+        )
+        model = RBLNQwen3MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen3-30B-A3B-Thinking-2507",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = Qwen3MoeWrapper

optimum/rbln/transformers/models/qwen3_moe/qwen3_moe_architecture.py

@@ -0,0 +1,100 @@
+# 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
+from torch import nn
+
+from ..decoderonly.configuration_decoderonly import RBLNLoRAConfig
+from ..decoderonly.decoderonly_architecture import DecoderOnlyAttention, DecoderOnlyLayer, DecoderOnlyWrapper
+
+
+class Qwen3MoeWrapper(DecoderOnlyWrapper):
+    def get_rbln_layer_class(self):
+        return Qwen3MoeLayer
+
+    def get_rbln_attn_class(self):
+        return Qwen3MoeAttention
+
+
+class Qwen3MoeAttention(DecoderOnlyAttention):
+    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
+
+
+class Qwen3MoeLayer(DecoderOnlyLayer):
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config: Optional[RBLNLoRAConfig] = None):
+        super().__init__(layer, self_attn, lora_config)
+        self.mlp = (
+            Qwen3MoeSparseMoeBlock(layer.mlp)
+            if layer.mlp.__class__.__name__ == "Qwen3MoeSparseMoeBlock"
+            else layer.mlp
+        )
+
+    def get_mlp(self) -> nn.Module:
+        return self.mlp
+
+
+class Qwen3MoeSparseMoeBlock(nn.Module):
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.num_experts = model.num_experts
+        self.top_k = model.top_k
+        self.norm_topk_prob = model.norm_topk_prob
+        self.gate = model.gate
+        self.experts = Qwen3MoeMLP(model.experts, self.top_k, self.norm_topk_prob)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        # router_logits: (batch * sequence_length, n_experts)
+        router_logits = self.gate(hidden_states)
+        final_hidden_states = self.experts(hidden_states, router_logits)
+
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states
+
+
+class Qwen3MoeMLP(nn.Module):
+    def __init__(self, expert_list, top_k, norm_topk_prob):
+        super().__init__()
+        self.hidden_size = expert_list[0].hidden_size
+        self.intermediate_size = expert_list[0].intermediate_size
+        self.top_k = top_k
+        self.norm_topk_prob = norm_topk_prob
+        self.num_experts = len(expert_list)
+        self.gate_proj = nn.Linear(self.hidden_size, self.num_experts * self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.num_experts * self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.num_experts * self.intermediate_size, self.hidden_size, bias=False)
+        self.gate_proj.weight.data = torch.stack([expert.gate_proj.weight.data for expert in expert_list], dim=0)
+        self.up_proj.weight.data = torch.stack([expert.up_proj.weight.data for expert in expert_list], dim=0)
+        self.down_proj.weight.data = torch.stack([expert.down_proj.weight.data for expert in expert_list], dim=0)
+
+    def forward(self, x, router_logits):
+        return torch.ops.rbln_custom_ops.custom_moe_glu(
+            x,
+            self.gate_proj.weight,
+            self.up_proj.weight,
+            self.down_proj.weight,
+            router_logits,
+            self.top_k,
+            self.norm_topk_prob,
+        )

optimum/rbln/transformers/models/resnet/configuration_resnet.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from typing import Optional
+from typing import Any, Optional, Tuple, Union
 
 from ...configuration_generic import RBLNModelForImageClassificationConfig
 
@@ -26,17 +26,23 @@ class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConf
     RBLN-optimized ResNet models for image classification tasks.
     """
 
-    def __init__(self, output_hidden_states: Optional[bool] = None, **kwargs):
+    def __init__(
+        self,
+        image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        batch_size: Optional[int] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs: Any,
+    ):
         """
         Args:
             image_size (Optional[Union[int, Tuple[int, int]]]): The size of input images.
                 Can be an integer for square images or a tuple (height, width).
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
-            output_hidden_states (bool, optional) — Whether or not to return the hidden states of all layers.
+            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers.
             kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.
         """
-        super().__init__(**kwargs)
+        super().__init__(image_size=image_size, batch_size=batch_size, **kwargs)
         self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py

@@ -268,13 +268,12 @@ class Seq2SeqDecoder(torch.nn.Module):
 
     def __init__(self, model, layers, **kwargs):
         super().__init__()
-        self._original_mod = model
         self.layers = nn.ModuleList(layers)
         self.embed_tokens = model.embed_tokens
-        self.final_layer_norm = getattr(model, "final_layer_norm", None)
-        self.__post_init__(**kwargs)
+        self.final_layer_norm = getattr(model, "final_layer_norm", None) or getattr(model, "layer_norm", None)
+        self.__post_init__(model, **kwargs)
 
-    def __post_init__(self, **kwargs):
+    def __post_init__(self, model: nn.Module, **kwargs):
         """
         Abstract method intended to be overridden by subclasses to modify or override
         the attributes of the original model after initialization.
@@ -344,12 +343,11 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
 
     def __init__(self, decoder_layer, self_attn, cross_attn):
         super().__init__()
-        self._original_mod = decoder_layer
         self.self_attn = self_attn
         self.cross_attn = cross_attn
-        self.__post_init__()
+        self.__post_init__(decoder_layer)
 
-    def __post_init__(self, **kwargs):
+    def __post_init__(self, decoder_layer: nn.Module, **kwargs):
         """
         Abstract method intended to be overridden by subclasses to modify or override
         the attributes of the original model after initialization.
@@ -423,10 +421,9 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
 class Seq2SeqSelfAttention(nn.Module):
     def __init__(self, attn, **kwargs):
         super().__init__()
-        self._original_mod = attn
-        self.__post_init__(**kwargs)
+        self.__post_init__(attn, **kwargs)
 
-    def __post_init__(self, **kwargs):
+    def __post_init__(self, attn: nn.Module, **kwargs):
         """
         Abstract method intended to be overridden by subclasses to modify or override
         the attributes of the original model after initialization.
@@ -495,8 +492,13 @@ class Seq2SeqSelfAttention(nn.Module):
 class Seq2SeqCrossAttention(nn.Module):
     def __init__(self, attn, **kwargs):
         super().__init__()
-        self._original_mod = attn
-        self.__post_init__(**kwargs)
+        self.__post_init__(attn, **kwargs)
+
+    def __post_init__(self, attn: nn.Module, **kwargs):
+        """
+        Optional post-init hook for subclasses (e.g., to register q/k/v/out projections).
+        """
+        pass
 
     def forward(
         self,

optimum/rbln/transformers/models/siglip/modeling_siglip.py

@@ -21,6 +21,7 @@ from transformers.modeling_outputs import BaseModelOutputWithPooling
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
+from ...modeling_outputs import _validate_output_attentions, _validate_output_hidden_states
 from .configuration_siglip import RBLNSiglipVisionModelConfig
 
 
@@ -52,7 +53,7 @@ class _SiglipVisionModel(torch.nn.Module):
             interpolate_pos_encoding=self.interpolate_pos_encoding,
             output_attentions=self.output_attentions,
         )
-        return tuple(x for x in enc_out if x is not None)
+        return enc_out
 
 
 class RBLNSiglipVisionModel(RBLNModel):
@@ -138,23 +139,8 @@ class RBLNSiglipVisionModel(RBLNModel):
             The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPooling object.
         """
 
-        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
-        )
-
-        if output_attentions != self.rbln_config.output_attentions:
-            raise ValueError(
-                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
-                f"Please compile again with the correct argument."
-            )
-
-        if output_hidden_states != self.rbln_config.output_hidden_states:
-            raise ValueError(
-                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
-                f"Please compile again with the correct argument."
-            )
-
+        output_attentions = _validate_output_attentions(output_attentions, self.rbln_config)
+        output_hidden_states = _validate_output_hidden_states(output_hidden_states, self.rbln_config)
         if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
             raise ValueError(
                 f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "

optimum/rbln/transformers/models/swin/configuration_swin.py

@@ -32,11 +32,6 @@ class RBLNSwinBackboneConfig(RBLNModelForImageClassificationConfig):
         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}")
-
-        self.image_size = image_size
+        super().__init__(batch_size=batch_size, image_size=image_size, **kwargs)
         self.output_hidden_states = output_hidden_states
         self.output_attentions = output_attentions

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -111,9 +111,9 @@ class T5ForConditionalGeneration(Seq2SeqForConditionalGeneration):
 class T5Decoder(Seq2SeqDecoder):
     has_pos_emb = False
 
-    def __post_init__(self, dec_max_seq_len: int = None):
-        self.invert_attention_mask = self._original_mod.invert_attention_mask
-        self._dec_position_bias = self.precompute_dec_position_bias(self._original_mod, dec_max_seq_len)
+    def __post_init__(self, model: nn.Module, dec_max_seq_len: int = None):
+        self.invert_attention_mask = model.invert_attention_mask
+        self._dec_position_bias = self.precompute_dec_position_bias(model, dec_max_seq_len)
 
     def precompute_dec_position_bias(self, model, dec_max_length):
         attn_layer = model.block[0].layer[0].SelfAttention
@@ -145,13 +145,12 @@ class T5Decoder(Seq2SeqDecoder):
 class T5Block(Seq2SeqDecoderLayer):
     def __init__(self, decoder_layer, self_attn):
         super().__init__(decoder_layer, self_attn, cross_attn=None)
-        self.__post_init__()
 
-    def __post_init__(self):
-        self.self_attn_layer_norm = self._original_mod.layer[0].layer_norm
-        self.encoder_attn_layer_norm = self._original_mod.layer[1].layer_norm
-        self.cross_attn = T5CrossAttention(self._original_mod.layer[1].EncDecAttention)
-        self.ff_layer = self._original_mod.layer[2]
+    def __post_init__(self, decoder_layer: nn.Module):
+        self.self_attn_layer_norm = decoder_layer.layer[0].layer_norm
+        self.encoder_attn_layer_norm = decoder_layer.layer[1].layer_norm
+        self.cross_attn = T5CrossAttention(decoder_layer.layer[1].EncDecAttention)
+        self.ff_layer = decoder_layer.layer[2]
 
     def pre_self_attn_layer_norm(self, hidden_states):
         return self.self_attn_layer_norm(hidden_states)
@@ -167,13 +166,13 @@ class T5Block(Seq2SeqDecoderLayer):
 
 
 class T5LayerSelfAttention(Seq2SeqSelfAttention):
-    def __post_init__(self):
-        self.q_proj = self._original_mod.q
-        self.k_proj = self._original_mod.k
-        self.v_proj = self._original_mod.v
-        self.out_proj = self._original_mod.o
-        self.num_heads = self._original_mod.n_heads
-        self.head_dim = self._original_mod.key_value_proj_dim
+    def __post_init__(self, attn: nn.Module):
+        self.q_proj = attn.q
+        self.k_proj = attn.k
+        self.v_proj = attn.v
+        self.out_proj = attn.o
+        self.num_heads = attn.n_heads
+        self.head_dim = attn.key_value_proj_dim
         self.attn_decode = torch.ops.rbln_custom_ops.paged_add_softmax_attn_decode
 
     def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

optimum/rbln/transformers/models/time_series_transformer/time_series_transformers_architecture.py

@@ -140,7 +140,6 @@ class TimeSeriesTransformersDecoderWrapper(torch.nn.Module):
 class TimeSeriesTransformersDecoder(nn.Module):
     def __init__(self, model, layers, **kwargs):
         super().__init__()
-        self._original_mod = model
         self.config = model.config
         self.layers = nn.ModuleList(layers)
         self.value_embedding = model.value_embedding
@@ -190,7 +189,6 @@ class TimeSeriesTransformersDecoder(nn.Module):
 class TimeSeriesTransformersDecoderLayer(nn.Module):
     def __init__(self, decoder_layer, self_attn, cross_attn):
         super().__init__()
-        self._original_mod = decoder_layer
         self.self_attn = self_attn
         self.encoder_attn = cross_attn
         self.embed_dim = decoder_layer.embed_dim
@@ -245,7 +243,6 @@ class TimeSeriesTransformersDecoderLayer(nn.Module):
 class TimeSeriesTransformersAttention(nn.Module):
     def __init__(self, attn, num_parallel_samples):
         super().__init__()
-        self._original_mod = attn
         self.q_proj = attn.q_proj
         self.k_proj = attn.k_proj
         self.v_proj = attn.v_proj

optimum/rbln/transformers/models/whisper/generation_whisper.py

@@ -51,22 +51,22 @@ class RBLNWhisperGenerationMixin(WhisperGenerationMixin, GenerationMixin):
         return_segments: Optional[bool] = None,
         return_timestamps: Optional[bool] = None,
         return_token_timestamps: Optional[bool] = None,
-        **kwargs,
+        **kwargs: Optional[Dict[str, Any]],
     ) -> Union[ModelOutput, Dict[str, Any], torch.LongTensor]:
         """
         The generate function is utilized in its standard form as in the HuggingFace transformers library. User can use this function to generate text from the model.
         Check the [HuggingFace transformers documentation](https://huggingface.co/docs/transformers/v4.57.1/en/model_doc/whisper#transformers.WhisperForConditionalGeneration.generate) for more details.
 
         Args:
-            input_features(torch.Tensor, optional): The input features to the model.
-            attention_mask(torch.Tensor, optional): Attention mask needs to be passed when doing long-form transcription using a batch size > 1.
-            generation_config(GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
+            input_features (torch.Tensor, optional): The input features to the model.
+            attention_mask (torch.Tensor, optional): Attention mask needs to be passed when doing long-form transcription using a batch size > 1.
+            generation_config (GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
                 If generation_config is not provided, the default will be used, which had the following loading priority: 1) from the generation_config.json model file, if it exists; 2) from the model configuration.
                 Please note that unspecified parameters will inherit [GenerationConfig](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationConfig)’s default values.
-            return_segments(bool, optional): Whether to return segments.
-            return_timestamps(bool, optional): Whether to return the timestamps with the text. For audios longer than 30 seconds, it is necessary to set return_timestamps=True.
-            return_token_timestamps(bool, optional): Whether to return token timestamps.
-            kwargs(dict[str, Any], optional): Additional arguments passed to the generate function.
+            return_segments (bool, optional): Whether to return segments.
+            return_timestamps (bool, optional): Whether to return the timestamps with the text. For audios longer than 30 seconds, it is necessary to set return_timestamps=True.
+            return_token_timestamps (bool, optional): Whether to return token timestamps.
+            kwargs (dict[str, Any], optional): Additional arguments passed to the generate function.
 
         Returns:
             Transcribes or translates log-mel input features to a sequence of auto-regressively generated token ids.

optimum/rbln/transformers/models/whisper/whisper_architecture.py

@@ -154,7 +154,6 @@ class WhisperDecoderWrapper(torch.nn.Module):
 class WhisperDecoder(nn.Module):
     def __init__(self, model, layers, **kwargs):
         super().__init__()
-        self._original_mod = model
         self.layers = nn.ModuleList(layers)
         self.embed_tokens = model.embed_tokens
         self.layer_norm = model.layer_norm
@@ -210,7 +209,6 @@ class WhisperDecoder(nn.Module):
 class WhisperDecoderLayer(nn.Module):
     def __init__(self, decoder_layer, self_attn, cross_attn):
         super().__init__()
-        self._original_mod = decoder_layer
         self.self_attn = self_attn
         self.encoder_attn = cross_attn
         self.self_attn_layer_norm = decoder_layer.self_attn_layer_norm
@@ -263,7 +261,6 @@ class WhisperDecoderLayer(nn.Module):
 class WhisperAttention(nn.Module):
     def __init__(self, attn):
         super().__init__()
-        self._original_mod = attn
         self.q_proj = attn.q_proj
         self.k_proj = attn.k_proj
         self.v_proj = attn.v_proj

optimum/rbln/transformers/utils/rbln_quantization.py

@@ -221,11 +221,12 @@ def load_weight_files(
     cache_dir: Optional[str] = None,
     force_download: bool = False,
     local_files_only: bool = False,
+    exception_keywords: Optional[List[str]] = None,
 ) -> list[str]:
     """
     Discover and download safetensors files for the given model id.
     """
-
+    exception_keywords = exception_keywords or []
     if os.path.isdir(model_id):
         safetensor_files = glob.glob(f"{model_id}/*.safetensors")
     else:
@@ -237,17 +238,24 @@ def load_weight_files(
 
             for file in repo_files:
                 if file.endswith(".safetensors"):
-                    # Download the safetensors file
-                    downloaded_file = hf_hub_download(
-                        repo_id=model_id,
-                        filename=file,
-                        revision=revision,
-                        token=use_auth_token,
-                        cache_dir=cache_dir,
-                        force_download=force_download,
-                        local_files_only=local_files_only,
-                    )
-                    safetensor_files.append(downloaded_file)
+                    exculde = False
+                    for except_key in exception_keywords:
+                        if except_key in file:
+                            exculde = True
+                            break
+
+                    if not exculde:
+                        # Download the safetensors file
+                        downloaded_file = hf_hub_download(
+                            repo_id=model_id,
+                            filename=file,
+                            revision=revision,
+                            token=use_auth_token,
+                            cache_dir=cache_dir,
+                            force_download=force_download,
+                            local_files_only=local_files_only,
+                        )
+                        safetensor_files.append(downloaded_file)
         except Exception as e:
             logger.error(f"Failed to download safetensors files from Hugging Face Hub: {e}")
             raise e

optimum/rbln/utils/deprecation.py

@@ -194,7 +194,7 @@ def deprecate_kwarg(
                 message = f"{message} {additional_message}"
 
             # update minimum_action if argument is ALREADY deprecated (current version >= deprecated version)
-            if is_greater_or_equal_version:
+            if is_greater_or_equal_version and message is not None:
                 # change to NOTIFY -> RAISE  in case we want to raise error for already deprecated arguments
                 if raise_if_greater_or_equal_version:
                     minimum_action = Action.RAISE
@@ -211,3 +211,80 @@ def deprecate_kwarg(
         return wrapped_func
 
     return wrapper
+
+
+def deprecate_method(
+    version: str,
+    new_method: Optional[str] = None,
+    raise_if_greater_or_equal_version: bool = True,
+    additional_message: Optional[str] = None,
+):
+    """
+    Decorator to mark a method as deprecated, optionally pointing to a replacement method.
+    This decorator allows you to:
+    - Notify users when a method is deprecated.
+    - Optionally specify a new method name that should be used instead.
+    - Raise an error if the deprecated method is called after the specified version.
+    Parameters:
+        version (`str`):
+            The version in which the method was (or will be) deprecated.
+        new_method (`Optional[str]`, *optional*):
+            The name of the new method to use instead. If specified, users will be directed to use this method.
+        raise_if_greater_or_equal_version (`bool`, *optional*, defaults to `True`):
+            Whether to raise `ValueError` if current `optimum.rbln` version is greater than or equal to the deprecated version.
+        additional_message (`Optional[str]`, *optional*):
+            An additional message to append to the default deprecation message.
+    Returns:
+        Callable:
+            A wrapped function that handles the deprecation warning or error.
+    Examples:
+        >>> class MyClass:
+        ...     @deprecate_method(version="0.12.0", new_method="from_pretrained")
+        ...     def load(self, path):
+        ...         return self.from_pretrained(path)
+    """
+
+    deprecated_version = packaging.version.parse(version)
+    current_version = packaging.version.parse(__version__)
+    is_greater_or_equal_version = current_version >= deprecated_version
+
+    if is_greater_or_equal_version:
+        version_message = f"and removed starting from version {version}"
+    else:
+        version_message = f"and will be removed in version {version}"
+
+    def wrapper(func):
+        sig = inspect.signature(func)
+        function_named_args = set(sig.parameters.keys())
+        is_instance_method = "self" in function_named_args
+        is_class_method = "cls" in function_named_args
+
+        @wraps(func)
+        def wrapped_func(*args, **kwargs):
+            # Get class + method name for better warning message
+            method_name = func.__name__
+            if is_instance_method:
+                method_name = f"{args[0].__class__.__name__}.{method_name}"
+            elif is_class_method:
+                method_name = f"{args[0].__name__}.{method_name}"
+
+            # Build deprecation message
+            if new_method is not None:
+                message = f"`{method_name}` is deprecated {version_message}. Use `{new_method}` instead."
+            else:
+                message = f"`{method_name}` is deprecated {version_message}."
+
+            if additional_message is not None:
+                message = f"{message} {additional_message}"
+
+            # Determine action based on version
+            if is_greater_or_equal_version and raise_if_greater_or_equal_version:
+                raise ValueError(message)
+            else:
+                logger.warning(message, stacklevel=2)
+
+            return func(*args, **kwargs)
+
+        return wrapped_func
+
+    return wrapper