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

optimum/rbln/transformers/utils/rbln_quantization.py

@@ -22,12 +22,95 @@
 # 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 +124,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 +145,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

@@ -23,13 +23,17 @@ def remove_compile_time_kwargs(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
+        compiled_image_size = self.vae.image_size
         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)
+            if kwargs["height"] == compiled_image_size[0] and kwargs["width"] == compiled_image_size[1]:
+                pass
+            else:
+                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)
 
         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,25 @@ class VersionCompat:
 
 
 RBLN_VERSION_COMPATS = {
+    "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/model_utils.py

@@ -0,0 +1,53 @@
+# 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

@@ -43,7 +43,7 @@ class RBLNPytorchRuntime:
         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)

optimum/rbln/utils/submodule.py

@@ -0,0 +1,114 @@
+import importlib
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List
+
+from ..modeling_config import RBLNConfig
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+    from ..modeling_base import RBLNBaseModel
+
+
+class SubModulesMixin:
+    """
+    _rbln_submodules = [
+        {"name": "vision_tower"},
+        {"name": "language_model"},
+    ]
+    """
+
+    _rbln_submodules: List[Dict[str, Any]] = []
+
+    def __init__(
+        self,
+        *,
+        rbln_submodules: List["RBLNBaseModel"] = [],
+        **kwargs,
+    ) -> None:
+        for submodule_meta, submodule in zip(self._rbln_submodules, rbln_submodules):
+            setattr(self, submodule_meta["name"], submodule)
+
+    @classmethod
+    def _export_submodules_from_model(
+        cls,
+        model: "PreTrainedModel",
+        model_save_dir: str,
+        rbln_kwargs: Dict[str, Any],
+        **kwargs,
+    ) -> List["RBLNBaseModel"]:
+        rbln_submodules = []
+        for submodule in cls._rbln_submodules:
+            submodule_name = submodule["name"]
+            torch_submodule: "PreTrainedModel" = getattr(model, submodule["name"])
+            cls_name = torch_submodule.__class__.__name__
+            submodule_cls: "RBLNBaseModel" = getattr(importlib.import_module("optimum.rbln"), f"RBLN{cls_name}")
+
+            if submodule_name in rbln_kwargs:
+                kwargs["rbln_config"] = rbln_kwargs[submodule_name]
+
+            rbln_submodule = submodule_cls.from_model(
+                model=torch_submodule,
+                subfolder=submodule_name,
+                model_save_dir=model_save_dir,
+                **kwargs,
+            )
+
+            rbln_submodules.append(rbln_submodule)
+
+        return rbln_submodules
+
+    @classmethod
+    def _load_submodules_from_compiled_models(
+        cls,
+        model_save_dir: str,
+        rbln_kwargs: Dict[str, Any],
+        **kwargs,
+    ):
+        rbln_submodules = []
+        for submodule in cls._rbln_submodules:
+            submodule_name = submodule["name"]
+
+            if submodule_name in rbln_kwargs:
+                kwargs["rbln_config"] = rbln_kwargs[submodule_name]
+
+            # Get cls name for call the constructor of the rbln class
+            submodule_rbln_config = RBLNConfig.load(Path(model_save_dir) / submodule_name)
+            submodule_cls_name = submodule_rbln_config.meta["cls"]
+            submodule_cls: "RBLNBaseModel" = getattr(importlib.import_module("optimum.rbln"), submodule_cls_name)
+
+            rbln_submodule = submodule_cls._from_pretrained(
+                model_id=model_save_dir,
+                config=None,
+                subfolder=submodule_name,
+                **kwargs,
+            )
+            rbln_submodules.append(rbln_submodule)
+        return rbln_submodules
+
+    @classmethod
+    def _load_submodules(
+        cls,
+        model_save_dir,
+        rbln_kwargs,
+        model=None,
+        **kwargs,
+    ):
+        # Two ways :
+        # 1. Compile from pytorch object
+        # 2. Load from compiled file
+        if model is not None:
+            return cls._export_submodules_from_model(
+                model=model,
+                model_save_dir=model_save_dir,
+                rbln_kwargs=rbln_kwargs,
+                **kwargs,
+            )
+
+        else:
+            return cls._load_submodules_from_compiled_models(
+                model_save_dir=model_save_dir,
+                rbln_kwargs=rbln_kwargs,
+                **kwargs,
+            )

