optimum-rbln 0.8.2a4__py3-none-any.whl → 0.8.2a6__py3-none-any.whl

optimum/rbln/transformers/models/pixtral/modeling_pixtral.py

@@ -0,0 +1,318 @@
+# 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 pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
+
+import rebel
+import torch
+import torch.nn as nn
+from transformers import PixtralVisionConfig, PixtralVisionModel
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_utils import no_init_weights
+from transformers.models.pixtral.modeling_pixtral import (
+    PixtralRMSNorm,
+    PixtralRotaryEmbedding,
+)
+
+from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from ....utils.runtime_utils import RBLNPytorchRuntime
+from .configuration_pixtral import RBLNPixtralVisionModelConfig
+from .pixtral_architecture import PixtralAttention
+
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+    from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
+
+
+class RBLNRuntimePixtralVisionModel(RBLNPytorchRuntime):
+    mandatory_members = ["main_input_name"]
+
+    def __init__(
+        self,
+        runtime: rebel.Runtime,
+        config: PixtralVisionConfig,
+        rbln_config: RBLNPixtralVisionModelConfig,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(runtime, **kwargs)
+        self.patch_positional_embedding = PixtralRotaryEmbedding(config)
+        self.patch_size = config.patch_size
+        self.image_size = config.image_size
+        self.hidden_size = config.hidden_size
+        self.max_image_size = rbln_config.max_image_size
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        image_sizes: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        if pixel_values.shape[2] > self.max_image_size[0] or pixel_values.shape[3] > self.max_image_size[1]:
+            raise ValueError("The height() and width of pixel_values can't be larger than max_image_size.")
+
+        if pixel_values.shape[2] != self.max_image_size[0] or pixel_values.shape[3] != self.max_image_size[1]:
+            padded_pixel_values = [
+                torch.nn.functional.pad(
+                    image,
+                    pad=(
+                        0,
+                        self.max_image_size[1] - pixel_values.shape[3],
+                        0,
+                        self.max_image_size[0] - pixel_values.shape[2],
+                    ),
+                )
+                for image in pixel_values
+            ]
+            pixel_values = torch.stack(padded_pixel_values)
+
+        batch_size, _, H_max, W_max = pixel_values.shape
+        H_max_p = H_max // self.patch_size
+        W_max_p = W_max // self.patch_size
+
+        final_hidden_states = None
+
+        last_hidden_state_list = []
+        if output_hidden_states:
+            batch_hidden_states_list = []
+
+        for i in range(batch_size):
+            h_patched_original = image_sizes[i, 0] // self.patch_size
+            w_patched_original = image_sizes[i, 1] // self.patch_size
+
+            single_pixel_values = pixel_values[i : i + 1]
+            patch_embed = self.patch_conv(single_pixel_values)
+            patch_embed_seq = patch_embed[:, :, :h_patched_original, :w_patched_original].flatten(2).transpose(1, 2)
+            patch_embed_seq = self.ln_pre(patch_embed_seq)
+            patch_embed_seq = nn.functional.pad(
+                patch_embed_seq, (0, 0, 0, H_max_p * W_max_p - patch_embed_seq.shape[1]), "constant", value=0
+            )
+
+            max_w_from_config = self.image_size // self.patch_size
+            mesh = torch.meshgrid(torch.arange(h_patched_original), torch.arange(w_patched_original), indexing="ij")
+            h_grid, v_grid = torch.stack(mesh, dim=-1).reshape(-1, 2).chunk(2, -1)
+            ids = h_grid * max_w_from_config + v_grid
+            position_ids = ids[:, 0]
+
+            position_embeddings = self.patch_positional_embedding(patch_embed_seq, position_ids)
+            cos = nn.functional.pad(
+                position_embeddings[0],
+                (0, 0, 0, H_max_p * W_max_p - position_embeddings[0].shape[0]),
+                "constant",
+                value=0,
+            )
+            sin = nn.functional.pad(
+                position_embeddings[1],
+                (0, 0, 0, H_max_p * W_max_p - position_embeddings[1].shape[0]),
+                "constant",
+                value=0,
+            )
+
+            attention_mask = torch.full(
+                (1, patch_embed_seq.shape[-2]), fill_value=torch.finfo(patch_embed_seq.dtype).min
+            )
+            attention_mask[:, : h_patched_original * w_patched_original] = 0
+
+            transformer_output = super().forward(patch_embed_seq, attention_mask, cos, sin)
+
+            last_hidden_state_list.append(transformer_output[0][:, : h_patched_original * w_patched_original, :])
+            hidden_states = transformer_output[1:]
+
+            if output_hidden_states:
+                batch_hidden_states_list.append(
+                    [hidden_state[:, : h_patched_original * w_patched_original, :] for hidden_state in hidden_states]
+                )
+
+        final_last_hidden_state = torch.cat(last_hidden_state_list, dim=1)
+
+        if output_hidden_states:
+            hidden_states = [
+                torch.cat(
+                    [batch_hidden_states[layer_idx] for batch_hidden_states in batch_hidden_states_list],
+                    dim=1,
+                )
+                for layer_idx in range(len(batch_hidden_states_list[0]))
+            ]
+
+            final_hidden_states = tuple(hidden_states)
+
+        if not return_dict:
+            return tuple(v for v in (final_last_hidden_state, final_hidden_states) if v is not None)
+
+        # TODO: output_attentions
+        return BaseModelOutput(
+            last_hidden_state=final_last_hidden_state,
+            hidden_states=final_hidden_states,
+        )
+
+
+class _PixtralVisionModel(torch.nn.Module):
+    def __init__(self, model: PixtralVisionModel, output_hidden_states: bool):
+        super().__init__()
+        self.transformer = self.convert_to_rbln_pixtral_vision_model(model)
+        self.output_hidden_states = output_hidden_states
+
+    def convert_to_rbln_pixtral_vision_model(self, model: nn.Module):
+        for layer in model.transformer.layers:
+            layer.attention = PixtralAttention(layer.attention)
+        return model.transformer
+
+    def forward(self, patch_embeds, attention_mask, position_embeddings_1, position_embeddings_2):
+        output = self.transformer(
+            inputs_embeds=patch_embeds,
+            attention_mask=attention_mask,
+            position_embeddings=(position_embeddings_1, position_embeddings_2),
+            output_hidden_states=self.output_hidden_states,
+            return_dict=False,
+        )
+        return output
+
+
+class RBLNPixtralVisionModel(RBLNModel):
+    """
+    RBLN optimized Pixtral vision encoder model.
+
+    This class provides hardware-accelerated inference for Pixtral vision encoders
+    on RBLN devices, supporting image encoding for multimodal tasks.
+    """
+
+    def __post_init__(self, **kwargs):
+        artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
+        with no_init_weights():
+            self.patch_conv = nn.Conv2d(
+                in_channels=self.config.num_channels,
+                out_channels=self.config.hidden_size,
+                kernel_size=self.config.patch_size,
+                stride=self.config.patch_size,
+                bias=False,
+            )
+            self.ln_pre = PixtralRMSNorm(self.config.hidden_size, eps=1e-5)
+        self.patch_conv.load_state_dict(artifacts["patch_conv"])
+        self.ln_pre.load_state_dict(artifacts["ln_pre"])
+        self.model = RBLNRuntimePixtralVisionModel(
+            self.model[0],
+            main_input_name="pixel_values",
+            config=self.config,
+            rbln_config=self.rbln_config,
+            patch_conv=self.patch_conv,
+            ln_pre=self.ln_pre,
+        )
+
+    @classmethod
+    def save_torch_artifacts(
+        cls,
+        model: "PreTrainedModel",
+        save_dir_path: Path,
+        subfolder: str,
+        rbln_config: RBLNModelConfig,
+    ):
+        save_dict = {}
+        save_dict["patch_conv"] = model.get_input_embeddings().state_dict()
+        save_dict["ln_pre"] = model.ln_pre.state_dict()
+        torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
+
+    @classmethod
+    def wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNPixtralVisionModelConfig
+    ) -> torch.nn.Module:
+        wrapper_cfg = {
+            "output_hidden_states": rbln_config.output_hidden_states,
+        }
+        return _PixtralVisionModel(model, **wrapper_cfg).eval()
+
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        return rbln_config
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "PixtralVisionConfig" = None,
+        rbln_config: Optional[RBLNPixtralVisionModelConfig] = None,
+    ) -> RBLNPixtralVisionModelConfig:
+        if rbln_config.max_image_size is None:
+            rbln_config.max_image_size = (model_config.image_size, model_config.image_size)
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
+        num_total_patches = (rbln_config.max_image_size[0] // model_config.patch_size) * (
+            rbln_config.max_image_size[1] // model_config.patch_size
+        )
+
+        rbln_compile_config = RBLNCompileConfig(
+            input_info=[
+                (
+                    "patch_embeds",
+                    [1, num_total_patches, model_config.hidden_size],
+                    "float32",
+                ),
+                ("attention_mask", [1, num_total_patches], "float32"),
+                (
+                    "position_embeddings_1",
+                    [
+                        num_total_patches,
+                        model_config.head_dim,
+                    ],
+                    "float32",
+                ),
+                (
+                    "position_embeddings_2",
+                    [
+                        num_total_patches,
+                        model_config.head_dim,
+                    ],
+                    "float32",
+                ),
+            ]
+        )
+
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+        return rbln_config
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutput]:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        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 = self.model(
+            pixel_values, image_sizes, output_hidden_states=output_hidden_states, return_dict=return_dict
+        )
+
+        return output

