optimum-rbln 0.7.3.post1__py3-none-any.whl → 0.7.4__py3-none-any.whl

optimum/rbln/modeling.py

@@ -14,15 +14,16 @@
 
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 import rebel
 import torch
 from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
 from transformers import AutoConfig, PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutput
 
+from .configuration_utils import DEFAULT_COMPILED_MODEL_NAME, RBLNModelConfig
 from .modeling_base import RBLNBaseModel
-from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, use_rbln_config
 from .utils.logging import get_logger
 
 
@@ -48,6 +49,9 @@ class RBLNModel(RBLNBaseModel):
         ```
     """
 
+    output_class = None
+    output_key = "last_hidden_state"
+
     @classmethod
     def update_kwargs(cls, kwargs):
         """
@@ -56,12 +60,7 @@ class RBLNModel(RBLNBaseModel):
         For example, `torchscript`=True should be set because torch.jit
         does not support `transformers` output instances as module output;
         """
-        kwargs.update(
-            {
-                "torchscript": True,
-                "return_dict": False,
-            }
-        )
+        kwargs.update({"torchscript": True})
         return kwargs
 
     @classmethod
@@ -70,7 +69,7 @@ class RBLNModel(RBLNBaseModel):
         model: "PreTrainedModel",
         save_dir_path: Path,
         subfolder: str,
-        rbln_config: RBLNConfig,
+        rbln_config: RBLNModelConfig,
     ):
         """
         If you are unavoidably running on a CPU rather than an RBLN device,
@@ -78,30 +77,29 @@ class RBLNModel(RBLNBaseModel):
         """
 
     @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: RBLNModelConfig) -> torch.nn.Module:
         # Wrap the model if needed.
         return model
 
     @classmethod
-    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
+    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         model = cls.wrap_model_if_needed(model, rbln_config)
         rbln_compile_config = rbln_config.compile_cfgs[0]
         compiled_model = cls.compile(model, rbln_compile_config=rbln_compile_config)
         return compiled_model
 
     @classmethod
-    @use_rbln_config
     def from_model(
         cls,
         model: "PreTrainedModel",
         config: Optional[PretrainedConfig] = None,
-        rbln_config: Dict[str, Any] = {},
+        rbln_config: Optional[RBLNModelConfig] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
         **kwargs,
     ):
         preprocessors = kwargs.pop("preprocessors", [])
-        rbln_kwargs = rbln_config
+        rbln_config, kwargs = cls.prepare_rbln_config(rbln_config=rbln_config, **kwargs)
 
         # Directory to save compile artifacts(.rbln) and original configs
         if model_save_dir is None:
@@ -123,8 +121,15 @@ class RBLNModel(RBLNBaseModel):
                 config = AutoConfig.from_pretrained(config._name_or_path, **kwargs)
 
         if hasattr(model, "can_generate") and model.can_generate():
+            import json
+
             generation_config = model.generation_config
-            generation_config.save_pretrained(save_dir_path / subfolder)
+            generation_config_path = save_dir_path / subfolder / "generation_config.json"
+
+            generation_config.save_pretrained(generation_config_path.parent)
+            local_config = json.loads(generation_config_path.read_text(encoding="utf-8"))
+            local_config["transformers_version"] = generation_config.transformers_version
+            generation_config_path.write_text(json.dumps(local_config, indent=2) + "\n", encoding="utf-8")
 
         if not isinstance(config, PretrainedConfig):  # diffusers config
             config = PretrainedConfig(**config)
@@ -134,14 +139,21 @@ class RBLNModel(RBLNBaseModel):
         for preprocessor in preprocessors:
             preprocessor.save_pretrained(save_dir_path / subfolder)
 
-        # ad-hoc
-        rbln_kwargs["n_model_params"] = sum(p.numel() for p in model.parameters())
+        # Load submodules
+        if len(cls._rbln_submodules) > 0:
+            rbln_submodules = cls._load_submodules(
+                model=model,
+                model_save_dir=save_dir,
+                rbln_config=rbln_config,
+                **kwargs,
+            )
+        else:
+            rbln_submodules = []
 
         # Get compilation arguments (e.g. input_info)
-        rbln_config: RBLNConfig = cls.get_rbln_config(
-            preprocessors=preprocessors, model_config=config, rbln_kwargs=rbln_kwargs
+        rbln_config: RBLNModelConfig = cls.update_rbln_config(
+            preprocessors=preprocessors, model=model, model_config=config, rbln_config=rbln_config
         )
-        # rbln_config.update_runtime_cfg(rbln_kwargs) # This is done in get_rbln_config
 
         compiled_model: Union[rebel.RBLNCompiledModel, Dict[str, rebel.RBLNCompiledModel]] = cls.get_compiled_model(
             model, rbln_config=rbln_config
@@ -160,17 +172,6 @@ class RBLNModel(RBLNBaseModel):
         # Save torch artifacts (e.g. embedding matrix if needed.)
         cls.save_torch_artifacts(model, save_dir_path=save_dir_path, subfolder=subfolder, rbln_config=rbln_config)
 
-        # Load submodules
-        if len(cls._rbln_submodules) > 0:
-            rbln_submodules = cls._load_submodules(
-                model=model,
-                model_save_dir=save_dir,
-                rbln_kwargs=rbln_kwargs,
-                **kwargs,
-            )
-        else:
-            rbln_submodules = []
-
         # Instantiate
         return cls._from_pretrained(
             model_id=save_dir_path,
@@ -194,8 +195,8 @@ class RBLNModel(RBLNBaseModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
-        # Some rbln-kwargs should be applied before loading torch module (i.e. quantized llm)
-        rbln_kwargs: Optional[Dict[str, Any]] = None,
+        # Some rbln-config should be applied before loading torch module (i.e. quantized llm)
+        rbln_config: Optional[RBLNModelConfig] = None,
         **kwargs,
     ) -> "PreTrainedModel":
         kwargs = cls.update_kwargs(kwargs)
@@ -215,18 +216,43 @@ class RBLNModel(RBLNBaseModel):
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNModelConfig,
     ) -> List[rebel.Runtime]:
-        if DEFAULT_COMPILED_MODEL_NAME not in rbln_device_map:
+        if DEFAULT_COMPILED_MODEL_NAME not in rbln_config.device_map:
             cls._raise_missing_compiled_file_error([DEFAULT_COMPILED_MODEL_NAME])
 
-        device = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
         return [
-            compiled_model.create_runtime(tensor_type="pt", device=device, activate_profiler=activate_profiler)
+            rebel.Runtime(
+                compiled_model,
+                tensor_type="pt",
+                device=rbln_config.device_map[DEFAULT_COMPILED_MODEL_NAME],
+                activate_profiler=rbln_config.activate_profiler,
+            )
             for compiled_model in compiled_models
         ]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
+    def forward(self, *args, return_dict: Optional[bool] = None, **kwargs):
+        if self.hf_library_name == "transformers":
+            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        else:
+            return_dict = True if return_dict is None else return_dict
+
+        # Get output from the model
         output = self.model[0](*args, **kwargs)
-        return output
+
+        # Format output according to task requirements
+        return self._prepare_output(output, return_dict)
+
+    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:
+            if self.output_class is None:
+                return BaseModelOutput(last_hidden_state=output)
+
+            # Create output with the appropriate class and key
+            return self.output_class(**{self.output_key: output})

optimum/rbln/modeling_base.py

@@ -18,18 +18,13 @@ import shutil
 from abc import ABC, abstractmethod
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 
 import rebel
 import torch
-from transformers import (
-    AutoConfig,
-    AutoModel,
-    GenerationConfig,
-    PretrainedConfig,
-)
-
-from .modeling_config import RBLNCompileConfig, RBLNConfig, use_rbln_config
+from transformers import AutoConfig, AutoModel, GenerationConfig, PretrainedConfig
+
+from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig
 from .utils.hub import PushToHubMixin, pull_compiled_model_from_hub, validate_files
 from .utils.logging import get_logger
 from .utils.runtime_utils import UnavailableRuntime
@@ -47,6 +42,10 @@ class PreTrainedModel(ABC):  # noqa: F811
     pass
 
 
+class RBLNBaseModelConfig(RBLNModelConfig):
+    pass
+
+
 class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     """
     An abstract base class for compiling, loading, and saving neural network models from the huggingface
@@ -85,15 +84,17 @@ 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,
         models: List[rebel.Runtime],
         config: "PretrainedConfig",
-        rbln_config: RBLNConfig,
+        rbln_config: RBLNModelConfig,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
         rbln_compiled_models: Optional[rebel.RBLNCompiledModel] = None,
@@ -103,6 +104,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         self.model = models
         self.config = config
         self.rbln_config = rbln_config
+        if not rbln_config.is_frozen():
+            raise RuntimeError("`rbln_config` must be frozen. Please call `rbln_config.freeze()` first.")
+
         self.compiled_models = rbln_compiled_models
 
         # Registers the RBLN classes into the transformers AutoModel classes to avoid warnings when creating
@@ -118,7 +122,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         else:
             self.generation_config = None
 
-        # self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None
         if self.generation_config is not None:
             self.generation_config.use_cache = True
 
@@ -181,11 +184,10 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return rbln_compiled_models
 
     @classmethod
-    @use_rbln_config
     def _from_pretrained(
         cls,
         model_id: Union[str, Path],
-        config: "PretrainedConfig" = None,
+        config: Optional["PretrainedConfig"] = None,
         use_auth_token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,
@@ -195,17 +197,12 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         trust_remote_code: bool = False,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         # passed from compile function
-        rbln_config: Optional[RBLNConfig] = None,
+        rbln_config: Optional[RBLNModelConfig] = None,
         rbln_compiled_models: Optional[Dict[str, rebel.RBLNCompiledModel]] = None,
         rbln_submodules: List["RBLNBaseModel"] = [],
         **kwargs,
     ) -> "RBLNBaseModel":
-        from_export_method = isinstance(rbln_config, RBLNConfig) and rbln_compiled_models is not None
-
-        if not from_export_method:
-            # from compiled dir
-            rbln_kwargs = rbln_config or {}
-
+        if rbln_compiled_models is None:
             model_path_subfolder = cls._load_compiled_model_dir(
                 model_id=model_id,
                 use_auth_token=use_auth_token,
@@ -216,16 +213,34 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 local_files_only=local_files_only,
             )
 
-            rbln_config = RBLNConfig.load(model_path_subfolder)
-            rbln_config.update_runtime_cfg(rbln_kwargs)
+            if isinstance(rbln_config, dict):
+                rbln_config_as_kwargs = {f"rbln_{key}": value for key, value in rbln_config.items()}
+                kwargs.update(rbln_config_as_kwargs)
+                rbln_config = None
+            elif isinstance(rbln_config, RBLNModelConfig) and rbln_config.rbln_model_cls_name != cls.__name__:
+                raise ValueError(
+                    f"Cannot use the passed rbln_config. Its model class name ({rbln_config.rbln_model_cls_name}) "
+                    f"does not match the expected model class name ({cls.__name__})."
+                )
+
+            rbln_config, kwargs = RBLNAutoConfig.load(
+                model_path_subfolder, passed_rbln_config=rbln_config, kwargs=kwargs, return_unused_kwargs=True
+            )
 
-            if rbln_config.meta["cls"] != cls.__name__:
+            if rbln_config.rbln_model_cls_name != cls.__name__:
                 raise NameError(
                     f"Cannot load the model. The model was originally compiled using "
-                    f"{rbln_config.meta['cls']}, but you are trying to load it with {cls.__name__}."
+                    f"{rbln_config.rbln_model_cls_name}, but you are trying to load it with {cls.__name__}."
                     "Please use the same model class that was used during compilation."
                 )
 
+            if len(cls._rbln_submodules) > 0:
+                rbln_submodules = cls._load_submodules(model_save_dir=model_id, rbln_config=rbln_config, **kwargs)
+            else:
+                rbln_submodules = []
+
+            rbln_config.freeze()
+
             if config is None:
                 if cls.hf_library_name == "transformers":
                     config = AutoConfig.from_pretrained(
@@ -258,15 +273,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
             rbln_compiled_models = cls._load_compiled_models(model_path_subfolder)
 
-            if len(cls._rbln_submodules) > 0:
-                rbln_submodules = cls._load_submodules(
-                    model_save_dir=model_id,
-                    rbln_kwargs=rbln_kwargs,
-                    **kwargs,
-                )
-            else:
-                rbln_submodules = []
-
             if subfolder != "":
                 model_save_dir = Path(model_path_subfolder).absolute().parent
             else:
@@ -286,7 +292,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     def _from_compiled_models(
         cls,
         rbln_compiled_models: Dict[str, rebel.RBLNCompiledModel],
-        rbln_config: RBLNConfig,
+        rbln_config: RBLNModelConfig,
         config: "PretrainedConfig",
         model_save_dir: Union[Path, str],
         subfolder: Union[Path, str],
@@ -303,16 +309,21 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         # create runtimes only if `rbln_create_runtimes` is enabled
         try:
             models = (
-                cls._create_runtimes(rbln_compiled_models, rbln_config.device_map, rbln_config.activate_profiler)
+                cls._create_runtimes(rbln_compiled_models, rbln_config)
                 if rbln_config.create_runtimes
                 else UnavailableRuntime()
             )
 
         except rebel.core.exception.RBLNRuntimeError as e:
-            logger.warning(
-                f"Failed to create the runtime for the model due to a runtime error: {e.__class__.__name__} - {e}"
+            error_msg = (
+                f"\nFailed to create RBLN runtime: {str(e)}\n\n"
+                f"If you only need to compile the model without loading it to NPU, you can use:\n"
+                f"  from_pretrained(..., rbln_create_runtimes=False) or\n"
+                f"  from_pretrained(..., rbln_config={{..., 'create_runtimes': False}})\n\n"
+                f"To check your NPU status, run the 'rbln-stat' command in your terminal.\n"
+                f"Make sure your NPU is properly installed and operational."
             )
-            models = UnavailableRuntime()
+            raise rebel.core.exception.RBLNRuntimeError(error_msg) from e
 
         return cls(
             models,
@@ -326,38 +337,31 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         )
 
     @classmethod
-    @use_rbln_config
-    def _export(
-        cls,
-        model_id: Union[str, Path],
-        rbln_config: Optional[Dict[str, Any]] = None,
-        **kwargs,
-    ) -> "RBLNBaseModel":
+    def _export(cls, model_id: Union[str, Path], **kwargs) -> "RBLNBaseModel":
         subfolder = kwargs.get("subfolder", "")
         model_save_dir = kwargs.pop("model_save_dir", None)
 
-        rbln_kwargs = rbln_config
-        model: "PreTrainedModel" = cls.get_pytorch_model(
-            model_id=model_id,
-            rbln_kwargs=rbln_kwargs,
-            **kwargs,
-        )
+        rbln_config, kwargs = cls.prepare_rbln_config(**kwargs)
+
+        model: "PreTrainedModel" = cls.get_pytorch_model(model_id=model_id, rbln_config=rbln_config, **kwargs)
         preprocessors = maybe_load_preprocessors(model_id, subfolder=subfolder)
         return cls.from_model(
-            model,
-            rbln_config=rbln_config,
-            preprocessors=preprocessors,
-            model_save_dir=model_save_dir,
-            **kwargs,
+            model, preprocessors=preprocessors, model_save_dir=model_save_dir, rbln_config=rbln_config, **kwargs
         )
 
     @classmethod
-    def from_pretrained(
-        cls,
-        model_id: Union[str, Path],
-        export: bool = False,
-        **kwargs,
-    ) -> "RBLNBaseModel":
+    def prepare_rbln_config(
+        cls, rbln_config: Optional[Union[Dict[str, Any], RBLNModelConfig]] = None, **kwargs
+    ) -> Tuple[RBLNModelConfig, Dict[str, Any]]:
+        """
+        Extract rbln-config from kwargs and convert it to RBLNModelConfig.
+        """
+        config_cls = cls.get_rbln_config_class()
+        rbln_config, kwargs = config_cls.initialize_from_kwargs(rbln_config, **kwargs)
+        return rbln_config, kwargs
+
+    @classmethod
+    def from_pretrained(cls, model_id: Union[str, Path], export: bool = False, **kwargs) -> "RBLNBaseModel":
         if isinstance(model_id, Path):
             model_id = model_id.as_posix()
         from_pretrained_method = cls._export if export else cls._from_pretrained
@@ -376,29 +380,26 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return compiled_model
 
     @classmethod
-    def get_rbln_config(
-        cls,
-        rbln_kwargs: Dict[str, Any],
-        **others,
-    ) -> RBLNConfig:
-        """
-        Make default rbln-config for the model.
-        kwargs for overriding model's config can be accepted.
-        Note that batch_size should be specified with proper input_info.
-        """
-        rbln_config = cls._get_rbln_config(**others, rbln_kwargs=rbln_kwargs)
+    def update_rbln_config(cls, **others) -> RBLNModelConfig:
+        rbln_config = cls._update_rbln_config(**others)
+        rbln_config.freeze()
+        if rbln_config.rbln_model_cls_name != cls.__name__:
+            raise NameError(
+                f"Cannot get the rbln config. {cls.__name__} is not the same as {rbln_config.rbln_model_cls_name}. "
+                "This is an internal error. Please report it to the developers."
+            )
         return rbln_config
 
     @classmethod
     def get_hf_class(cls):
         """
-        Lazily loads and caches the corresponding Hugging Face model class.
+        Lazily loads and caches the corresponding HuggingFace model class.
         Removes 'RBLN' prefix from the class name to get the original class name
         (e.g., RBLNLlamaForCausalLM -> LlamaForCausalLM) and imports it from
         the transformers/diffusers module.
 
         Returns:
-            type: The original Hugging Face model class
+            type: The original HuggingFace model class
         """
         if cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
@@ -406,12 +407,42 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             cls._hf_class = getattr(library, hf_cls_name, None)
         return cls._hf_class
 
+    @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:
+            rbln_config_class_name = cls.__name__ + "Config"
+            library = importlib.import_module("optimum.rbln")
+            cls._rbln_config_class = getattr(library, rbln_config_class_name, None)
+            if cls._rbln_config_class is None:
+                raise ValueError(
+                    f"RBLN config class {rbln_config_class_name} not found. This is an internal error. "
+                    "Please report it to the developers."
+                )
+        return cls._rbln_config_class
+
     def can_generate(self):
         return False
 
     def to(self, *args, **kwargs):
         return self
 
+    def parameters(self):
+        """
+        Provides a dummy parameter generator for compatibility.
+
+        This method mimics the interface of torch.nn.Module.parameters()
+        specifically for code that uses `next(model.parameters())` to infer
+        the device or dtype. It yields a single dummy tensor on CPU with float32 dtype.
+
+        Warning:
+            This does NOT yield the actual model parameters used by the RBLN runtime.
+            Code relying on iterating through all model parameters will not work as expected.
+        """
+        yield torch.tensor([1.0], dtype=torch.float32, device=torch.device("cpu"))
+
     def __call__(self, *args, **kwargs):
         return self.forward(*args, **kwargs)
 
@@ -448,7 +479,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             save_directory (`Union[str, Path]`):
                 Directory where to save the model file.
             push_to_hub (`bool`, *optional*, defaults to `False`):
-                Whether or not to push your model to the Hugging Face model hub after saving it.
+                Whether or not to push your model to the HuggingFace model hub after saving it.
 
         """
         if os.path.isfile(save_directory):
@@ -481,11 +512,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             # First copy everything to a temporary directory
             shutil.copytree(real_save_dir, tmp_dir)
 
-            # Save configs to the temporary directory
-            self.config.save_pretrained(tmp_dir)
-            if self.generation_config is not None:
-                self.generation_config.save_pretrained(tmp_dir)
-
             # If everything succeeded, atomically replace the target directory
             if os.path.exists(save_directory_path):
                 shutil.rmtree(save_directory_path)
@@ -521,7 +547,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @classmethod
     @abstractmethod
-    def _get_rbln_config(cls, **rbln_config_kwargs) -> RBLNConfig:
+    def _update_rbln_config(cls, **rbln_config_kwargs) -> RBLNModelConfig:
         pass
 
     @classmethod
@@ -529,8 +555,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNModelConfig,
     ) -> List[rebel.Runtime]:
         # compiled_models -> runtimes
         pass
@@ -542,11 +567,11 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @classmethod
     @abstractmethod
-    @use_rbln_config
     def from_model(
         cls,
         model: "PreTrainedModel",
-        rbln_config: Dict[str, Any] = {},
+        config: Optional[PretrainedConfig] = None,
+        rbln_config: Optional[RBLNModelConfig] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
         **kwargs,

optimum/rbln/ops/__init__.py

@@ -12,10 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .attn import (
-    register_rbln_custom_add_softmax_attention,
-    register_rbln_custom_paged_attention,
-    register_rbln_custom_paged_causal_attention,
-)
-from .flash_attn import register_rbln_custom_paged_flash_attention, register_rbln_custom_paged_flash_causal_attention
-from .kv_cache_update import register_rbln_custom_cache_update
+from .attn import *
+from .flash_attn import *
+from .kv_cache_update import *
+from .linear import linear