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

optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py

@@ -6,22 +6,28 @@ import torch.nn as nn
 from transformers import PreTrainedModel
 
 from ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper, apply_rotary_pos_emb
+from .configuration_qwen2_5_vl import RBLNQwen2_5_VisionTransformerPretrainedModelConfig
 
 
 class Qwen2_5_VisionTransformerWrapper(nn.Module):
-    def __init__(self, model: torch.nn.Module):
+    def __init__(self, model: torch.nn.Module, rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig):
         super().__init__()
-        self._original_mod = model
         self.fullatt_block_indexes = model.fullatt_block_indexes
         self.merger = model.merger
+        self.rbln_config = rbln_config
         window_seq_len = (model.window_size // model.patch_size) ** 2
-        self.blocks = self.wrap_vision_blocks(model.blocks, window_seq_len)
+        self.blocks = self.wrap_vision_blocks(model.blocks, window_seq_len, rbln_config)
 
-    def wrap_vision_blocks(self, blocks: torch.nn.ModuleList, window_seq_len: int):
+    def wrap_vision_blocks(
+        self,
+        blocks: torch.nn.ModuleList,
+        window_seq_len: int,
+        rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
+    ):
         wrapped_blocks = []
         for i, block in enumerate(blocks):
             is_full_attn = True if i in self.fullatt_block_indexes else False
-            wrapped_blocks.append(Qwen2_5_VLVisionBlock(block, is_full_attn, window_seq_len))
+            wrapped_blocks.append(Qwen2_5_VLVisionBlock(block, is_full_attn, window_seq_len, rbln_config))
         return nn.ModuleList(wrapped_blocks)
 
     def forward(
@@ -32,8 +38,8 @@ class Qwen2_5_VisionTransformerWrapper(nn.Module):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ):
-        full_attn_masks = (1 - full_attn_masks) * torch.finfo(torch.float32).min
-        window_attn_masks = (1 - window_attn_masks) * torch.finfo(torch.float32).min
+        full_attn_masks = (1.0 - full_attn_masks) * torch.finfo(hidden_states.dtype).min
+        window_attn_masks = (1.0 - window_attn_masks) * torch.finfo(hidden_states.dtype).min
 
         for i, block in enumerate(self.blocks):
             attn_masks = full_attn_masks if i in self.fullatt_block_indexes else window_attn_masks
@@ -45,16 +51,23 @@ class Qwen2_5_VisionTransformerWrapper(nn.Module):
 
 
 class Qwen2_5_VLVisionBlock(torch.nn.Module):
-    def __init__(self, model: torch.nn.Module, is_full_attn: bool, window_seq_len: int):
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        is_full_attn: bool,
+        window_seq_len: int,
+        rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
+    ):
         super().__init__()
         self._origin_model = model
+        self.rbln_config = rbln_config
         self.norm1 = model.norm1
         self.norm2 = model.norm2
 
         if is_full_attn:
-            self.attn = Qwen2_5_VLVisionFullAttention(model.attn)
+            self.attn = Qwen2_5_VLVisionFullAttention(model.attn, rbln_config)
         else:
-            self.attn = Qwen2_5_VLVisionWindowAttention(model.attn, window_seq_len)
+            self.attn = Qwen2_5_VLVisionWindowAttention(model.attn, window_seq_len, rbln_config)
         self.mlp = model.mlp
 
     def forward(
@@ -73,13 +86,15 @@ class Qwen2_5_VLVisionBlock(torch.nn.Module):
 
 
 class Qwen2_5_VLVisionFullAttention(nn.Module):
-    def __init__(self, model: nn.Module) -> None:
+    def __init__(self, model: nn.Module, rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig) -> None:
         super().__init__()
         self._origin_model = model
+        self.rbln_config = rbln_config
         self.num_heads = model.num_heads
         self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
         self.qkv = model.qkv
         self.proj = model.proj
+        self.scale = torch.tensor(1 / math.sqrt(self.head_dim), dtype=rbln_config.dtype)
 
     def forward(
         self,
@@ -96,9 +111,9 @@ class Qwen2_5_VLVisionFullAttention(nn.Module):
         cos, sin = position_embeddings
         q, k = apply_rotary_pos_emb(q, k, cos, sin)
 
-        attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)
+        attn_weights = torch.matmul(q, k.transpose(2, 3)) * self.scale
         attn_weights = attn_weights + attn_masks
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=hidden_states.dtype)
         attn_output = torch.matmul(attn_weights, v)
         attn_output = attn_output.transpose(1, 2)
         attn_output = attn_output.reshape(1, seq_length, -1)
