optimum-rbln 0.8.4a0__tar.gz → 0.8.4a1__tar.gz

optimum_rbln-0.8.4a1/.github/version.yaml

@@ -0,0 +1 @@
+rebel_compiler_version: 0.8.4.dev278+ge147357b

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.4a0
+Version: 0.8.4a1
 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.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.8.4a0'
-__version_tuple__ = version_tuple = (0, 8, 4, 'a0')
+__version__ = version = '0.8.4a1'
+__version_tuple__ = version_tuple = (0, 8, 4, 'a1')
 
 __commit_id__ = commit_id = None

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/configuration_utils.py

@@ -248,9 +248,6 @@ class RBLNAutoConfig:
             if key[5:] not in RUNTIME_KEYWORDS and key[5:] not in cls.submodules
         }
 
-        if len(rbln_kwargs) > 0:
-            raise ValueError(f"Cannot set the following arguments: {list(rbln_kwargs.keys())}")
-
         # Process submodule's rbln_config
         for submodule in cls.submodules:
             if submodule not in config_file:
@@ -265,6 +262,16 @@ class RBLNAutoConfig:
 
         config_file.update(rbln_runtime_kwargs)
 
+        rbln_config = cls(**config_file)
+
+        if len(rbln_kwargs) > 0:
+            for key, value in rbln_kwargs.items():
+                if getattr(rbln_config, key) != value:
+                    raise ValueError(
+                        f"Cannot set the following arguments: {list(rbln_kwargs.keys())} "
+                        f"Since the value is already set to {getattr(rbln_config, key)}"
+                    )
+
         if return_unused_kwargs:
             return cls(**config_file), kwargs
         else:

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/diffusers/modeling_diffusers.py

@@ -130,7 +130,7 @@ class RBLNDiffusionMixin:
         cls,
         model_id: str,
         *,
-        export: bool = False,
+        export: bool = None,
         model_save_dir: Optional[PathLike] = None,
         rbln_config: Dict[str, Any] = {},
         lora_ids: Optional[Union[str, List[str]]] = None,
@@ -181,6 +181,20 @@ class RBLNDiffusionMixin:
         """
         rbln_config, kwargs = cls.get_rbln_config_class().initialize_from_kwargs(rbln_config, **kwargs)
 
+        if export is None:
+            export = any(
+                not RBLNModel._is_compiled(
+                    model_id,
+                    token=kwargs.get("token"),
+                    revision=kwargs.get("revision"),
+                    force_download=kwargs.get("force_download", False),
+                    cache_dir=kwargs.get("cache_dir"),
+                    subfolder=submodule_name,
+                    local_files_only=kwargs.get("local_files_only", False),
+                )
+                for submodule_name in cls._submodules
+            )
+
         if export:
             # keep submodules if user passed any of them.
             passed_submodules = {

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/diffusers/pipelines/auto_pipeline.py

@@ -14,7 +14,8 @@
 
 
 import importlib
-from typing import Type
+from pathlib import Path
+from typing import Type, Union
 
 from diffusers.models.controlnets import ControlNetUnionModel
 from diffusers.pipelines.auto_pipeline import (
@@ -42,7 +43,13 @@ class RBLNAutoPipelineBase:
     _model_mapping_names = None
 
     @classmethod
-    def get_rbln_cls(cls, pretrained_model_name_or_path, export=True, **kwargs):
+    def get_rbln_cls(cls, pretrained_model_name_or_path: Union[str, Path], export: bool = None, **kwargs):
+        if isinstance(pretrained_model_name_or_path, Path):
+            pretrained_model_name_or_path = pretrained_model_name_or_path.as_posix()
+
+        if export is None:
+            export = not cls._is_compiled_pipeline(pretrained_model_name_or_path, **kwargs)
+
         if export:
             hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
             rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
@@ -66,7 +73,7 @@ class RBLNAutoPipelineBase:
         return rbln_cls
 
     @classmethod
-    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path, **kwargs):
+    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path: Union[str, Path], **kwargs):
         """
         Retrieve the path to the compiled model directory for a given RBLN model.
 
@@ -86,10 +93,36 @@ class RBLNAutoPipelineBase:
 
         return model_index_config["_class_name"]
 