optimum/rbln/transformers/models/pixtral/pixtral_architecture.py

@@ -0,0 +1,73 @@
+# 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, Tuple
+
+import torch
+import torch.nn as nn
+
+from ..decoderonly.decoderonly_architecture import apply_rotary_pos_emb
+
+
+class PixtralAttention(nn.Module):
+    def __init__(self, self_attention):
+        super().__init__()
+        self.original_model = self_attention
+        self.num_heads = getattr(self.original_model, "num_heads", None) or getattr(
+            self.original_model.config, "num_attention_heads"
+        )
+        self.head_dim = self.original_model.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.__post_init__()
+
+    def __post_init__(self):
+        self.q_proj = self.original_model.q_proj
+        self.k_proj = self.original_model.k_proj
+        self.v_proj = self.original_model.v_proj
+        self.o_proj = self.original_model.o_proj
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+    ):
+        batch_size, patches, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # TODO: return output attention
+        query_states = query_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
+        key_states = key_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
+        value_states = value_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
+
+        cos, sin = position_embeddings
+        cos = cos[None, None, None, :, :]
+        sin = sin[None, None, None, :, :]
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(3, 4)) * self.scaling
+        attn_weights = attn_weights + attention_mask
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 3)
+
+        attn_output = attn_output.reshape(batch_size, patches, -1)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, _

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig
-from .modeling_qwen2 import RBLNQwen2ForCausalLM
+from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig, RBLNQwen2ModelConfig
+from .modeling_qwen2 import RBLNQwen2ForCausalLM, RBLNQwen2Model

