optimum-rbln 0.7.4a4__py3-none-any.whl → 0.7.4a5__py3-none-any.whl

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -12,28 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
-from transformers import (
-    CLIPTextConfig,
-    CLIPTextModel,
-    CLIPVisionConfig,
-    CLIPVisionModel,
-)
-from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers import CLIPTextConfig, CLIPTextModel, CLIPVisionConfig, CLIPVisionModel
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
 
-from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
+from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
+from .configuration_clip import RBLNCLIPTextModelConfig, RBLNCLIPVisionModelConfig
 
 
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, CLIPTextModel
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, CLIPTextModel, PreTrainedModel
+
+    from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 
 class _TextEncoder(torch.nn.Module):
@@ -48,53 +44,55 @@ class _TextEncoder(torch.nn.Module):
 
 class RBLNCLIPTextModel(RBLNModel):
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
         return _TextEncoder(model).eval()
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_config: str
+    ) -> "RBLNDiffusionMixinConfig":
         return rbln_config
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "CLIPTextConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-        rbln_batch_size: Optional[int] = None,
-    ) -> RBLNConfig:
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        model_config.return_dict = False
-
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "CLIPTextConfig" = None,
+        rbln_config: Optional[RBLNCLIPTextModelConfig] = None,
+    ) -> RBLNCLIPTextModelConfig:
         input_info = [
             (
                 "input_ids",
                 [
-                    rbln_batch_size,
+                    rbln_config.batch_size,
                     model_config.max_position_embeddings,
                 ],
                 "int64",
             ),
         ]
 
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
         return rbln_config
 
-    def forward(self, input_ids: "torch.Tensor", **kwargs):
-        text_output = super().forward(input_ids)
-        return CLIPTextModelOutput(
-            text_embeds=text_output[0],
-            last_hidden_state=text_output[1],
-            hidden_states=text_output[2:],
-        )
+    def forward(self, input_ids: torch.LongTensor, return_dict: bool = None, **kwargs) -> torch.FloatTensor:
+        # To ignore using attention_mask, we override forward method.
+        output = super().forward(input_ids, 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:
+            return CLIPTextModelOutput(
+                text_embeds=output[0],
+                last_hidden_state=output[1],
+                hidden_states=output[2:],
+            )
 
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
@@ -113,30 +111,30 @@ class _VisionEncoder(torch.nn.Module):
 
 class RBLNCLIPVisionModel(RBLNModel):
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
         return _VisionEncoder(model).eval()
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
         return rbln_config
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "CLIPVisionConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_batch_size = rbln_kwargs.get("batch_size", 1)
-        rbln_image_size = rbln_kwargs.get("image_size", None)
-
-        if rbln_image_size is None:
-            rbln_image_size = getattr(model_config, "image_size", None)
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "CLIPVisionConfig" = None,
+        rbln_config: Optional[RBLNCLIPVisionModelConfig] = None,
+    ) -> RBLNCLIPVisionModelConfig:
+        if rbln_config.image_size is None:
+            rbln_config.image_size = getattr(model_config, "image_size", None)
 
-        if isinstance(rbln_image_size, int):
-            rbln_image_size = (rbln_image_size, rbln_image_size)
+        if isinstance(rbln_config.image_size, int):
+            rbln_config.image_size = (rbln_config.image_size, rbln_config.image_size)
 
-        if rbln_image_size is None:
+        if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
         rbln_compile_config = RBLNCompileConfig(
@@ -144,45 +142,44 @@ class RBLNCLIPVisionModel(RBLNModel):
                 (
                     "pixel_values",
                     [
-                        rbln_batch_size,
+                        rbln_config.batch_size,
                         3,
-                        rbln_image_size[0],
-                        rbln_image_size[1],
+                        rbln_config.image_height,
+                        rbln_config.image_width,
                     ],
                     "float32",
                 )
             ]
         )
 
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update(
-            {
-                "batch_size": rbln_batch_size,
-                "image_size": rbln_image_size,
-            }
-        )
-
+        rbln_config.set_compile_cfgs([rbln_compile_config])
         return rbln_config
 
     def forward(
         self,
         pixel_values: Optional[torch.FloatTensor] = None,
+        return_dict: bool = None,
         **kwargs,
-    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+    ) -> Union[Tuple, CLIPVisionModelOutput]:
         if len(kwargs) > 0 and any(kwargs.values()):
             logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
 