@@ -108,14 +123,18 @@ class Qwen2_5_VLVisionFullAttention(nn.Module):
 
 
 class Qwen2_5_VLVisionWindowAttention(nn.Module):
-    def __init__(self, model: nn.Module, window_seq_len: int) -> None:
+    def __init__(
+        self, model: nn.Module, window_seq_len: int, rbln_config: RBLNQwen2_5_VisionTransformerPretrainedModelConfig
+    ) -> None:
         super().__init__()
         self._origin_model = model
+        self.rbln_config = rbln_config
         self.num_heads = model.num_heads
         self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
         self.qkv = model.qkv
         self.proj = model.proj
         self.window_seq_len = window_seq_len
+        self.scale = torch.tensor(1 / math.sqrt(self.head_dim), dtype=rbln_config.dtype)
 
     def forward(
         self,
@@ -142,10 +161,10 @@ class Qwen2_5_VLVisionWindowAttention(nn.Module):
         sin = sin.reshape(num_windows, 1, seq_length // num_windows, -1)
         q, k = apply_rotary_pos_emb(q, k, cos, sin)
 
-        attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)
+        attn_weights = torch.matmul(q, k.transpose(2, 3)) * self.scale
 
         attn_weights = attn_weights + attn_masks
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
         attn_output = torch.matmul(attn_weights, v)
         attn_output = attn_output.transpose(1, 2)
         attn_output = attn_output.reshape(1, seq_length, -1)
@@ -155,6 +174,12 @@ class Qwen2_5_VLVisionWindowAttention(nn.Module):
 
 
 class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
+    def get_decoder_layers(self, model: PreTrainedModel):
+        return model.model.language_model.layers if hasattr(model, "model") else model.language_model.layers
+
+    def get_model_layer(self, model: PreTrainedModel):
+        return model.model.language_model if hasattr(model, "model") else model.language_model
+
     def prepare_forward_args(self, *args):
         args = list(args)
         input_ids = None if self.rbln_config.use_inputs_embeds else args.pop(0)
@@ -163,10 +188,10 @@ class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
         global_block_tables = args.pop(0)
         local_block_tables = None
         position_embeds = args.pop(0)
-        query_position = args.pop(0) if self.phase == "prefill" else None
+        query_position = args.pop(0) if self.phase == "prefill" and self.rbln_config.logits_to_keep > 0 else None
         position_ids = None
-        lora_int_id = None
         attention_mask = args.pop(0) if self.rbln_config.use_attention_mask else None
+        lora_int_id = args.pop(0) if self.rbln_config.lora_config else None
         past_key_values = args
 
         if len(past_key_values) != 2 * self.num_hidden_layers:
@@ -197,24 +222,3 @@ class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
             past_key_values,
             position_embeds,
         )
-
-    def convert_to_rbln_class(self, model: PreTrainedModel, max_seq_len: int):
-        new_layers = []
-
-        for layer_idx, layer in enumerate(model.model.language_model.layers):
-            is_sliding = layer_idx in self.rbln_config.sliding_window_layers
-            new_self_attn = self.get_rbln_attn_class()(
-                self.get_attn_layer(layer), self.rbln_config, is_sliding=is_sliding
-            )
-            new_layer = self.get_rbln_layer_class()(layer, new_self_attn)
-            new_layers.append(new_layer)
-
-        new_model = self.get_rbln_model_class()(
-            model.model.language_model,
-            new_layers,
-            self.rbln_config,
-            use_learned_pos_emb=self.__class__._use_learned_pos_emb,
-        )
-
-        new_model = self.get_rbln_causal_lm_class()(model.model, new_model)
-        return new_model

optimum/rbln/transformers/models/qwen2_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_qwen2_moe import RBLNQwen2MoeForCausalLMConfig
+from .modeling_qwen2_moe import RBLNQwen2MoeForCausalLM