optimum/rbln/transformers/models/qwen2/configuration_qwen2.py

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

optimum/rbln/transformers/models/qwen2/modeling_qwen2.py

@@ -15,7 +15,11 @@
 from transformers import PretrainedConfig
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModel,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
 from .qwen2_architecture import QWEN2Wrapper
 
 
@@ -95,3 +99,25 @@ class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         rbln_config.sliding_window = model_config.sliding_window
         rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
         return rbln_config
+
+
+class RBLNQwen2Model(RBLNDecoderOnlyModel):
+    """
+    The Qwen2 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 = QWEN2Wrapper
+
+    @classmethod
+    def _update_sliding_window_config(
+        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
+    ):
+        # https://github.com/huggingface/transformers/issues/35896
+        # There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
+        # we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
+
+        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/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -34,7 +34,7 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
-from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput
+from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyForCausalLMOutput, RBLNDecoderOnlyModelForCausalLM
 from .configuration_qwen2_5_vl import (
     RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
     RBLNQwen2_5_VLForConditionalGenerationConfig,
@@ -595,7 +595,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         generate_idx: Optional[torch.Tensor] = None,
         return_dict: Optional[bool] = None,
         **kwargs,
-    ) -> RBLNDecoderOnlyOutput:
+    ) -> RBLNDecoderOnlyForCausalLMOutput:
         # Prefill
         if cache_position is None:
             inputs_embeds, position_embed, rope_deltas = self._preprocess_prefill(
@@ -637,7 +637,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         if not return_dict:
             return logits, generate_idx
         else:
-            return RBLNDecoderOnlyOutput(
+            return RBLNDecoderOnlyForCausalLMOutput(
                 logits=logits,
                 generate_idx=generate_idx,
             )

optimum/rbln/transformers/models/qwen3/configuration_qwen3.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNQwen3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -43,7 +43,7 @@ class RBLNQwen3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     """
 
 
-class RBLNQwen3ModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+class RBLNQwen3ModelConfig(RBLNDecoderOnlyModelConfig):
     """
     Configuration class for RBLN Qwen3 models.