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

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

@@ -0,0 +1,220 @@
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper, apply_rotary_pos_emb
+
+
+class Qwen2_5_VisionTransformerWrapper(nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self._original_mod = model
+        self.fullatt_block_indexes = model.fullatt_block_indexes
+        self.merger = model.merger
+        window_seq_len = (model.window_size // model.patch_size) ** 2
+        self.blocks = self.wrap_vision_blocks(model.blocks, window_seq_len)
+
+    def wrap_vision_blocks(self, blocks: torch.nn.ModuleList, window_seq_len: int):
+        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))
+        return nn.ModuleList(wrapped_blocks)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        full_attn_masks: torch.Tensor,
+        window_attn_masks: torch.Tensor,
+        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
+
+        for i, block in enumerate(self.blocks):
+            attn_masks = full_attn_masks if i in self.fullatt_block_indexes else window_attn_masks
+            hidden_states = block(hidden_states, attn_masks, [cos, sin])
+
+        hidden_states = self.merger(hidden_states)
+
+        return hidden_states
+
+
+class Qwen2_5_VLVisionBlock(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module, is_full_attn: bool, window_seq_len: int):
+        super().__init__()
+        self._origin_model = model
+        self.norm1 = model.norm1
+        self.norm2 = model.norm2
+
+        if is_full_attn:
+            self.attn = Qwen2_5_VLVisionFullAttention(model.attn)
+        else:
+            self.attn = Qwen2_5_VLVisionWindowAttention(model.attn, window_seq_len)
+        self.mlp = model.mlp
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_masks: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            attn_masks,
+            position_embeddings,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+        return hidden_states
+
+
+class Qwen2_5_VLVisionFullAttention(nn.Module):
+    def __init__(self, model: nn.Module) -> None:
+        super().__init__()
+        self._origin_model = model
+        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
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_masks: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        hidden_states = hidden_states.unsqueeze(0)
+        q, k, v = (
+            self.qkv(hidden_states).reshape(1, seq_length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0)
+        )
+
+        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 = attn_weights + attn_masks
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(1, seq_length, -1)
+        attn_output = self.proj(attn_output).squeeze(0)
+
+        return attn_output
+
+
+class Qwen2_5_VLVisionWindowAttention(nn.Module):
+    def __init__(self, model: nn.Module, window_seq_len: int) -> None:
+        super().__init__()
+        self._origin_model = model
+        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
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_masks: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        num_windows = seq_length // self.window_seq_len
+
+        window_hidden_states = []
+        for i in range(0, seq_length, self.window_seq_len):
+            window_hidden_states.append(hidden_states[i : i + self.window_seq_len])
+        hidden_states = torch.stack(window_hidden_states)
+
+        q, k, v = (
+            self.qkv(hidden_states)
+            .reshape(num_windows, self.window_seq_len, 3, self.num_heads, -1)
+            .permute(2, 0, 3, 1, 4)
+            .unbind(0)
+        )
+        cos, sin = position_embeddings
+        cos = cos.reshape(num_windows, 1, seq_length // num_windows, -1)
+        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 = attn_weights + attn_masks
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(1, seq_length, -1)
+        attn_output = self.proj(attn_output).squeeze(0)
+
+        return attn_output
+
+
+class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
+    def prepare_forward_args(self, *args):
+        args = list(args)
+        input_ids = None if self.rbln_config.use_inputs_embeds else args.pop(0)
+        inputs_embeds = args.pop(0) if self.rbln_config.use_inputs_embeds else None
+        cache_position = args.pop(0)
+        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
+        position_ids = None
+        lora_int_id = None
+        attention_mask = args.pop(0) if self.rbln_config.use_attention_mask else None
+        past_key_values = args
+
+        if len(past_key_values) != 2 * self.num_hidden_layers:
+            raise ValueError(
+                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.num_hidden_layers}"
+            )
+
+        # [key, value] * n_layer -> ( (key, value) ) * n_layer
+        # cache shape : batch, n_heads, 1, max_seq_len, head_dim
+        _past_key_values = []
+        for i in range(self.config.num_hidden_layers):
+            key_states = past_key_values[i * 2]
+            value_states = past_key_values[i * 2 + 1]
+            past_key_value = [key_states, value_states]
+            _past_key_values.append(past_key_value)
+        past_key_values = _past_key_values
+
+        return (
+            input_ids,
+            inputs_embeds,
+            cache_position,
+            global_block_tables,
+            local_block_tables,
+            query_position,
+            attention_mask,
+            position_ids,
+            lora_int_id,
+            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_vl/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_qwen2_vl import (
+    RBLNQwen2VisionTransformerPretrainedModelConfig,
+    RBLNQwen2VLForConditionalGenerationConfig,
+)
+from .modeling_qwen2_vl import RBLNQwen2VisionTransformerPretrainedModel, RBLNQwen2VLForConditionalGeneration

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

@@ -0,0 +1,88 @@
+# 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 Any, Dict, List, Optional, Union
+
+from ....configuration_utils import RBLNModelConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+
+
+class RBLNQwen2VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    submodules = ["visual"]
+
+    def __init__(
+        self,
+        use_inputs_embeds: bool = True,
+        visual: Optional[RBLNModelConfig] = None,
+        **kwargs: Dict[str, Any],
+    ):
+        """
+        Args:
+            use_inputs_embeds (bool): Whether or not to use `inputs_embeds` as input. Defaults to `True`.
+            visual (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            kwargs: Additional arguments passed to the parent `RBLNDecoderOnlyModelForCausalLMConfig`.
+
+        Raises:
+            ValueError: If `use_inputs_embeds` is False.
+            ValueError: If the visual configuration is provided but contains invalid settings, such as an invalid max_seq_lens (e.g., not a positive integer or insufficient for the expected resolution).
+            ValueError: If visual is None and no default vision configuration can be inferred for the model architecture.
+            ValueError: If any inherited parameters violate constraints defined in the parent class, such as batch_size not being a positive integer, prefill_chunk_size not being divisible by 64, or max_seq_len not meeting requirements for Flash Attention.
+        """
+        super().__init__(use_inputs_embeds=use_inputs_embeds, **kwargs)
+        if not self.use_inputs_embeds:
+            raise ValueError(
+                "RBLNQwen2VLForConditionalGenerationConfig does not allow `use_inputs_embeds` to be set to False, "
+                "as RBLNQwen2VLForConditionalGeneration accepts only `inputs_embeds` as input."
+            )
+        self.visual = visual
+
+
+class RBLNQwen2VisionTransformerPretrainedModelConfig(RBLNModelConfig):
+    def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Dict[str, Any]):
+        """
+        Args:
+            max_seq_lens (Optional[Union[int, List[int]]]): Maximum sequence lengths for Vision
+                Transformer attention. Can be an integer or list of integers, each indicating
+                the number of patches in a sequence for an image or video. For example, an image
+                of 224x224 pixels with patch size 14 results in (224/14) * (224/14) = 256 patches,
+                so `max_seq_lens` must be at least 256. RBLN optimization runs inference per image
+                or video frame, so set `max_seq_lens` to match the maximum expected resolution to
+                optimize computation. If not provided, a `ValueError` is raised.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+            ValueError: If `max_seq_lens` (or any value in the list) is not a positive integer.
+            ValueError: If `max_seq_lens` is insufficient for the expected image/video resolution.
+            ValueError: If `batch_size` (inherited from RBLNModelConfig) is not a positive integer.
+
+        Max Seq Lens:
+            Since `Qwen2VLForConditionalGeneration` performs inference on a per-image or per-frame basis,
+            `max_seq_lens` should be set based on the maximum expected resolution of the input images or video frames.
+
+            The value must be greater than or equal to the number of patches generated from the input image.
+            For example, a 224x224 image with a patch size of 14 results in (224 / 14) * (224 / 14) = 256 patches.
+            Therefore, `max_seq_lens` must be at least 256.
+        """
+        super().__init__(**kwargs)
+
+        if max_seq_lens is not None:
+            if isinstance(max_seq_lens, int):
+                max_seq_lens = [max_seq_lens]
+            elif isinstance(max_seq_lens, list):
+                max_seq_lens.sort(reverse=True)
+        else:
+            raise ValueError("'max_seq_lens' must be specified.")
+
+        self.max_seq_lens = max_seq_lens