optimum/rbln/transformers/models/qwen2_moe/configuration_qwen2_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 RBLNQwen2MoeForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLN Qwen2 Moe models.
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNQwen2MoeForCausalLM, RBLNQwen2MoeForCausalLMConfig
+    # Create a configuration object
+    config = RBLNQwen2MoeForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=8192,
+        tensor_parallel_size=4
+    )
+    # Use the configuration with from_pretrained
+    model = RBLNQwen2MoeForCausalLM.from_pretrained(
+        "Qwen/Qwen1.5-MoE-A2.7B",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/qwen2_moe/modeling_qwen2_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 .qwen2_moe_architecture import Qwen2MoeWrapper
+
+
+class RBLNQwen2MoeForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Qwen2MoE is a Mixture-of-Experts (MoE) variant of Qwen2, 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 Qwen2MoeForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen2MoeForCausalLM 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 [`RBLNQwen2MoeForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen2MoeForCausalLMConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNQwen2MoeForCausalLMConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen2MoeForCausalLM
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen2MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen1.5-MoE-A2.7B",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 8192,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNQwen2MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen1.5-MoE-A2.7B",
+            export=True,
+            rbln_config=rbln_config
+        )
+        # Using a RBLNQwen2ForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNQwen2MoeForCausalLMConfig
+        config = RBLNQwen2MoeForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=8192,
+            tensor_parallel_size=4
+        )
+        model = RBLNQwen2MoeForCausalLM.from_pretrained(
+            "Qwen/Qwen1.5-MoE-A2.7B",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    _decoder_wrapper_cls = Qwen2MoeWrapper

optimum/rbln/transformers/models/qwen2_moe/qwen2_moe_architecture.py

@@ -0,0 +1,94 @@
+# 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 Qwen2MoeWrapper(DecoderOnlyWrapper):
+    def get_rbln_layer_class(self):
+        return Qwen2MoeLayer
+
+
+class Qwen2MoeLayer(DecoderOnlyLayer):
+    def __init__(self, layer, self_attn: DecoderOnlyAttention, lora_config: Optional[RBLNLoRAConfig] = None):
+        super().__init__(layer, self_attn, lora_config)
+        self.mlp = (
+            Qwen2MoeSparseMoeBlock(layer.mlp)
+            if layer.mlp.__class__.__name__ == "Qwen2MoeSparseMoeBlock"
+            else layer.mlp
+        )
+
+    def get_mlp(self) -> nn.Module:
+        return self.mlp
+
+
+class Qwen2MoeSparseMoeBlock(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.shared_expert = model.shared_expert
+        self.shared_expert_gate = model.shared_expert_gate
+        self.experts = Qwen2MoeMLP(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)
+        shared_expert_output = self.shared_expert(hidden_states)
+        shared_expert_output = (
+            torch.nn.functional.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
+        )
+        final_hidden_states = final_hidden_states + shared_expert_output
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states
+
+
+class Qwen2MoeMLP(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/qwen2_vl/__init__.py

@@ -15,5 +15,10 @@
 from .configuration_qwen2_vl import (
     RBLNQwen2VisionTransformerPretrainedModelConfig,
     RBLNQwen2VLForConditionalGenerationConfig,
+    RBLNQwen2VLModelConfig,
+)
+from .modeling_qwen2_vl import (
+    RBLNQwen2VisionTransformerPretrainedModel,
+    RBLNQwen2VLForConditionalGeneration,
+    RBLNQwen2VLModel,
 )
-from .modeling_qwen2_vl import RBLNQwen2VisionTransformerPretrainedModel, RBLNQwen2VLForConditionalGeneration

optimum/rbln/transformers/models/qwen2_vl/configuration_qwen2_vl.py

@@ -15,7 +15,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNQwen2VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -48,6 +48,16 @@ class RBLNQwen2VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausalLMC
         self.visual = visual
 
 
+class RBLNQwen2VLModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLNQwen2VLModel.
+    """
+
+    def __init__(self, visual: Optional[RBLNModelConfig] = None, **kwargs: Dict[str, Any]):
+        super().__init__(**kwargs)
+        self.visual = self.initialize_submodule_config(submodule_config=visual)
+
+
 class RBLNQwen2VisionTransformerPretrainedModelConfig(RBLNModelConfig):
     def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Dict[str, Any]):
         """