optimum-rbln 0.7.5a0__py3-none-any.whl → 0.7.5a1__py3-none-any.whl

optimum/rbln/transformers/models/opt/modeling_opt.py

@@ -0,0 +1,78 @@
+# 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.
+
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from ....utils import logging
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from .opt_architecture import OPTWrapper
+
+
+logger = logging.get_logger(__name__)
+
+
+class MLP(nn.Module):
+    def __init__(self, fc1, fc2, activation_fn):
+        super(MLP, self).__init__()
+        self.fc1 = fc1
+        self.fc2 = fc2
+        self.activation_fn = activation_fn
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.activation_fn(x)
+        x = self.fc2(x)
+        return x
+
+
+class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The OPT 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 OPTForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers OPTForCausalLM 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.
+    """
+
+    _decoder_wrapper_cls = OPTWrapper
+    _use_rotary_emb = False
+
+    def modify_opt_decoder_layer(layer):
+        mlp = MLP(layer.fc1, layer.fc2, layer.activation_fn)
+        layer.mlp = mlp
+        del layer.fc1
+        del layer.fc2
+        del layer.activation_fn
+
+        return layer
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+        wrapper_cfg = {
+            "max_seq_len": rbln_config.max_seq_len,
+            "attn_impl": rbln_config.attn_impl,
+            "kvcache_partition_len": rbln_config.kvcache_partition_len,
+            "kvcache_block_size": rbln_config.kvcache_block_size,
+            "use_rotary_emb": cls._use_rotary_emb,
+            "use_attention_mask": rbln_config.use_attention_mask,
+        }
+
+        for i in range(len(model.model.decoder.layers)):
+            model.model.decoder.layers[i] = cls.modify_opt_decoder_layer(model.model.decoder.layers[i])
+
+        return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()

optimum/rbln/transformers/models/opt/opt_architecture.py

@@ -0,0 +1,74 @@
+# 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
+
+import torch.nn as nn
+
+from ...models.decoderonly.decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyForCausalLM,
+    DecoderOnlyLayer,
+    DecoderOnlyModel,
+    DecoderOnlyWrapper,
+)
+
+
+if TYPE_CHECKING:
+    from transformers import OPTForCausalLM
+
+
+class OPTWrapper(DecoderOnlyWrapper):
+    def convert_to_rbln_causal_lm(self, causal_lm: "OPTForCausalLM", max_seq_len: int):
+        if self.attn_impl != "eager":
+            raise NotImplementedError(f"flash attention ({self.attn_impl}) is not implemented for {self.__class__}")
+
+        new_layers = []
+
+        for layer in causal_lm.model.decoder.layers:
+            new_self_attn = OPTAttention(
+                layer.self_attn, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+            )
+            new_layer = OPTDecoderLayer(layer, new_self_attn)
+            new_layers.append(new_layer)
+        new_model = OPTModel(causal_lm.model.decoder, new_layers, max_seq_len=max_seq_len, use_learned_pos_emb=True)
+        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
+        return new_causal_lm
+
+
+class OPTAttention(DecoderOnlyAttention):
+    def __post_init__(self):
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.q_proj = self._original_mod.q_proj
+        self.o_proj = self._original_mod.out_proj
+
+
+class OPTModel(DecoderOnlyModel):
+    def get_embedding(self) -> nn.Embedding:
+        return self._original_mod.embed_tokens
+
+    def get_pos_embedding(self):
+        return self._original_mod.embed_positions
+
+    def get_last_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.final_layer_norm
+
+
+class OPTDecoderLayer(DecoderOnlyLayer):
+    def get_pre_attention_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.self_attn_layer_norm
+
+    def get_post_attention_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.final_layer_norm

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

@@ -371,6 +371,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         query_length: int,
         use_inputs_embeds: bool,
         use_attention_mask: bool,
+        use_position_ids: bool,
         max_seq_len: int,
         kvcache_block_size: int,
         kvcache_num_blocks: int,
@@ -384,6 +385,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             query_length,
             use_inputs_embeds,
             use_attention_mask,
+            use_position_ids,
             max_seq_len,
             kvcache_block_size,
             kvcache_num_blocks,
@@ -392,8 +394,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             hidden_size,
             head_dim,
         )
-        pos_idx = 5 if query_length > 1 else 4
-        pos_idx = pos_idx if use_attention_mask else pos_idx - 1
+        pos_idx = 4 if query_length > 1 else 5
         input_info.insert(pos_idx, ("position_emb", [2, batch_size, 1, query_length, head_dim], "float32"))
 
         return input_info
@@ -562,7 +563,8 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
         video_grid_thw: Optional[torch.LongTensor] = None,
         cache_position: Optional[torch.LongTensor] = None,
         second_per_grid_ts: Optional[torch.Tensor] = None,
-        generate_idx: torch.Tensor = None,
+        generate_idx: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
         **kwargs,
     ) -> RBLNDecoderOnlyOutput:
         # Prefill
@@ -584,25 +586,29 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
             for b_idx in range(batch_size):
                 cache_position = torch.arange(0, generate_idx[b_idx].item(), dtype=torch.int32).unsqueeze(0)
 
-                logit = self.prefill_decoder(
+                output = self.prefill_decoder(
                     inputs_embeds=inputs_embeds[b_idx : b_idx + 1],
                     attention_mask=attention_mask[b_idx] if attention_mask is not None else None,
                     cache_position=cache_position,
                     batch_idx=b_idx,
                     position_embed=position_embed[:, b_idx : b_idx + 1],
                 )
-                logits.append(logit)
+                logits.append(output.logits)
             logits = torch.cat(logits, dim=0)
         # Decoder
         else:
             inputs_embeds, position_embed = self._preprocess_decoder(input_ids, cache_position)
-            logits = self.decoder(
+            output = self.decoder(
                 inputs_embeds=inputs_embeds,
                 cache_position=cache_position,
                 position_embed=position_embed,
             )
+            logits = output.logits
 
-        return RBLNDecoderOnlyOutput(
-            logits=logits,
-            generate_idx=generate_idx,
-        )
+        if not return_dict:
+            return logits, generate_idx
+        else:
+            return RBLNDecoderOnlyOutput(
+                logits=logits,
+                generate_idx=generate_idx,
+            )

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

@@ -157,58 +157,41 @@ class Qwen2_5_VLVisionWindowAttention(nn.Module):
 
 
 class Qwen2_5_VL_LanguageModelWrapper(DecoderOnlyWrapper):
-    def forward(self, *args):
-        if self.phase == "decode":
-            if self.use_attention_mask:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    attention_mask,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-            else:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-                attention_mask = None
-            query_position = None
-        elif self.phase == "prefill":
-            if self.use_attention_mask:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    attention_mask,
-                    query_position,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-            else:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    query_position,
-                    block_tables,
-                    position_emb,
-                    *past_key_values,
-                ) = args
-                attention_mask = None
-
-        else:
-            raise ValueError(f"Unknown phase: {self.phase}")
-
-        return self.forward_common(
-            input_ids_or_inputs_embeds,
+    def prepare_forward_args(self, *args):
+        args = list(args)
+        input_ids = None if self.use_inputs_embeds else args.pop(0)
+        inputs_embeds = args.pop(0) if self.use_inputs_embeds else None
+        cache_position = args.pop(0)
+        block_tables = args.pop(0)
+        position_embeds = args.pop(0)
+        query_position = args.pop(0) if self.phase == "prefill" else None
+        position_ids = None
+        attention_mask = args.pop(0) if self.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,
-            attention_mask,
-            query_position,
             block_tables,
-            position_emb,
-            *past_key_values,
+            query_position,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            position_embeds,
         )

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

@@ -476,6 +476,8 @@ class Seq2SeqSelfAttention(nn.Module):
         ]
         if attention_mask is not None:
             args.insert(3, attention_mask.unsqueeze(2))
+        else:
+            args.append(None)
 
         attn_output = self.attn_decode(*args)
 

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

@@ -0,0 +1,20 @@
+# 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_siglip import (
+    RBLNSiglipVisionModelConfig,
+)
+from .modeling_siglip import (
+    RBLNSiglipVisionModel,
+)

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

@@ -0,0 +1,66 @@
+# 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
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNSiglipVisionModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        image_size: Optional[int] = None,
+        interpolate_pos_encoding: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for image processing. Defaults to 1.
+            image_size (Optional[int]): The size of input images. Can be an integer for square images,
+                a tuple/list (height, width), or a dictionary with 'height' and 'width' keys.
+            interpolate_pos_encoding (Optional[bool]): Whether to interpolate the position encoding.
+            output_hidden_states: (Optional[bool]): Whether to return hidden states.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        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
+        self.interpolate_pos_encoding = interpolate_pos_encoding or False
+        self.output_hidden_states = output_hidden_states
+
+    @property
+    def image_width(self):
+        if isinstance(self.image_size, int):
+            return self.image_size
+        elif isinstance(self.image_size, (list, tuple)):
+            return self.image_size[1]
+        else:
+            return self.image_size["width"]
+
+    @property
+    def image_height(self):
+        if isinstance(self.image_size, int):
+            return self.image_size
+        elif isinstance(self.image_size, (list, tuple)):
+            return self.image_size[0]
+        else:
+            return self.image_size["height"]

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

@@ -0,0 +1,146 @@
+# 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, Tuple, Union
+
+import torch
+from transformers import SiglipVisionConfig, SiglipVisionModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.siglip.modeling_siglip import SiglipVisionModelOutput
+
+from ....configuration_utils import RBLNCompileConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from .configuration_siglip import RBLNSiglipVisionModelConfig
+
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+    from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
+
+
+class _SiglipVisionModel(torch.nn.Module):
+    def __init__(self, model: SiglipVisionModel, interpolate_pos_encoding: bool, output_hidden_states: bool):
+        super().__init__()
+        self.vision_model = model.vision_model
+        self.interpolate_pos_encoding = interpolate_pos_encoding
+        self.output_hidden_states = output_hidden_states
+
+    def forward(self, inp):
+        enc_out = self.vision_model(
+            inp,
+            output_hidden_states=self.output_hidden_states,
+            return_dict=False,
+            interpolate_pos_encoding=self.interpolate_pos_encoding,
+        )
+        return tuple(x for x in enc_out if x is not None)
+
+
+class RBLNSiglipVisionModel(RBLNModel):
+    @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig) -> torch.nn.Module:
+        wrapper_cfg = {
+            "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
+            "output_hidden_states": rbln_config.output_hidden_states,
+        }
+        return _SiglipVisionModel(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: "SiglipVisionConfig" = None,
+        rbln_config: Optional[RBLNSiglipVisionModelConfig] = None,
+    ) -> RBLNSiglipVisionModelConfig:
+        if rbln_config.image_size is None:
+            rbln_config.image_size = getattr(model_config, "image_size", None)
+
+        if isinstance(rbln_config.image_size, int):
+            rbln_config.image_size = (rbln_config.image_size, rbln_config.image_size)
+        if rbln_config.image_size is None:
+            raise ValueError("`rbln_image_size` should be specified!")
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = model_config.output_hidden_states
+
+        rbln_compile_config = RBLNCompileConfig(
+            input_info=[
+                (
+                    "pixel_values",
+                    [
+                        rbln_config.batch_size,
+                        3,
+                        rbln_config.image_height,
+                        rbln_config.image_width,
+                    ],
+                    "float32",
+                )
+            ]
+        )
+
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+        return rbln_config
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        return_dict: bool = None,
+        interpolate_pos_encoding: bool = False,
+        **kwargs,
+    ) -> Union[Tuple, SiglipVisionModelOutput]:
+        if len(kwargs) > 0 and any(kwargs.values()):
+            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
+
+        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}"
+                f"Please compile again with the correct argument."
+            )
+        output = super().forward(pixel_values, return_dict=return_dict)
+        return output
+
+    def _prepare_output(self, output, return_dict):
+        """
+        Prepare model output based on return_dict flag.
+        This method can be overridden by subclasses to provide task-specific output handling.
+        """
+        if not return_dict:
+            return (output,) if not isinstance(output, (tuple, list)) else output
+        else:
+            last_hidden_state = (
+                output[0]
+                if self.rbln_config.interpolate_pos_encoding or self.rbln_config.output_hidden_states
+                else output
+            )
+            pooler_output = output[1] if self.rbln_config.interpolate_pos_encoding else None
+            if self.rbln_config.output_hidden_states:
+                hidden_states = (output[2:] if self.rbln_config.interpolate_pos_encoding else output[1:],)
+            else:
+                hidden_states = None
+
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+            )

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

@@ -313,6 +313,7 @@ class WhisperSelfAttention(WhisperAttention):
             args["mask"] = attention_mask.unsqueeze(2)
             attn_output = torch.ops.rbln_custom_ops.paged_attn_decode(**args)
         else:
+            args["mask"] = None
             attn_output = torch.ops.rbln_custom_ops.paged_causal_attn_decode(**args)
 
         attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)