+    @classmethod
+    def _is_compiled_pipeline(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+        cache_dir=None,
+        force_download=False,
+        proxies=None,
+        token=None,
+        local_files_only=False,
+        revision=None,
+        **kwargs,
+    ):
+        config: dict = cls.load_config(
+            pretrained_model_name_or_path,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            token=token,
+            local_files_only=local_files_only,
+            revision=revision,
+        )
+        for value in config.values():
+            if isinstance(value, list) and len(value) > 0 and value[0] == "optimum.rbln":
+                return True
+        return False
+
     @classmethod
     def infer_hf_model_class(
         cls,
-        pretrained_model_or_path,
+        pretrained_model_or_path: Union[str, Path],
         cache_dir=None,
         force_download=False,
         proxies=None,

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/modeling_base.py

@@ -343,11 +343,37 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         rbln_config, kwargs = config_cls.initialize_from_kwargs(rbln_config, **kwargs)
         return rbln_config, kwargs
 
+    @classmethod
+    def _is_compiled(
+        cls,
+        model_id: Union[str, Path],
+        token: Optional[Union[bool, str]] = None,
+        revision: Optional[str] = None,
+        force_download: bool = False,
+        cache_dir: Optional[str] = None,
+        subfolder: str = "",
+        local_files_only: bool = False,
+    ) -> bool:
+        # Check if the model is already compiled.
+        try:
+            cls._load_compiled_model_dir(
+                model_id=model_id,
+                token=token,
+                revision=revision,
+                force_download=force_download,
+                cache_dir=cache_dir,
+                subfolder=subfolder,
+                local_files_only=local_files_only,
+            )
+            return True
+        except (FileNotFoundError, KeyError):
+            return False
+
     @classmethod
     def from_pretrained(
         cls: Type["RBLNBaseModel"],
         model_id: Union[str, Path],
-        export: bool = False,
+        export: bool = None,
         rbln_config: Optional[Union[Dict, RBLNModelConfig]] = None,
         **kwargs: Any,
     ) -> "RBLNBaseModel":
@@ -357,7 +383,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
         Args:
             model_id: The model id of the pre-trained model to be loaded. It can be downloaded from the HuggingFace model hub or a local path, or a model id of a compiled model using the RBLN Compiler.
-            export: A boolean flag to indicate whether the model should be compiled.
+            export: A boolean flag to indicate whether the model should be compiled. If None, it will be determined based on the existence of the compiled model files in the model_id.
             rbln_config: Configuration for RBLN model compilation and runtime. This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
                 For detailed configuration options, see the specific model's configuration class documentation.
 
@@ -369,6 +395,18 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
         if isinstance(model_id, Path):
             model_id = model_id.as_posix()
+
+        if export is None:
+            export = not cls._is_compiled(
+                model_id=model_id,
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                force_download=kwargs.get("force_download", False),
+                cache_dir=kwargs.get("cache_dir"),
+                subfolder=kwargs.get("subfolder", ""),
+                local_files_only=kwargs.get("local_files_only", False),
+            )
+
         from_pretrained_method = cls._export if export else cls._from_pretrained
         return from_pretrained_method(model_id=model_id, **kwargs, rbln_config=rbln_config)
 

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/transformers/models/auto/auto_factory.py

@@ -14,9 +14,10 @@
 import importlib
 import inspect
 import warnings
-from typing import Type
+from pathlib import Path
+from typing import Any, Type, Union
 
-from transformers import AutoConfig, PretrainedConfig
+from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
 from transformers.models.auto.auto_factory import _get_model_class
 
@@ -43,10 +44,10 @@ class _BaseAutoModelClass:
     @classmethod
     def get_rbln_cls(
         cls,
-        pretrained_model_name_or_path,
-        *args,
-        export=True,
-        **kwargs,
+        pretrained_model_name_or_path: Union[str, Path],
+        *args: Any,
+        export: bool = None,
+        **kwargs: Any,
     ):
         """
         Determine the appropriate RBLN model class based on the given model ID and configuration.
@@ -59,6 +60,20 @@ class _BaseAutoModelClass:
         Returns:
             RBLNBaseModel: The corresponding RBLN model class.
         """
+        if isinstance(pretrained_model_name_or_path, Path):
+            pretrained_model_name_or_path = pretrained_model_name_or_path.as_posix()
+
+        if export is None:
+            export = not RBLNBaseModel._is_compiled(
+                model_id=pretrained_model_name_or_path,
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                force_download=kwargs.get("force_download", False),
+                cache_dir=kwargs.get("cache_dir"),
+                subfolder=kwargs.get("subfolder", ""),
+                local_files_only=kwargs.get("local_files_only", False),
+            )
+
         if export:
             hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
             rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
@@ -85,9 +100,9 @@ class _BaseAutoModelClass:
     @classmethod
     def infer_hf_model_class(
         cls,
-        pretrained_model_name_or_path,
-        *args,
-        **kwargs,
+        pretrained_model_name_or_path: Union[str, Path],
+        *args: Any,
+        **kwargs: Any,
     ):
         """
         Infer the HuggingFace model class based on the configuration or model name.
@@ -140,7 +155,7 @@ class _BaseAutoModelClass:
         return model_class
 
     @classmethod
-    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path, **kwargs):
+    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path: Union[str, Path], **kwargs):
         """
         Retrieve the path to the compiled model directory for a given RBLN model.
 
@@ -163,17 +178,17 @@ class _BaseAutoModelClass:
         return rbln_config.rbln_model_cls_name
 
     @classmethod
-    def from_pretrained(cls, model_id, *args, **kwargs):
+    def from_pretrained(cls, model_id: Union[str, Path], *args, **kwargs):
         rbln_cls = cls.get_rbln_cls(model_id, *args, **kwargs)
         return rbln_cls.from_pretrained(model_id, *args, **kwargs)
 
     @classmethod
-    def from_model(cls, model, *args, **kwargs):
+    def from_model(cls, model: PreTrainedModel, *args, **kwargs):
         rbln_cls = get_rbln_model_cls(f"RBLN{model.__class__.__name__}")
         return rbln_cls.from_model(model, *args, **kwargs)
 
     @staticmethod
-    def register(rbln_cls: Type[RBLNBaseModel], exist_ok=False):
+    def register(rbln_cls: Type[RBLNBaseModel], exist_ok: bool = False):
         """
         Register a new RBLN model class.
 

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/transformers/models/grounding_dino/grounding_dino_architecture.py

@@ -357,10 +357,16 @@ class _GroundingDinoMultiscaleDeformableAttention(torch.nn.Module):
         batch_size, num_queries, _ = hidden_states.shape
         batch_size, sequence_length, _ = encoder_hidden_states.shape
         # Ignore copy
-        if (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() != sequence_length:
-            raise ValueError(
-                "Make sure to align the spatial shapes with the sequence length of the encoder hidden states"
+        if torch.compiler.is_exporting():
+            torch._check(
+                (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum().item() == sequence_length,
+                "Make sure to align the spatial shapes with the sequence length of the encoder hidden states",
             )
+        else:
+            if (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() != sequence_length:
+                raise ValueError(
+                    "Make sure to align the spatial shapes with the sequence length of the encoder hidden states"
+                )
 
         value = self.value_proj(encoder_hidden_states)
         if attention_mask is not None:

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/src/optimum/rbln/transformers/models/time_series_transformer/time_series_transformers_architecture.py

@@ -162,7 +162,13 @@ class TimeSeriesTransformersDecoder(nn.Module):
         attention_mask = _prepare_4d_causal_attention_mask(attention_mask, input_shape, inputs_embeds, cache_position)
 
         hidden_states = self.value_embedding(inputs_embeds)
-        embed_pos = self.embed_positions.weight[cache_position + self.config.context_length]
+        embed_idx = cache_position + self.config.context_length
+        if torch.compiler.is_exporting():
+            embed_idx = embed_idx.item()
+            torch._check_is_size(embed_idx)
+            torch._check(embed_idx >= 0)
+            torch._check(embed_idx < len(self.embed_positions.weight))
+        embed_pos = self.embed_positions.weight[embed_idx]
         hidden_states = self.layernorm_embedding(hidden_states + embed_pos)
 
         # iterate decoder_layer

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/tests/test_base.py

@@ -111,7 +111,6 @@ class BaseHubTest:
 
             _ = self.RBLN_CLASS.from_pretrained(
                 f"{HF_USER_ID}/{self.get_hf_remote_dir()}",
-                export=False,
                 **self.HF_CONFIG_KWARGS,
                 rbln_device=self.DEVICE,
                 token=HF_AUTH_TOKEN,
@@ -148,7 +147,6 @@ class BaseTest:
 
             cls.model = cls.RBLN_CLASS.from_pretrained(
                 cls.HF_MODEL_ID,
-                export=True,
                 model_save_dir=cls.get_rbln_local_dir(),
                 rbln_device=cls.DEVICE,
                 **cls.RBLN_CLASS_KWARGS,
@@ -221,7 +219,6 @@ class BaseTest:
                     # Test load
                     _ = self.RBLN_CLASS.from_pretrained(
                         tmpdir,
-                        export=False,
                         **self.HF_CONFIG_KWARGS,
                     )
 
@@ -230,7 +227,6 @@ class BaseTest:
                     self.model.save_pretrained(tmpdir)
                     _ = self.RBLN_CLASS.from_pretrained(
                         tmpdir,
-                        export=False,
                         rbln_create_runtimes=False,
                         **self.HF_CONFIG_KWARGS,
                     )
@@ -244,7 +240,6 @@ class BaseTest:
                 # Test model_save_dir
                 _ = self.RBLN_CLASS.from_pretrained(
                     self.get_rbln_local_dir(),
-                    export=False,
                     rbln_create_runtimes=False,
                     **self.HF_CONFIG_KWARGS,
                 )
@@ -305,7 +300,6 @@ class DisallowedTestBase:
             try:
                 _ = self.RBLN_CLASS.from_pretrained(
                     self.HF_MODEL_ID,
-                    export=True,
                     model_save_dir=self.get_rbln_local_dir(),
                     **self.RBLN_CLASS_KWARGS,
                     **self.HF_CONFIG_KWARGS,

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/tests/test_config.py

@@ -60,7 +60,7 @@ def test_stable_diffusion_config(stable_diffusion_model):
 def test_explicit_config_parameters(model_id):
     """Test loading model with explicit configuration parameters."""
     model = RBLNResNetForImageClassification.from_pretrained(
-        model_id, export=True, rbln_image_size=224, rbln_batch_size=2, rbln_create_runtimes=False
+        model_id, rbln_image_size=224, rbln_batch_size=2, rbln_create_runtimes=False
     )
     assert model is not None
     assert hasattr(model, "rbln_config")
@@ -71,7 +71,7 @@ def test_config_dict(model_id):
     """Test loading model with configuration passed as a dictionary."""
     rbln_config = {"create_runtimes": False, "optimize_host_memory": True, "image_size": 64}
 
-    model = RBLNResNetForImageClassification.from_pretrained(model_id, export=True, rbln_config=rbln_config)
+    model = RBLNResNetForImageClassification.from_pretrained(model_id, rbln_config=rbln_config)
     assert model is not None
     assert hasattr(model, "rbln_config")
     assert model.rbln_config.image_size == 64
@@ -88,7 +88,7 @@ def test_config_object(model_id):
     compile_cfg = RBLNCompileConfig(input_info=[("pixel_values", (1, 3, 224, 224), "float32")])
     config.set_compile_cfgs([compile_cfg])
 
-    model = RBLNResNetForImageClassification.from_pretrained(model_id, export=True, rbln_config=config)
+    model = RBLNResNetForImageClassification.from_pretrained(model_id, rbln_config=config)
     assert model is not None
     assert hasattr(model, "rbln_config")
     # Pre-configured object should be properly applied
@@ -125,7 +125,7 @@ def test_config_persistence_after_reload(model_id, tmp_path):
 
     # Use distinctive values to ensure we can detect them
     original_model = RBLNResNetForImageClassification.from_pretrained(
-        model_id, export=True, rbln_image_size=112, rbln_batch_size=3, rbln_create_runtimes=False
+        model_id, rbln_image_size=112, rbln_batch_size=3, rbln_create_runtimes=False
     )
     original_model.save_pretrained(save_dir)
 
@@ -173,7 +173,7 @@ def test_config_priority(model_id):
 def test_invalid_config_parameters(model_id, invalid_param):
     """Test robust handling of various invalid configuration parameters."""
     with pytest.raises((ValueError, TypeError)):
-        _ = RBLNResNetForImageClassification.from_pretrained(model_id, export=True, **invalid_param)
+        _ = RBLNResNetForImageClassification.from_pretrained(model_id, **invalid_param)
 
 
 def test_custom_class(model_id):
@@ -203,7 +203,7 @@ def test_custom_class(model_id):
 
     RBLNAutoModel.register(RBLNResNetModel)
     RBLNAutoConfig.register(RBLNResNetModelConfig)
-    my_model = RBLNResNetModel.from_pretrained(model_id, export=True, rbln_device=-1)
+    my_model = RBLNResNetModel.from_pretrained(model_id, rbln_device=-1)
     random_image_input = torch.randn(1, 3, 64, 64)
     _ = my_model(random_image_input)
 

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/tests/test_diffusers.py

@@ -281,11 +281,7 @@ class TestKandinskyV22Model(BaseTest.TestModel):
         rbln_class_kwargs_copy = self.RBLN_CLASS_KWARGS.copy()
         rbln_class_kwargs_copy["rbln_config"] = rbln_config
         with self.subTest():
-            _ = self.RBLN_CLASS.from_pretrained(
-                model_id=self.HF_MODEL_ID,
-                export=True,
-                **rbln_class_kwargs_copy,
-            )
+            _ = self.RBLN_CLASS.from_pretrained(model_id=self.HF_MODEL_ID, **rbln_class_kwargs_copy)
         with self.subTest():
             self.assertEqual(_.prior_text_encoder.rbln_config.batch_size, 2)
             self.assertEqual(_.prior_prior.rbln_config.batch_size, 4)

{optimum_rbln-0.8.4a0 → optimum_rbln-0.8.4a1}/tests/test_llm.py

@@ -674,6 +674,7 @@ class TestLlamaForCausalLM_fp8(LLMTest.TestLLM):
             "tensor_parallel_size": 1,
         },
     }
+    TEST_LEVEL = TestLevel.DISABLED
 
     def test_generate(self):
         # Cannot generate output with fp8 quantization in ATOM™

optimum_rbln-0.8.4a0/.github/version.yaml

@@ -1 +0,0 @@
-rebel_compiler_version: 0.8.4.dev42+g3cd3a7b5