optimum-rbln 0.8.2a2__tar.gz → 0.8.2a4__tar.gz

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/.github/workflows/rbln_optimum_pytest.yaml

@@ -20,12 +20,19 @@ on:
         required: false
         type: string
         default: "default"
+      enable_hf_hub_tests:
+        description: "Whether to enable HF Hub tests (requires HF credentials)"
+        required: false
+        type: boolean
+        default: false
 
 env:
   REBEL_PYPI_ENDPOINT: ${{ vars.REBEL_PYPI_INTERNAL_ENDPOINT }}
   REBEL_PYPI_USERNAME: ${{ secrets.REBEL_PYPI_USERNAME }}
   REBEL_PYPI_PASSWORD: ${{ secrets.REBEL_PYPI_PASSWORD }}
   HF_HOME: ${{ secrets.HF_HOME }}
+  HF_USER_ID: ${{ inputs.enable_hf_hub_tests && secrets.HF_USER_ID || '' }}
+  HF_AUTH_TOKEN: ${{ inputs.enable_hf_hub_tests && secrets.HF_AUTH_TOKEN || '' }}
 
 jobs:
   pytest:

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/.github/workflows/rbln_scheduled_test.yaml

@@ -47,6 +47,7 @@ jobs:
       ref: main
       rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
       test_level: "full"
+      enable_hf_hub_tests: true
     secrets: inherit
 
   optimum-rbln-inference-test:

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.2a2
+Version: 0.8.2a4
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.2a2'
-__version_tuple__ = version_tuple = (0, 8, 2, 'a2')
+__version__ = version = '0.8.2a4'
+__version_tuple__ = version_tuple = (0, 8, 2, 'a4')

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/diffusers/modeling_diffusers.py

@@ -70,8 +70,6 @@ class RBLNDiffusionMixin:
     _submodules = []
     _optional_submodules = []
     _prefix = {}