-        output = super().forward(pixel_values)
-        return BaseModelOutputWithPooling(
-            last_hidden_state=output[0],
-            pooler_output=output[1],
-            hidden_states=output[2:],
-        )
+        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:
+            return CLIPVisionModelOutput(
+                image_embeds=output[0],
+                last_hidden_state=output[1],
+                hidden_states=output[2:],
+            )
 
 
 class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):

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

@@ -22,4 +22,5 @@ from ....ops import (
     paged_flash_causal_attn_decode,
     paged_flash_causal_attn_prefill,
 )
+from .configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 from .modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM

optimum/rbln/transformers/models/decoderonly/configuration_decoderonly.py

@@ -0,0 +1,90 @@
+# 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, Optional
+
+import rebel
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+from ...utils.rbln_quantization import QuantizationManager
+
+
+logger = get_logger()
+
+
+class RBLNDecoderOnlyModelForCausalLMConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        max_seq_len: Optional[int] = None,
+        use_inputs_embeds: Optional[bool] = None,
+        use_attention_mask: Optional[bool] = None,
+        attn_impl: Optional[str] = None,
+        kvcache_partition_len: Optional[int] = None,
+        kvcache_block_size: Optional[int] = None,
+        quantization: Optional[Dict[str, Any]] = None,
+        prefill_chunk_size: Optional[int] = None,
+        kvcache_num_blocks: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            max_seq_len (Optional[int]): The maximum sequence length supported by the model.
+            use_inputs_embeds (Optional[bool]): Whether to use input embeddings directly. Defaults to False.
+            use_attention_mask (Optional[bool]): Whether to use attention masks. This is automatically set to True
+                for RBLN-CA02 devices.
+            attn_impl (Optional[str]): The attention implementation to use.
+            kvcache_partition_len (Optional[int]): The length of each KV cache partition.
+            kvcache_block_size (Optional[int]): The block size for KV cache.
+            quantization (Optional[Dict[str, Any]]): Configuration for model quantization.
+            prefill_chunk_size (Optional[int]): The chunk size for prefilling the KV cache. Defaults to 128,
+                and must be a positive integer divisible by 64.
+            kvcache_num_blocks (Optional[int]): The number of blocks in the KV cache.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer or if prefill_chunk_size is not
+                a positive integer divisible by 64.
+        """
+        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.max_seq_len = max_seq_len
+        self.use_inputs_embeds = use_inputs_embeds or False
+
+        self.use_attention_mask = use_attention_mask
+        npu = self.npu or rebel.get_npu_name()
+        if npu == "RBLN-CA02":
+            if self.use_attention_mask is False:
+                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
+            self.use_attention_mask = True
+        else:
+            self.use_attention_mask = self.use_attention_mask or False
+
+        self.attn_impl = attn_impl
+        self.kvcache_partition_len = kvcache_partition_len
+        self.kvcache_block_size = kvcache_block_size
+        self.quantization = quantization or {}
+        if self.quantization:
+            QuantizationManager.validate_quantization_config(self.quantization)
+
+        self.prefill_chunk_size = prefill_chunk_size or 128
+        if self.prefill_chunk_size % 64 != 0 or self.prefill_chunk_size <= 0:
+            raise ValueError("`prefill_chunk_size` must be a positive integer divisible by 64.")
+
+        self.kvcache_num_blocks = kvcache_num_blocks

optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -32,30 +32,39 @@ MIN_FLASH_ATTN_PARTITION_LENGTH = 4_096
 MAX_FLASH_ATTN_PARTITION_LENGTH = 32_768
 
 
-def validate_attention_method(
-    rbln_attn_impl: str, rbln_kvcache_partition_len: int, rbln_kvcache_block_size: int, rbln_max_seq_len: int
-) -> Tuple[str, int]:
-    if rbln_kvcache_partition_len is not None:
-        if rbln_attn_impl == "eager":
-            raise ValueError(
-                f"`rbln_kvcache_partition_len` is set to {rbln_kvcache_partition_len}, but KV cache partitioning"
-                " is not supported with 'eager' attention. Please set `rbln_kvcache_partition_len` to None, "
-                "or switch `rbln_attn_impl` to 'flash_attn' to use KV cache partitioning."
-            )
-        elif rbln_attn_impl is None:
-            rbln_attn_impl = "flash_attn"
+def set_default_values(
+    attn_impl: Optional[str] = None,
+    kvcache_partition_len: Optional[int] = None,
+    kvcache_block_size: Optional[int] = None,
+    max_seq_len: Optional[int] = None,
+) -> Tuple[str, int, int]:
+    if attn_impl is None:
+        attn_impl = "eager"
+
+    if kvcache_partition_len is not None:
+        if attn_impl == "eager":
+            attn_impl = "flash_attn"
             logger.warning(
-                "A non-null `rbln_kvcache_partition_len` was provided, but `rbln_attn_impl` was not explicitly set. "
-                "Since KV cache partitioning is only supported with flash attention, "
-                "`rbln_attn_impl` has been automatically switched to 'flash_attn'."
+                "A non-null `kvcache_partition_len` was provided, but `attn_impl` was not explicitly set or "
+                "set to 'eager'. Since KV cache partitioning is only supported with flash attention, "
+                "`attn_impl` has been automatically switched to 'flash_attn'."
             )
 
-    rbln_attn_impl = "eager" if rbln_attn_impl is None else rbln_attn_impl
-    if rbln_attn_impl not in ["eager", "flash_attn"]:
-        raise ValueError(f"Unknown `rbln_attn_impl` : {rbln_attn_impl}. (Available : 'eager', 'flash_attn`)")
+    if kvcache_partition_len is None and attn_impl == "flash_attn":
+        kvcache_partition_len = DEFAULT_FLASH_ATTN_PARTITION_LENGTH
+
+    if kvcache_block_size is None:
+        if attn_impl == "eager":
+            kvcache_block_size = max_seq_len
+        else:
+            kvcache_block_size = kvcache_partition_len
 
-    if rbln_kvcache_partition_len is None and rbln_attn_impl == "flash_attn":
-        rbln_kvcache_partition_len = DEFAULT_FLASH_ATTN_PARTITION_LENGTH
+    return attn_impl, kvcache_partition_len, kvcache_block_size
+
+
+def validate_attention_method(attn_impl: str, kvcache_partition_len: int, kvcache_block_size: int, max_seq_len: int):
+    if attn_impl not in ["eager", "flash_attn"]:
+        raise ValueError(f"Unknown `attn_impl` : {attn_impl}. (Available : 'eager', 'flash_attn`)")
 
     ## Checking Constraints...
     # Constraint of eager attention:
@@ -65,47 +74,45 @@ def validate_attention_method(
     # 1. `max_seq_len` should be multiple of `partition_len`.
     # 2. 4k <= `partition_len` <= 32k.
     # 3. `max_seq_len` should be larger then 8k.
-    if rbln_attn_impl == "eager" and rbln_max_seq_len > DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH:
+    if attn_impl == "eager" and max_seq_len > DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH:
         raise ValueError(
-            f"`rbln_max_seq_len` is set to {rbln_max_seq_len}, "
+            f"`max_seq_len` is set to {max_seq_len}, "
             f"which exceeds the limit of {DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH} for 'eager' attention. "
-            f"Please reduce the `rbln_max_seq_len` to {DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH} or lower,"
-            " or consider switching `rbln_attn_impl` to 'flash_attn' for larger sequence lengths."
+            f"Please reduce the `max_seq_len` to {DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH} or lower,"
+            " or consider switching `attn_impl` to 'flash_attn' for larger sequence lengths."
         )
 
-    if rbln_attn_impl == "flash_attn":
-        if rbln_max_seq_len // rbln_kvcache_partition_len < 2 or rbln_max_seq_len % rbln_kvcache_partition_len != 0:
+    if attn_impl == "flash_attn":
+        if max_seq_len // kvcache_partition_len < 2 or max_seq_len % kvcache_partition_len != 0:
             raise ValueError(
-                f"`rbln_max_seq_len` ({rbln_max_seq_len}) must be a multiple of `rbln_kvcache_partition_len` ({rbln_kvcache_partition_len}) "
+                f"`max_seq_len` ({max_seq_len}) must be a multiple of `kvcache_partition_len` ({kvcache_partition_len}) "
                 f"when using 'flash_attn'. Please adjust either value to meet this requirement."
             )
-        elif not (MIN_FLASH_ATTN_PARTITION_LENGTH <= rbln_kvcache_partition_len <= MAX_FLASH_ATTN_PARTITION_LENGTH):
+        elif not (MIN_FLASH_ATTN_PARTITION_LENGTH <= kvcache_partition_len <= MAX_FLASH_ATTN_PARTITION_LENGTH):
             raise ValueError(
-                f"`rbln_kvcache_partition_len` ({rbln_kvcache_partition_len}) is out of the supported range for 'flash_attn' "
-                f"({MIN_FLASH_ATTN_PARTITION_LENGTH} <= `rbln_kvcache_partition_len` <= {MAX_FLASH_ATTN_PARTITION_LENGTH}). "
+                f"`kvcache_partition_len` ({kvcache_partition_len}) is out of the supported range for 'flash_attn' "
+                f"({MIN_FLASH_ATTN_PARTITION_LENGTH} <= `kvcache_partition_len` <= {MAX_FLASH_ATTN_PARTITION_LENGTH}). "
                 f"Please provide a valid value within this range."
             )
-        elif rbln_max_seq_len < MIN_FLASH_ATTN_MAX_SEQ_LEN:
+        elif max_seq_len < MIN_FLASH_ATTN_MAX_SEQ_LEN:
             raise ValueError(
-                f"`rbln_max_seq_len` ({rbln_max_seq_len}) is too small for 'flash_attn'. The minimum "
-                f"supported value is {MIN_FLASH_ATTN_MAX_SEQ_LEN}. Please increase `rbln_max_seq_len` to meet "
-                "this requirement, or consider switching `rbln_attn_impl` to 'eager' for shorter lengths."
+                f"`max_seq_len` ({max_seq_len}) is too small for 'flash_attn'. The minimum "
+                f"supported value is {MIN_FLASH_ATTN_MAX_SEQ_LEN}. Please increase `max_seq_len` to meet "
+                "this requirement, or consider switching `attn_impl` to 'eager' for shorter lengths."
             )
 
-    if rbln_kvcache_block_size is not None:
-        if rbln_attn_impl == "flash_attn" and rbln_kvcache_partition_len != rbln_kvcache_block_size:
+    if kvcache_block_size is not None:
+        if attn_impl == "flash_attn" and kvcache_partition_len != kvcache_block_size:
             raise ValueError(
-                f" When using 'flash attention', the `rbln_kvcache_block_size` ({rbln_kvcache_block_size})  "
-                f"must always be set equal to the `rbln_kvcache_partition_len` {rbln_kvcache_partition_len}."
+                f" When using 'flash attention', the `kvcache_block_size` ({kvcache_block_size})  "
+                f"must always be set equal to the `kvcache_partition_len` {kvcache_partition_len}."
             )
-        elif rbln_attn_impl == "eager" and rbln_kvcache_block_size != rbln_max_seq_len:
+        elif attn_impl == "eager" and kvcache_block_size != max_seq_len:
             raise ValueError(
-                f" When using 'eager attention', the `rbln_kvcache_block_size` ({rbln_kvcache_block_size})  "
-                f"must always be set equal to the `rbln_max_seq_len` {rbln_max_seq_len}."
+                f" When using 'eager attention', the `kvcache_block_size` ({kvcache_block_size})  "
+                f"must always be set equal to the `max_seq_len` {max_seq_len}."
             )
 
-    return rbln_attn_impl, rbln_kvcache_partition_len, rbln_kvcache_block_size
-
 
 class DecoderOnlyWrapper(nn.Module):
     """A wrapper class for decoder-only language models that handles RBLN-specific optimizations and requirements.