optimum_rbln-0.1.15.dist-info/METADATA

@@ -0,0 +1,106 @@
+Metadata-Version: 2.3
+Name: optimum-rbln
+Version: 0.1.15
+Summary: Optimum RBLN is the interface between the Hugging Face 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
+Author-email: "Rebellions Inc." <support@rebellions.ai>
+License: Apache
+Keywords: atom,diffusers,inference,rbln,rebel,transformers
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: <3.13,>=3.9
+Requires-Dist: accelerate>=1.0.1
+Requires-Dist: diffusers<=0.31.0
+Requires-Dist: einops>=0.8.0
+Requires-Dist: halo>=0.0.31
+Requires-Dist: packaging>=24.1
+Requires-Dist: torch<=2.5.1
+Requires-Dist: torchaudio<=2.5.1
+Requires-Dist: torchvision<=0.20.1
+Requires-Dist: transformers==4.45.2
+Description-Content-Type: text/markdown
+
+
+# Optimum RBLN
+
+<div align="center">
+
+<img src="assets/rbln_logo.png" width="60%"/>
+
+[![PyPI version](https://badge.fury.io/py/optimum-rbln.svg)](https://badge.fury.io/py/optimum-rbln)
+[![License](https://img.shields.io/github/license/rebellions-sw/optimum-rbln)](https://github.com/rebellions-sw/optimum-rbln/blob/main/LICENSE)
+
+</div>
+
+🤗 Optimum RBLN provides an interface between Hugging Face libraries ([Transformers](https://huggingface.co/docs/transformers), [Diffusers](https://huggingface.co/docs/diffusers/index)) and RBLN Accelerators, including [ATOM](https://rebellions.ai/rebellions-product/rbln-ca25/) and [REBEL](https://rebellions.ai/rebellions-product/rebel/).
+
+This library enables seamless integration between the Hugging Face ecosystem and RBLN's NPU acceleration through a comprehensive toolkit for model loading and inference across single- and multi-Accelerator environments. While we maintain a list of [officially validated models and tasks](https://docs.rbln.ai/software/optimum/optimum_rbln.html), users can easily adapt other models and tasks with minimal modifications.
+
+## Key Features
+
+🚀 **High Performance Inference**
+- Optimized model execution on RBLN NPUs through RBLN SDK compilation
+- Support for both single-NPU and multi-NPU inference
+- Integrated with RBLN Runtime for optimal performance
+
+🔧 **Easy Integration**
+- Seamless compatibility with Huggingface model hub
+- Drop-in replacement for existing Huggingface pipelines
+- Minimal code changes required for NPU acceleration
+
+
+## Documentation
+
+Check out [the documentation of Optimum RBLN](https://docs.rbln.ai/software/optimum/optimum_rbln.html) for more advanced usage.
+
+## Getting Started
+
+### Install from PyPI
+
+To install the latest release of this package:
+
+- Export environment variables to access to RBLN private PyPI.
+  ```bash
+  export REBEL_PYPI_USERNAME=<username>
+  export REBEL_PYPI_PASSWORD=<password>
+  ```
+
+- Install optimum-rbln package:
+  ```bash
+  pip install --index-url https://pypi.rebellions.in/simple optimum-rbln
+  ```
+
+### Install from source
+
+#### Prerequisites
+
+- Install [uv](https://docs.astral.sh/uv/) (refer [this link](https://docs.astral.sh/uv/getting-started/installation/) for detailed commands)
+
+The below command installs optimum-rbln along with its dependencies.
+
+```bash
+git clone https://github.com/rebellions-sw/optimum-rbln.git
+cd optimum-rbln
+./scripts/uv-sync.sh
+```
+
+If you want to install local rebel-compiler as editable mode in uv environment,
+```bash
+uv pip install -e /path/to/rebel_compiler/python
+```
+
+### Need Help?
+
+- Join our [Developer Community](https://discuss.rebellions.ai/)
+- Contact maintainers at [support@rebellions.ai](mailto:support@rebellions.ai)