-    _rbln_config_class = None
-    _hf_class = None
 
     @staticmethod
     def _maybe_apply_and_fuse_lora(
@@ -114,14 +112,14 @@ class RBLNDiffusionMixin:
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
         # Lazily loads and caches the corresponding RBLN model config class.
-        if cls._rbln_config_class is None:
+        if "_rbln_config_class" not in cls.__dict__ or cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
             cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
 
     @classmethod
     def get_hf_class(cls):
-        if cls._hf_class is None:
+        if "_hf_class" not in cls.__dict__ or cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
             library = importlib.import_module("diffusers")
             cls._hf_class = getattr(library, hf_cls_name, None)

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/modeling.py

@@ -35,8 +35,6 @@ logger = get_logger(__name__)
 
 
 class RBLNModel(RBLNBaseModel):
-    _output_class = None
-
     @classmethod
     def update_kwargs(cls, kwargs):
         # Update user-given kwargs to get proper pytorch model.
@@ -289,7 +287,7 @@ class RBLNModel(RBLNBaseModel):
     @classmethod
     def get_hf_output_class(cls):
         # Dynamically gets the output class from the corresponding HuggingFace model class.
-        if cls._output_class:
+        if "_output_class" in cls.__dict__ and cls._output_class is not None:
             return cls._output_class
 
         hf_class = cls.get_hf_class()

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/modeling_base.py

@@ -23,9 +23,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 import rebel
 import torch
 from transformers import AutoConfig, AutoModel, GenerationConfig, PretrainedConfig
+from transformers.utils.hub import PushToHubMixin
 
 from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig, get_rbln_config_class
-from .utils.hub import PushToHubMixin, pull_compiled_model_from_hub, validate_files
+from .utils.hub import pull_compiled_model_from_hub, validate_files
 from .utils.logging import get_logger
 from .utils.runtime_utils import UnavailableRuntime, tp_and_devices_are_ok
 from .utils.save_utils import maybe_load_preprocessors
@@ -50,11 +51,8 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     model_type = "rbln_model"
     auto_model_class = AutoModel
     config_class = AutoConfig
-
     config_name = "config.json"
     hf_library_name = "transformers"
-    _hf_class = None
-    _rbln_config_class = None
 
     def __init__(
         self,
@@ -115,7 +113,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     def _load_compiled_model_dir(
         cls,
         model_id: Union[str, Path],
-        use_auth_token: Optional[Union[bool, str]] = None,
+        token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,
         cache_dir: Optional[str] = None,
@@ -134,7 +132,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             model_path = pull_compiled_model_from_hub(
                 model_id=model_id,
                 subfolder=subfolder,
-                use_auth_token=use_auth_token,
+                token=token,
                 revision=revision,
                 cache_dir=cache_dir,
                 force_download=force_download,
@@ -172,7 +170,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         cls,
         model_id: Union[str, Path],
         config: Optional["PretrainedConfig"] = None,
-        use_auth_token: Optional[Union[bool, str]] = None,
+        token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,
         cache_dir: Optional[str] = None,
@@ -189,7 +187,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         if rbln_compiled_models is None:
             model_path_subfolder = cls._load_compiled_model_dir(
                 model_id=model_id,
-                use_auth_token=use_auth_token,
+                token=token,
                 revision=revision,
                 force_download=force_download,
                 cache_dir=cache_dir,
@@ -232,7 +230,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                         cache_dir=cache_dir,
                         force_download=force_download,
                         revision=revision,
-                        token=use_auth_token,
+                        token=token,
                         trust_remote_code=trust_remote_code,
                     )
                 elif cls.hf_library_name == "diffusers":
@@ -250,7 +248,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                         force_download=force_download,
                         local_files_only=local_files_only,
                         revision=revision,
-                        token=use_auth_token,
+                        token=token,
                         subfolder=subfolder,
                     )
                     config = PretrainedConfig(**config)
@@ -421,7 +419,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
         # Returns:
         #     type: The original HuggingFace model class
-        if cls._hf_class is None:
+        if "_hf_class" not in cls.__dict__ or cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
             library = importlib.import_module(cls.hf_library_name)
             cls._hf_class = getattr(library, hf_cls_name, None)
@@ -430,7 +428,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
         # Lazily loads and caches the corresponding RBLN model config class.
-        if cls._rbln_config_class is None:
+        if "_rbln_config_class" not in cls.__dict__ or cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
             cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
@@ -507,6 +505,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 f"Please ensure the model directory exists and you have the necessary permissions to access it."
             )
 
+        if isinstance(self.config, PretrainedConfig):
+            self.config.save_pretrained(real_save_dir)
+
         if save_directory_path == real_save_dir:
             raise FileExistsError(
                 f"Cannot save model to '{save_directory}'. This directory already exists and contains the model files."
@@ -534,7 +535,10 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             raise e  # Re-raise the exception after cleanup
 
         if push_to_hub:
-            return super().push_to_hub(str(save_directory_path), **kwargs)
+            repo_id = kwargs.pop("repo_id", None)
+            if repo_id is None:
+                raise ValueError("`repo_id` must be provided to push the model to the HuggingFace model hub.")
+            return super().push_to_hub(repo_id=repo_id, **kwargs)
 
     @staticmethod
     def _raise_missing_compiled_file_error(missing_files: List[str]):

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/transformers/models/bart/bart_architecture.py

@@ -56,10 +56,7 @@ class BartDecoderWrapper(Seq2SeqDecoderWrapper):
 
 
 class BartForConditionalGeneration(Seq2SeqForConditionalGeneration):
-    has_rescaling = False
-
-    def __post_init__(self):
-        self.scaling = self.config.d_model**-0.5
+    pass
 
 
 class BartDecoder(Seq2SeqDecoder):

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/transformers/models/clip/configuration_clip.py

@@ -43,7 +43,15 @@ class RBLNCLIPTextModelWithProjectionConfig(RBLNCLIPTextModelConfig):
 
 
 class RBLNCLIPVisionModelConfig(RBLNModelConfig):
-    def __init__(self, batch_size: Optional[int] = None, image_size: Optional[int] = None, **kwargs: Dict[str, Any]):
+    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,
+        output_attentions: Optional[bool] = None,
+        **kwargs: Dict[str, Any],
+    ):
         """
         Args:
             batch_size (Optional[int]): The batch size for image processing. Defaults to 1.
@@ -60,6 +68,9 @@ class RBLNCLIPVisionModelConfig(RBLNModelConfig):
             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
+        self.output_attentions = output_attentions
 
     @property
     def image_width(self):

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -16,6 +16,7 @@ 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.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
 
 from ....configuration_utils import RBLNCompileConfig
@@ -111,12 +112,27 @@ class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
 
 
 class _VisionEncoder(torch.nn.Module):
-    def __init__(self, enc: CLIPVisionModel):
+    def __init__(
+        self,
+        enc: CLIPVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
         super().__init__()
         self.enc = enc
+        self.interpolate_pos_encoding = interpolate_pos_encoding
+        self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     def forward(self, inp):
-        enc_out = self.enc(inp, output_hidden_states=True, return_dict=False)
+        enc_out = self.enc(
+            inp,
+            output_hidden_states=self.output_hidden_states,
+            interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
+            return_dict=False,
+        )
         return enc_out
 
 
@@ -130,7 +146,12 @@ class RBLNCLIPVisionModel(RBLNModel):
 
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
-        return _VisionEncoder(model).eval()
+        wrapper_cfg = {
+            "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
+            "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
+        }
+        return _VisionEncoder(model, **wrapper_cfg).eval()
 
     @classmethod
     def update_rbln_config_using_pipe(
@@ -155,6 +176,12 @@ class RBLNCLIPVisionModel(RBLNModel):
         if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
         rbln_compile_config = RBLNCompileConfig(
             input_info=[
                 (
@@ -176,27 +203,76 @@ class RBLNCLIPVisionModel(RBLNModel):
     def forward(
         self,
         pixel_values: Optional[torch.FloatTensor] = None,
-        return_dict: bool = None,
+        return_dict: bool = True,
+        output_attentions: bool = None,
+        output_hidden_states: bool = None,
+        interpolate_pos_encoding: bool = False,
         **kwargs,
-    ) -> Union[Tuple, CLIPVisionModelOutput]:
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
         if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
             logger.warning(
                 f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
             )
+
+        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        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."
+            )
+
+        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.
+        last_hidden_state = output.pop(0)
+        pooler_output = output.pop(0)
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
 
         if not return_dict:
-            return (output,) if not isinstance(output, (tuple, list)) else output
+            return tuple(
+                item for item in (last_hidden_state, pooler_output, hidden_states, attentions) if item is not None
+            )
         else:
-            return CLIPVisionModelOutput(
-                image_embeds=output[0],
-                last_hidden_state=output[1],
-                hidden_states=output[2:],
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+                attentions=attentions,
             )
 
 
@@ -208,21 +284,40 @@ class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
     multimodal embedding alignment tasks.
     """
 
-    def forward(
-        self,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        **kwargs,
-    ) -> 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)
-        image_embeds = output[0]
-        last_hidden_state = output[1]
-        hidden_states = output[2:]
-
-        return CLIPVisionModelOutput(
-            image_embeds=image_embeds,
-            last_hidden_state=last_hidden_state,
-            hidden_states=hidden_states,
-        )
+    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.
+
+        image_embeds = output.pop(0) if isinstance(output, (tuple, list)) else output
+        last_hidden_state = output.pop(0)
+
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
+
+        if not return_dict:
+            return tuple(
+                item for item in (image_embeds, last_hidden_state, hidden_states, attentions) if item is not None
+            )
+
+        else:
+            return CLIPVisionModelOutput(
+                image_embeds=image_embeds,
+                last_hidden_state=last_hidden_state,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -303,6 +303,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed = (
                 position_embed[:, :, :, attention_mask.bool(), :] if attention_mask is not None else position_embed
             )
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids[:, attention_mask.bool()] if attention_mask is not None else token_type_ids
 
         query_length = inputs.shape[1]
         if query_length > self.rbln_config.max_seq_len:
@@ -352,8 +354,11 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             if position_embed is not None:
                 position_embed = torch.nn.functional.pad(position_embed, (0, 0, 0, padding_size))
 
+            if token_type_ids is not None:
+                token_type_ids = torch.nn.functional.pad(token_type_ids, (0, padding_size), value=-1)
+
         # Overwrite position_ids and padded_cache_lengths
-        position_ids = None
+        position_ids = cache_position.clone()
         padded_cache_lengths = 0
 
         return (
@@ -365,6 +370,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed,
             padded_cache_lengths,
             query_length,
+            token_type_ids,
         )
 
     def prefill_forward(
@@ -393,6 +399,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             position_embed,
             padded_cache_lengths,
             query_length,
+            token_type_ids,
         ) = self._prepare_prefill_inputs(
             inputs, cache_position, attention_mask, position_embed, token_type_ids=token_type_ids
         )

{optimum_rbln-0.8.2a2 → optimum_rbln-0.8.2a4}/src/optimum/rbln/transformers/models/gemma3/configuration_gemma3.py

@@ -23,22 +23,21 @@ from ..siglip.configuration_siglip import RBLNSiglipVisionModelConfig
 class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     def __init__(
         self,
-        prefill_chunk_size: Optional[int] = None,
         use_position_ids: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
+        image_prefill_chunk_size: Optional[int] = None,
         **kwargs: Dict[str, Any],
     ):
         # use_attention_mask and use_position_ids are always True for Gemma3
         use_attention_mask = use_attention_mask or True
         use_position_ids = use_position_ids or True
-        prefill_chunk_size = prefill_chunk_size or 256
 
         super().__init__(
-            prefill_chunk_size=prefill_chunk_size,
             use_attention_mask=use_attention_mask,
             use_position_ids=use_position_ids,
             **kwargs,
         )
+        self.image_prefill_chunk_size = image_prefill_chunk_size
 
         npu = self.npu or rebel.get_npu_name()
         if npu == "RBLN-CA02":