optimum-rbln 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py

@@ -22,12 +22,12 @@
 # from Rebellions Inc.
 
 import logging
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import torch
-from transformers import PretrainedConfig, PreTrainedModel, XLMRobertaConfig, XLMRobertaModel
+from transformers import PretrainedConfig
 
-from ....modeling_base import RBLNModel
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
 
 
@@ -38,38 +38,6 @@ if TYPE_CHECKING:
 
 
 class RBLNXLMRobertaModel(RBLNModel):
-    original_model_class = XLMRobertaModel
-    original_config_class = XLMRobertaConfig
-
-    @classmethod
-    def get_pytorch_model(
-        cls,
-        model_id: str,
-        use_auth_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,
-        trust_remote_code: bool = False,
-        rbln_kwargs: Optional[Dict[str, Any]] = None,
-        **kwargs,
-    ) -> "PreTrainedModel":
-        model: "PreTrainedModel" = super().get_pytorch_model(
-            model_id=model_id,
-            use_auth_token=use_auth_token,
-            revision=revision,
-            force_download=force_download,
-            cache_dir=cache_dir,
-            subfolder=subfolder,
-            local_files_only=local_files_only,
-            trust_remote_code=trust_remote_code,
-            rbln_kwargs=rbln_kwargs,
-            library_name="transformers",
-        )
-
-        return model
-
     @classmethod
     def _get_rbln_config(
         cls,

optimum/rbln/transformers/utils/rbln_quantization.py

@@ -21,13 +21,95 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-
-from typing import Any
+import functools
+import glob
+import os
+from typing import Any, Callable, Dict, Optional
 
 import torch
+from safetensors.torch import load_file
 from torch.nn import Linear, Parameter
 from torch.nn import functional as F
 
+from ...utils.logging import get_logger
+
+
+logger = get_logger()
+
+SUPPORTED_QUANTIZATIONS: Dict[str, list[str]] = {
+    "rbln": ["w4a16"],
+}
+
+
+class QuantizationManager:
+    # The RBLN_QUANT_BITS environment variable defines the precision of each layer during the graph compilation process.
+    # It specifies the quantization bit depth. For instance, setting RBLN_QUANT_BITS=4 will apply 4-bit precision for quantization.
+    RBLN_QUANT_BITS_ENV = "RBLN_QUANT_BITS"
+
+    @staticmethod
+    def _raise_invalid_config_error(
+        key: str, value: str, valid_values: list[str], context: Optional[str] = None
+    ) -> None:
+        context_info = f" for {context}" if context else ""
+        valid_values_str = ", ".join(valid_values)
+        raise ValueError(f"Invalid {key}: {value}{context_info}. " f"Supported values are: {valid_values_str}")
+
+    @staticmethod
+    def validate_quantization_config(quantize_config: Optional[dict]) -> Optional[dict]:
+        if not quantize_config:
+            return None
+
+        q_format = quantize_config.get("format")
+        q_precision = quantize_config.get("precision")
+
+        if q_format not in SUPPORTED_QUANTIZATIONS:
+            QuantizationManager._raise_invalid_config_error(
+                "quantization format", q_format, list(SUPPORTED_QUANTIZATIONS.keys())
+            )
+
+        if q_precision not in SUPPORTED_QUANTIZATIONS[q_format]:
+            QuantizationManager._raise_invalid_config_error(
+                "precision", q_precision, SUPPORTED_QUANTIZATIONS[q_format], q_format
+            )
+
+        return quantize_config
+
+    @classmethod
+    def _set_env_var(cls, name: str, value: str) -> None:
+        os.environ[name] = value
+
+    @classmethod
+    def _unset_env_var(cls, name: str) -> None:
+        os.environ.pop(name, None)
+
+    @classmethod
+    def set_quantization_env(cls, quantize_config: Optional[dict]) -> Optional[str]:
+        quantize_config = cls.validate_quantization_config(quantize_config)
+        if quantize_config:
+            q_precision: str = quantize_config["precision"]
+            quant_bits = q_precision.split("w")[1].split("a")[0]
+            cls._set_env_var(cls.RBLN_QUANT_BITS_ENV, quant_bits)
+            return cls.RBLN_QUANT_BITS_ENV
+        return None
+
+    @classmethod
+    def reset_quantization_env(cls, env_var_name: Optional[str]) -> None:
+        if env_var_name:
+            cls._unset_env_var(env_var_name)
+
+    @classmethod
+    def with_quantization_env(cls, func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            quantize_config = kwargs.get("quantize_config")
+            quantize_env_var = cls.set_quantization_env(quantize_config)
+            try:
+                return func(*args, **kwargs)
+            finally:
+                cls.reset_quantization_env(quantize_env_var)
+
+        return wrapper
+
 
 # Constants
 QUANTIZED_WEIGHTS = {
@@ -41,7 +123,15 @@ QUANTIZED_WEIGHTS = {
 }
 
 
-def update_layers_to_quantized(module: torch.nn.Module) -> None:
+def prepare_model_for_quantization(model: torch.nn.Module, model_id: str, n_layer: Optional[int] = None) -> None:
+    """
+    Prepare the model for quantization by updating specified linear layers to quantized (qlinear) layers.
+    """
+    update_layers_to_quantize(model)
+    load_weights(model, model_id, n_layer)
+
+
+def update_layers_to_quantize(module: torch.nn.Module) -> None:
     """
     Updates specified linear layers to quantized (qlinear) layers in the given module.
     """
@@ -54,7 +144,33 @@ def update_layers_to_quantized(module: torch.nn.Module) -> None:
             processed_layers.append(name)
 
     if processed_layers:
-        print(f"Updated the following linear layers to quantized layers:\n {{{', '.join(processed_layers)}}}")
+        logger.debug(f"Updated the following linear layers to quantized layers:\n {{{', '.join(processed_layers)}}}")
+
+
+def load_weights(model, model_id, n_layer=None):
+    """
+    Load safetensor file data directly into the model, filtering by layer if n_layer is provided.
+    """
+
+    model_params = dict(model.named_parameters(recurse=True))
+    model_buffers = dict(model.named_buffers(recurse=True))
+    safetensor_files = glob.glob(f"{model_id}/*.safetensors")
+
+    target_layers = list(range(n_layer)) if n_layer is not None else None
+
+    for safetensor_file in safetensor_files:
+        file_data = load_file(safetensor_file)
+        for key, value in file_data.items():
+            if target_layers is not None:
+                parts = key.split(".")
+
+                if len(parts) > 2 and parts[2].isdigit() and (int(parts[2]) not in target_layers):
+                    continue
+
+            if key in model_params:
+                model_params[key].data.copy_(value)
+            elif key in model_buffers:
+                model_buffers[key].data.copy_(value)
 
 
 def is_target_for_qlinear_replacement(layer_name: str, layer: torch.nn.Module) -> bool:

optimum/rbln/utils/decorator_utils.py

@@ -1,3 +1,27 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+import inspect
 from functools import wraps
 
 from .logging import get_logger
@@ -12,8 +36,8 @@ def remove_compile_time_kwargs(func):
 
     For RBLN-optimized pipelines, several parameters must be determined during compilation
     and cannot be modified during inference. This decorator:
-    1. Removes and warns about LoRA scale in cross_attention_kwargs
-    2. Removes and warns about image dimension parameters (height, width)
+    1. Removes and warns about image dimension parameters (height, width)
+    2. Removes and warns about LoRA scale in cross_attention_kwargs
 
     Args:
         func: The pipeline's __call__ method to be wrapped
@@ -21,15 +45,31 @@ def remove_compile_time_kwargs(func):
 
     @wraps(func)
     def wrapper(self, *args, **kwargs):
-        height_exists = "height" in kwargs and kwargs["height"] is not None
-        width_exists = "width" in kwargs and kwargs["width"] is not None
-        if height_exists or width_exists:
-            logger.warning(
-                "Image dimension parameters (`height`, `width`) will be ignored during inference. "
-                "Image dimensions must be specified during model compilation using from_pretrained()."
-            )
-            kwargs.pop("width", None)
-            kwargs.pop("height", None)
+        check_params = {"height", "width"}
+        params = inspect.signature(self.original_class.__call__).parameters
+
+        # If height and width exist in the base pipeline's __call__ method arguments
+        # Otherwise, if there is no height or width of kwargs, it is filled based on the compiled size.
+        if check_params.issubset(params):
+            compiled_image_size = self.get_compiled_image_size()
+            if compiled_image_size is not None:
+                height_exists = "height" in kwargs and kwargs["height"] is not None
+                width_exists = "width" in kwargs and kwargs["width"] is not None
+                if height_exists or width_exists:
+                    if not (
+                        kwargs.get("height", None) == compiled_image_size[0]
+                        and kwargs.get("width", None) == compiled_image_size[1]
+                    ):
+                        logger.warning(
+                            "Image dimension parameters (`height`, `width`) will be ignored during inference. "
+                            "Image dimensions (%s, %s) must be specified during model compilation using from_pretrained(), (%s, %s).",
+                            str(kwargs.get("height", None)),
+                            str(kwargs.get("width", None)),
+                            str(compiled_image_size[0]),
+                            str(compiled_image_size[1]),
+                        )
+                kwargs["height"] = compiled_image_size[0]
+                kwargs["width"] = compiled_image_size[1]
 
         if "cross_attention_kwargs" in kwargs:
             cross_attention_kwargs = kwargs.get("cross_attention_kwargs")

optimum/rbln/utils/hub.py

@@ -0,0 +1,131 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+import os
+from pathlib import Path
+from typing import List, Optional, Union
+
+from huggingface_hub import HfApi, HfFolder, hf_hub_download
+
+
+class PushToHubMixin:
+    def push_to_hub(
+        self,
+        save_directory: str,
+        repository_id: str,
+        private: Optional[bool] = None,
+        use_auth_token: Union[bool, str] = True,
+    ) -> str:
+        huggingface_token = _get_huggingface_token(use_auth_token)
+        api = HfApi()
+
+        api.create_repo(
+            token=huggingface_token,
+            repo_id=repository_id,
+            exist_ok=True,
+            private=private,
+        )
+        for path, subdirs, files in os.walk(save_directory):
+            for name in files:
+                local_file_path = os.path.join(path, name)
+                _, hub_file_path = os.path.split(local_file_path)
+                # FIXME: when huggingface_hub fixes the return of upload_file
+                try:
+                    api.upload_file(
+                        token=huggingface_token,
+                        repo_id=f"{repository_id}",
+                        path_or_fileobj=os.path.join(os.getcwd(), local_file_path),
+                        path_in_repo=hub_file_path,
+                    )
+                except KeyError:
+                    pass
+                except NameError:
+                    pass
+
+
+def pull_compiled_model_from_hub(
+    model_id: Union[str, Path],
+    subfolder: str,
+    use_auth_token: Optional[Union[bool, str]],
+    revision: Optional[str],
+    cache_dir: Optional[str],
+    force_download: bool,
+    local_files_only: bool,
+) -> Path:
+    """Pull model files from the Hugging Face Hub."""
+    huggingface_token = _get_huggingface_token(use_auth_token)
+    repo_files = list(
+        map(
+            Path,
+            HfApi().list_repo_files(model_id, revision=revision, token=huggingface_token),
+        )
+    )
+
+    pattern_rbln = "*.rbln" if subfolder == "" else f"{subfolder}/*.rbln"
+    rbln_files = [p for p in repo_files if p.match(pattern_rbln)]
+
+    pattern_config = "rbln_config.json" if subfolder == "" else f"{subfolder}/rbln_config.json"
+    rbln_config_filenames = [p for p in repo_files if p.match(pattern_config)]
+
+    validate_files(rbln_files, rbln_config_filenames, f"repository {model_id}")
+
+    filenames = [str(path) for path in repo_files]
+
+    for filename in filenames:
+        rbln_config_cache_path = hf_hub_download(
+            repo_id=model_id,
+            filename=filename,
+            subfolder=subfolder,
+            use_auth_token=use_auth_token,
+            revision=revision,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            local_files_only=local_files_only,
+        )
+
+    return Path(rbln_config_cache_path).parent
+
+
+def validate_files(
+    files: List[Path],
+    config_files: List[Path],
+    location: str,
+):
+    """Validate the presence and count of required files."""
+    if len(files) == 0:
+        raise FileNotFoundError(f"Could not find any rbln model file in {location}")
+
+    if len(config_files) == 0:
+        raise FileNotFoundError(f"Could not find `rbln_config.json` file in {location}")
+
+    if len(config_files) > 1:
+        raise FileExistsError(f"Multiple rbln_config.json files found in {location}. This is not expected.")
+
+
+def _get_huggingface_token(use_auth_token: Union[bool, str]) -> str:
+    if isinstance(use_auth_token, str):
+        return use_auth_token
+    elif use_auth_token:
+        return HfFolder.get_token()
+    else:
+        raise ValueError("`use_auth_token` must be provided to interact with the Hugging Face Hub.")

optimum/rbln/utils/import_utils.py

@@ -37,11 +37,32 @@ class VersionCompat:
 
 
 RBLN_VERSION_COMPATS = {
+    "0.2.0": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.7.1",
+            max_version="0.7.2",
+        ),
+    ],
+    "0.1.15": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.6.2",
+            max_version="0.6.3",
+        ),
+    ],
+    "0.1.14": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.6.2",
+            max_version="0.6.3",
+        ),
+    ],
     "0.1.13": [
         VersionCompat(
             package_name="rebel-compiler",
             min_version="0.6.0",
-            max_version="0.6.1",
+            max_version="0.6.2",
         ),
     ],
     "0.1.12": [

optimum/rbln/utils/logging.py

@@ -1,3 +1,40 @@
+# Copyright 2020 Optuna, Hugging Face
+#
+# 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.
+
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
 """
 Logging utilities.
 Modified from `transformers.utils.logging.py`

optimum/rbln/utils/model_utils.py

@@ -0,0 +1,52 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+# Prefix used for RBLN model class names
+RBLN_PREFIX = "RBLN"
+
+
+def convert_hf_to_rbln_model_name(hf_model_name: str):
+    """
+    Convert Hugging Face model name to RBLN model name.
+
+    Args:
+        hf_model_name (str): The Hugging Face model name.
+
+    Returns:
+        str: The corresponding RBLN model name.
+    """
+    return RBLN_PREFIX + hf_model_name
+
+
+def convert_rbln_to_hf_model_name(rbln_model_name: str):
+    """
+    Convert RBLN model name to Hugging Face model name.
+
+    Args:
+        rbln_model_name (str): The RBLN model name.
+
+    Returns:
+        str: The corresponding Hugging Face model name.
+    """
+
+    return rbln_model_name.removeprefix(RBLN_PREFIX)

optimum/rbln/utils/runtime_utils.py

@@ -35,15 +35,15 @@ class RBLNPytorchRuntime:
         self.runtime = runtime
         for key, value in kwargs.items():
             setattr(self, key, value)
-        for mandatory_member in __class__.mandatory_members:
+        for mandatory_member in self.mandatory_members:
             if mandatory_member not in kwargs:
-                raise AttributeError(f"`{mandatory_member}` should be assigned to {__class__.__name__} objects.")
+                raise AttributeError(f"`{mandatory_member}` should be assigned to {self.__class__.__name__} objects.")
 
     def __call__(self, *args: Any, **kwds: Any) -> Any:
         return self.forward(*args, **kwds)
 
     def forward(self, *args: List["torch.Tensor"], **kwargs: Dict[str, "torch.Tensor"]):
-        # filtering uselss args or kwarg such as None.
+        # filtering useless args or kwarg such as None.
         args = list(filter(lambda arg: isinstance(arg, torch.Tensor), args))
         kwargs = dict(filter(lambda kwarg: isinstance(kwarg[1], torch.Tensor) or kwarg[0] == "out", kwargs.items()))
         output = self.runtime(*args, **kwargs)
@@ -76,17 +76,21 @@ class UnavailableRuntime:
 class ContextRblnConfig:
     _local = threading.local()
 
-    def __init__(self, device=None, device_map=None, create_runtimes=None, optimize_host_mem=None):
+    def __init__(
+        self, device=None, device_map=None, create_runtimes=None, optimize_host_mem=None, activate_profiler=None
+    ):
         self.device = device
         self.device_map = device_map
         self.create_runtimes = create_runtimes
         self.optimize_host_mem = optimize_host_mem
+        self.activate_profiler = activate_profiler
 
     def __enter__(self):
         self._local.device = self.device
         self._local.device_map = self.device_map
         self._local.create_runtimes = self.create_runtimes
         self._local.optimize_host_memory = self.optimize_host_mem
+        self._local.activate_profiler = self.activate_profiler
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -94,6 +98,7 @@ class ContextRblnConfig:
         self._local.device_map = None
         self._local.create_runtimes = None
         self._local.optimize_host_memory = None
+        self._local.activate_profiler = None
 
     @classmethod
     def get_current_context(cls):
@@ -102,4 +107,5 @@ class ContextRblnConfig:
             "device_map": getattr(cls._local, "device_map", None),
             "create_runtimes": getattr(cls._local, "create_runtimes", None),
             "optimize_host_memory": getattr(cls._local, "optimize_host_memory", None),
+            "activate_profiler": getattr(cls._local, "activate_profiler", None),
         }

optimum/rbln/utils/save_utils.py

@@ -1,3 +1,16 @@
+# Copyright 2022 The HuggingFace Team. 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.
 # Copyright 2024 Rebellions Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -21,6 +34,10 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
+"""
+Refer to huggingface/optimum/blob/4fdeea77d71e79451ba53e0c1f9d8f37e9704268/optimum/utils/save_utils.py
+"""
+
 import logging
 from pathlib import Path
 from typing import List, Union