optimum-rbln 0.1.9__py3-none-any.whl → 0.1.12__py3-none-any.whl

optimum/rbln/modeling_base.py

@@ -21,16 +21,19 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
+import importlib
+import inspect
 import logging
 import os
 import shutil
 from abc import ABC, abstractmethod
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import rebel
 import torch
+import transformers
 from huggingface_hub import HfApi, HfFolder, hf_hub_download
 from optimum.exporters import TasksManager
 from optimum.modeling_base import OptimizedModel
@@ -46,9 +49,9 @@ from transformers import (
     PretrainedConfig,
 )
 
-from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
+from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig, use_rbln_config
 from .utils.runtime_utils import UnavailableRuntime
-from .utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
+from .utils.save_utils import maybe_load_preprocessors
 
 
 if TYPE_CHECKING:
@@ -62,7 +65,111 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class RBLNBaseModel(OptimizedModel, ABC):
+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: "RBLNModel" = 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)
+
+            config = OptimizedModel._load_config(Path(model_save_dir) / submodule_name)
+            rbln_submodule = submodule_cls._from_pretrained(
+                model_id=model_save_dir,
+                config=config,
+                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,
+            )
+
+
+class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
     """
     An abstract base class for compiling, loading, and saving neural network models from the huggingface
     transformers and diffusers libraries to run on RBLN NPU devices.
@@ -110,6 +217,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
         rbln_compiled_models: Optional[rebel.RBLNCompiledModel] = None,
+        rbln_submodules: List["RBLNBaseModel"] = [],
         **kwargs,
     ):
         super().__init__(models, config)
@@ -127,11 +235,18 @@ class RBLNBaseModel(OptimizedModel, ABC):
             self.auto_model_class.register(AutoConfig, self.__class__)
 
         # copied from tranformers PreTrainedModel __init__
-        self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None
+        if self.can_generate():
+            gen_config_dir = model_save_dir.name if isinstance(model_save_dir, TemporaryDirectory) else model_save_dir
+            self.generation_config = GenerationConfig.from_pretrained(gen_config_dir, trust_remote_code=True)
+        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
 
         self.device = torch.device("cpu")
+        self.training = False
 
         # FIXME :: model_save_dir is not used after initialized. (This can be used when save/load)
         # This attribute is needed to keep one reference on the temporary directory, since garbage collecting it
@@ -146,11 +261,9 @@ class RBLNBaseModel(OptimizedModel, ABC):
             self.model_save_dir = model_save_dir
         self.subfolder = subfolder
 
+        self.rbln_submodules = rbln_submodules
         self.__post_init__(**kwargs)
 
-    def __post_init__(self, **kwargs):
-        pass
-
     def _save_pretrained(self, save_directory: Union[str, Path]):
         """
         Saves a model and its configuration file to a directory, so that it can be re-loaded using the
@@ -180,27 +293,18 @@ class RBLNBaseModel(OptimizedModel, ABC):
             )
 
     @classmethod
-    def _from_pretrained(
+    def _load_compiled_model_dir(
         cls,
         model_id: Union[str, Path],
-        config: "PretrainedConfig",
         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,
-        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
-        # Runtime - related kwargs
-        rbln_device: Optional[List[int]] = None,
-        rbln_device_map: Optional[Dict[str, int]] = None,
-        rbln_create_runtimes: Optional[bool] = None,
-        # passed from compile function
-        rbln_config: Optional[RBLNConfig] = None,
-        rbln_compiled_models: Optional[List[rebel.RBLNCompiledModel]] = None,
-        rbln_optimize_host_memory: Optional[bool] = None,
-        **kwargs,
-    ) -> "RBLNBaseModel":
+    ):
+        # Find compiled model
+        # And prepare or download cache folder from HF Hub if needed.
         model_path = Path(model_id)
         if model_path.is_dir():
             model_path = model_path / subfolder
@@ -236,16 +340,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
             )
 
         if model_path.is_dir():
-            if rbln_compiled_models is None:
-                rbln_config = RBLNConfig.load(str(model_path))
-                rbln_compiled_models = [
-                    rebel.RBLNCompiledModel(model_path / f"{compiled_model_name}.rbln")
-                    for compiled_model_name in rbln_config
-                ]
-                new_model_save_dir = model_path
-            else:
-                pass
-
+            model_path = str(model_path)
         else:
             rbln_config_filename = rbln_config_filenames[0]
             rbln_config_cache_path = hf_hub_download(
@@ -258,48 +353,106 @@ class RBLNBaseModel(OptimizedModel, ABC):
                 force_download=force_download,
                 local_files_only=local_files_only,
             )
-            rbln_config = RBLNConfig.load(Path(rbln_config_cache_path).parent)
-            rbln_compiled_models = []
-            for compiled_model_name in rbln_config:
-                model_cache_path = hf_hub_download(
-                    repo_id=model_id,
-                    filename=f"{compiled_model_name}.rbln",
-                    subfolder=subfolder,
-                    use_auth_token=use_auth_token,
-                    revision=revision,
-                    cache_dir=cache_dir,
-                    force_download=force_download,
-                    local_files_only=local_files_only,
+            model_path = Path(rbln_config_cache_path).parent
+
+        return model_path
+
+    @classmethod
+    def _load_compiled_models(cls, model_path: str):
+        compiled_models = Path(model_path).glob("*.rbln")
+        rbln_compiled_models = {cm.stem: rebel.RBLNCompiledModel(cm) for cm in compiled_models}
+        return rbln_compiled_models
+
+    @classmethod
+    @use_rbln_config
+    def _from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        config: "PretrainedConfig",
+        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,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
+        # passed from compile function
+        rbln_config: Optional[RBLNConfig] = 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 {}
+
+            model_path_subfolder = cls._load_compiled_model_dir(
+                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,
+            )
+
+            rbln_config = RBLNConfig.load(model_path_subfolder)
+            rbln_config.update_runtime_cfg(rbln_kwargs)
+
+            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,
                 )
-                rbln_compiled_models.append(rebel.RBLNCompiledModel(model_cache_path))
-            new_model_save_dir = Path(rbln_config_cache_path).parent
+            else:
+                rbln_submodules = []
 
-        preprocessors = maybe_load_preprocessors(model_id, subfolder=subfolder)
+            if subfolder != "":
+                model_save_dir = Path(model_path_subfolder).absolute().parent
+            else:
+                model_save_dir = Path(model_path_subfolder).absolute()
 
-        if model_save_dir is None:
-            model_save_dir = new_model_save_dir
-
-        # Create runtimes
-        if rbln_create_runtimes is None:
-            rbln_create_runtimes = rebel.npu_is_available()
-        if rbln_device_map is None:
-            rbln_device_map = {}
-            device_val = 0 if rbln_device is None else rbln_device
-            for key in rbln_config:
-                rbln_device_map[key] = device_val
-        else:
-            rbln_device_map = rbln_device_map
+        return cls._from_compiled_models(
+            rbln_compiled_models=rbln_compiled_models,
+            rbln_config=rbln_config,
+            config=config,
+            model_save_dir=model_save_dir,
+            subfolder=subfolder,
+            rbln_submodules=rbln_submodules,
+            **kwargs,
+        )
+
+    @classmethod
+    def _from_compiled_models(
+        cls,
+        rbln_compiled_models: Dict[str, rebel.RBLNCompiledModel],
+        rbln_config: RBLNConfig,
+        config: "PretrainedConfig",
+        model_save_dir: Union[Path, str],
+        subfolder: Union[Path, str],
+        rbln_submodules: List["RBLNBaseModel"] = [],
+        **kwargs,
+    ):
+        if isinstance(model_save_dir, str):
+            model_save_dir = Path(model_save_dir)
+        preprocessors = maybe_load_preprocessors(model_save_dir.name, subfolder=subfolder)
+
+        # FIXME:: Should we convert it?
+        compiled_model_names = [cfg.compiled_model_name for cfg in rbln_config.compile_cfgs]
+        rbln_compiled_models = [rbln_compiled_models[cm_name] for cm_name in compiled_model_names]
 
         # create runtimes only if `rbln_create_runtimes` is enabled
         models = (
-            cls._create_runtimes(rbln_compiled_models, rbln_device_map)
-            if rbln_create_runtimes
+            cls._create_runtimes(rbln_compiled_models, rbln_config.device_map)
+            if rbln_config.create_runtimes
             else UnavailableRuntime()
         )
 
-        if rbln_optimize_host_memory is None:
-            rbln_optimize_host_memory = True
-
         return cls(
             models,
             config,
@@ -307,99 +460,50 @@ class RBLNBaseModel(OptimizedModel, ABC):
             preprocessors,
             model_save_dir=model_save_dir,
             subfolder=subfolder,
-            rbln_compiled_models=(None if rbln_optimize_host_memory else rbln_compiled_models),
+            rbln_compiled_models=(None if rbln_config.optimize_host_memory else rbln_compiled_models),
+            rbln_submodules=rbln_submodules,
             **kwargs,
         )
 
     def __repr__(self):
-        return repr(self.model)
+        return repr(self.model) + repr(self.rbln_submodules)
 
     @classmethod
-    def compile(cls, model, rbln_runtime_config: Optional[RBLNRuntimeConfig] = None):
+    def compile(cls, model, rbln_compile_config: Optional[RBLNCompileConfig] = None):
         compiled_model = rebel.compile_from_torch(
             model,
-            input_info=rbln_runtime_config.input_info,
-            batch_size=rbln_runtime_config.batch_size,
-            fusion=rbln_runtime_config.fusion,
-            npu=rbln_runtime_config.npu,
-            tensor_parallel_size=rbln_runtime_config.tensor_parallel_size,
+            input_info=rbln_compile_config.input_info,
+            fusion=rbln_compile_config.fusion,
+            npu=rbln_compile_config.npu,
+            tensor_parallel_size=rbln_compile_config.tensor_parallel_size,
         )
         return compiled_model
 
     @classmethod
     def get_rbln_config(
         cls,
-        **rbln_config_kwargs,
+        rbln_kwargs: Dict[str, Any],
+        **others,
     ) -> RBLNConfig:
         """
         Make default rbln-config for the model.
-
-        if `input_info` specified,
-            other kwargs but `input_info`, `batch_size` and `fusion` are ignored.
-
         kwargs for overriding model's config can be accepted.
-
         Note that batch_size should be specified with proper input_info.
         """
-
-        input_info = rbln_config_kwargs.pop("rbln_input_info", None)
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", None)
-        fusion = rbln_config_kwargs.pop("rbln_fusion", None)
-        npu = rbln_config_kwargs.pop("rbln_npu", None)
-        tensor_parallel_size = rbln_config_kwargs.pop("rbln_tensor_parallel_size", None)
-
-        if input_info is not None:
-            rbln_runtime_config = RBLNRuntimeConfig(
-                input_info=input_info,
-                batch_size=batch_size,
-                fusion=fusion,
-                npu=npu,
-                tensor_parallel_size=tensor_parallel_size,
-            )
-            rbln_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config])
-        else:
-            rbln_config = cls._get_rbln_config(rbln_batch_size=batch_size, **rbln_config_kwargs)
-            for k, rcfgs in rbln_config.items():
-                for rcfg in rcfgs:
-                    rcfg: RBLNRuntimeConfig
-                    rcfg.fusion = fusion
-                    rcfg.npu = npu
-                    rcfg.tensor_parallel_size = tensor_parallel_size
-
+        rbln_config = cls._get_rbln_config(**others, rbln_kwargs=rbln_kwargs)
         return rbln_config
 
-    @staticmethod
-    def pop_rbln_kwargs_from_kwargs(kwargs: dict):
-        keys = list(kwargs.keys())
-        rbln_constructor_kwargs = {
-            key: kwargs.pop(key)
-            for key in keys
-            if key
-            in [
-                "rbln_device",
-                "rbln_device_map",
-                "rbln_create_runtimes",
-                "rbln_optimize_host_memory",
-            ]
-        }
-
-        keys = list(kwargs.keys())
-        rbln_config_kwargs = {key: kwargs.pop(key) for key in keys if key.startswith("rbln_")}
-        return rbln_config_kwargs, rbln_constructor_kwargs
-
     def can_generate(self):
         return False
 
     def to(self, *args, **kwargs):
-        pass
+        return self
 
     def __call__(self, *args, **kwargs):
         return self.forward(*args, **kwargs)
 
-    @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
-        # Wrap the model if needed.
-        return model
+    def __post_init__(self, **kwargs):
+        self.dtype = torch.float32
 
     @classmethod
     def _from_transformers(cls, *args, **kwargs) -> "RBLNBaseModel":
@@ -410,8 +514,14 @@ class RBLNBaseModel(OptimizedModel, ABC):
         return cls._export(*args, **kwargs)
 
     @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+        # Wrap the model if needed.
+        return model
+
+    @classmethod
+    @abstractmethod
     def _get_rbln_config(cls, **rbln_config_kwargs) -> RBLNConfig:
-        raise NotImplementedError
+        pass
 
     @abstractmethod
     def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
@@ -429,25 +539,49 @@ class RBLNBaseModel(OptimizedModel, ABC):
 
     @classmethod
     @abstractmethod
-    def _export(
+    def get_pytorch_model(cls, *args, **kwargs):
+        pass
+
+    @classmethod
+    @abstractmethod
+    @use_rbln_config
+    def from_model(
         cls,
-        model_id: Union[str, Path],
-        config: "PretrainedConfig",
-        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,
+        model: "PreTrainedModel",
+        rbln_config: Dict[str, Any] = {},
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
+        subfolder: str = "",
         **kwargs,
     ):
-        """
-        Exports a vanilla Transformers model into a rbln-compiled Module.
-        """
         pass
 
+    @classmethod
+    @use_rbln_config
+    def _export(
+        cls,
+        model_id: Union[str, Path],
+        config: "PretrainedConfig",  # FIXME : optimum passes config, but we ignore it.
+        rbln_config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> "RBLNModel":
+        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,
+        )
+        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,
+        )
+
 
 class RBLNModel(RBLNBaseModel):
     """
@@ -491,8 +625,8 @@ class RBLNModel(RBLNBaseModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
-        rbln_config_kwargs: Optional[Dict[str, Any]] = None,
-        rbln_constructor_kwargs: Optional[Dict[str, Any]] = None,
+        # Some rbln-kwargs should be applied before loading torch module (i.e. quantized llm)
+        rbln_kwargs: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> "PreTrainedModel":
         task = kwargs.pop("task", None)
@@ -517,36 +651,40 @@ class RBLNModel(RBLNBaseModel):
 
         return model
 
+    @classmethod
+    def save_torch_artifacts(
+        cls,
+        model: "PreTrainedModel",
+        save_dir_path: Path,
+        subfolder: str,
+        rbln_config: RBLNConfig,
+    ):
+        """
+        If you are unavoidably running on a CPU rather than an RBLN device,
+        store the torch tensor, weight, etc. in this function.
+        """
+
     @classmethod
     def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
         model = cls.wrap_model_if_needed(model, rbln_config)
-        rbln_runtime_configs = list(rbln_config.values())
-        if len(rbln_runtime_configs) != 1:
-            raise ValueError
-        rbln_runtime_config = rbln_runtime_configs[0]
-        if len(rbln_runtime_config) != 1:
-            raise ValueError
-        rbln_runtime_config = rbln_runtime_config[0]
-
-        compiled_model = cls.compile(model, rbln_runtime_config=rbln_runtime_config)
+        rbln_compile_config = rbln_config.compile_cfgs[0]
+        compiled_model = cls.compile(model, rbln_compile_config=rbln_compile_config)
         return compiled_model
 
     @classmethod
-    @torch.no_grad()
-    def _export(
+    @use_rbln_config
+    def from_model(
         cls,
-        model_id: str,
-        config: "PretrainedConfig",
-        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,
+        model: "PreTrainedModel",
+        rbln_config: Dict[str, Any] = {},
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
+        subfolder: str = "",
         **kwargs,
-    ) -> "RBLNModel":
+    ):
+        preprocessors = kwargs.pop("preprocessors", [])
+        rbln_kwargs = rbln_config
+
+        # Directory to save compile artifacts(.rbln) and original configs
         if model_save_dir is None:
             save_dir = TemporaryDirectory()
             save_dir_path = Path(save_dir.name)
@@ -558,63 +696,63 @@ class RBLNModel(RBLNBaseModel):
                 save_dir_path = Path(model_save_dir)
                 save_dir_path.mkdir(exist_ok=True)
 
-        rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-
-        model: "PreTrainedModel" = cls.get_pytorch_model(
-            model_id=model_id,
-            subfolder=subfolder,
-            revision=revision,
-            cache_dir=cache_dir,
-            use_auth_token=use_auth_token,
-            local_files_only=local_files_only,
-            force_download=force_download,
-            trust_remote_code=trust_remote_code,
-            rbln_config_kwargs=rbln_config_kwargs,
-            rbln_constructor_kwargs=rbln_constructor_kwargs,
-            **kwargs,
-        )
+        # (Optional) Save preprocessors (tokenizer, image preprocessors, etc)
+        for preprocessor in preprocessors:
+            preprocessor.save_pretrained(save_dir_path)
 
-        # FIXME :: optimum passes AutoConfig.
+        # Save configs
+        # FIXME :: optimum passes AutoConfig. But here we ignore it.
         config = model.config
-
+        if hasattr(model, "can_generate") and model.can_generate():
+            generation_config = model.generation_config
+            generation_config.save_pretrained(save_dir_path / subfolder)
         if not isinstance(config, PretrainedConfig):  # diffusers config
             config = PretrainedConfig(**config)
-
         config.save_pretrained(save_dir_path / subfolder)
-        preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
 
-        # Get compilation arguments
-        if (rbln_config := rbln_config_kwargs.pop("rbln_config", None)) is None:
-            rbln_config = cls.get_rbln_config(preprocessors=preprocessors, model_config=config, **rbln_config_kwargs)
-        compiled_model = cls.get_compiled_model(model, rbln_config=rbln_config)
+        # 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.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
+        )
 
-        # Save compiled models
+        # Save compiled models (.rbln)
         (save_dir_path / subfolder).mkdir(exist_ok=True)
-        if isinstance(compiled_model, Iterable):
-            # compiled_model is an Iterable instance
-            for single_compiled_model, compiled_model_name in zip(compiled_model, rbln_config):
-                single_compiled_model.save(save_dir_path / subfolder / f"{compiled_model_name}.rbln")
+        if not isinstance(compiled_model, dict):
+            compiled_models = {DEFAULT_COMPILED_MODEL_NAME: compiled_model}
+        else:
             compiled_models = compiled_model
+        for compiled_model_name, cm in compiled_models.items():
+            cm.save(save_dir_path / subfolder / f"{compiled_model_name}.rbln")
+        rbln_config.save(save_dir_path / subfolder)
+
+        # 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:
-            compiled_model.save(save_dir_path / subfolder / f"{DEFAULT_COMPILED_MODEL_NAME}.rbln")
-            compiled_models = [compiled_model]
-        rbln_config.save(save_dir_path / subfolder)
+            rbln_submodules = []
 
         # Instantiate
         return cls._from_pretrained(
             model_id=save_dir_path,
             config=config,
             model_save_dir=save_dir,
-            use_auth_token=use_auth_token,
-            revision=revision,
-            force_download=force_download,
-            cache_dir=cache_dir,
             subfolder=subfolder,
-            local_files_only=local_files_only,
             rbln_config=rbln_config,
             rbln_compiled_models=compiled_models,
-            **rbln_constructor_kwargs,
+            rbln_submodules=rbln_submodules,
             **kwargs,
         )
 
@@ -633,18 +771,20 @@ class RBLNModel(RBLNBaseModel):
 
 
 class RBLNModelForQuestionAnswering(RBLNModel):
-    model_type = "rbln_model"
     auto_model_class = AutoModelForQuestionAnswering
+    rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
 
     @classmethod
     def _get_rbln_config(
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_model_input_names: Optional[List[str]] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+
         if rbln_max_seq_len is None:
             for tokenizer in preprocessors:
                 if hasattr(tokenizer, "model_max_length"):
@@ -656,19 +796,34 @@ class RBLNModelForQuestionAnswering(RBLNModel):
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
-        if rbln_model_input_names is not None:
-            cls.rbln_model_input_names = rbln_model_input_names
+        if rbln_model_input_names is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_input_names"):
+                    rbln_model_input_names = tokenizer.model_input_names
+                    break
+            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
+                rbln_model_input_names = cls.rbln_model_input_names
+            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
+                original_model_class = getattr(transformers, model_config.architectures[0])
+                input_names_order = inspect.signature(original_model_class.forward).parameters.keys()
+                raise ValueError(
+                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
+                    f"and be sure to make the order of the inputs same as QuestionAnswering forward() arguments like ({list(input_names_order)})"
+                )
 
         input_info = [
             (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
-            for model_input_name in cls.rbln_model_input_names
+            for model_input_name in rbln_model_input_names
         ]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
-        rbln_runtime_config.batch_size = rbln_batch_size
-        meta = {"rbln_max_seq_len": rbln_max_seq_len}
-
-        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config
 
 
 class RBLNModelForImageClassification(RBLNModel):
@@ -676,7 +831,6 @@ class RBLNModelForImageClassification(RBLNModel):
     This is a generic model class that will be instantiated as one of the model classes of the library (with a image classification head) when created with the from_pretrained() class method
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForImageClassification
 
     @classmethod
@@ -684,33 +838,52 @@ class RBLNModelForImageClassification(RBLNModel):
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_image_size: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
+        rbln_image_size = rbln_kwargs.get("image_size", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
         if rbln_image_size is None:
             for processor in preprocessors:
                 if hasattr(processor, "size"):
-                    rbln_image_size = processor.size["shortest_edge"]
+                    if all(required_key in processor.size.keys() for required_key in ["height", "width"]):
+                        rbln_image_size = (processor.size["height"], processor.size["width"])
+                    elif "shortest_edge" in processor.size.keys():
+                        rbln_image_size = (processor.size["shortest_edge"], processor.size["shortest_edge"])
+                    elif "longest_edge" in processor.size.keys():
+                        rbln_image_size = (processor.size["longest_edge"], processor.size["longest_edge"])
                     break
+
+            if rbln_image_size is None:
+                rbln_image_size = model_config.image_size
+
             if rbln_image_size is None:
                 raise ValueError("`rbln_image_size` should be specified!")
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
+        if isinstance(rbln_image_size, int):
+            rbln_image_height, rbln_image_width = rbln_image_size, rbln_image_size
+        elif isinstance(rbln_image_size, (list, tuple)):
+            rbln_image_height, rbln_image_width = rbln_image_size[0], rbln_image_size[1]
+        elif isinstance(rbln_image_size, dict):
+            rbln_image_height, rbln_image_width = rbln_image_size["height"], rbln_image_size["width"]
+        else:
+            raise ValueError(
+                "`rbln_image_size` should be `int` (ex. 224), `tuple` (ex. 224, 224), `dict` (ex. {'height': 224, 'width': 224}) format"
+            )
+
         input_info = [
             (
                 "pixel_values",
-                [rbln_batch_size, 3, rbln_image_size, rbln_image_size],
+                [rbln_batch_size, 3, rbln_image_height, rbln_image_width],
                 "float32",
             )
         ]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
-        rbln_runtime_config.batch_size = rbln_batch_size
-        meta = {"rbln_image_size": rbln_image_size}
-
-        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        return RBLNConfig(rbln_cls=cls.__name__, compile_cfgs=[rbln_compile_config], rbln_kwargs=rbln_kwargs)
 
 
 class RBLNModelForAudioClassification(RBLNModel):
@@ -726,7 +899,6 @@ class RBLNModelForAudioClassification(RBLNModel):
     Currently, this model class only supports the 'AST' model from the transformers library. Future updates may include support for additional model types.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForAudioClassification
 
     @classmethod
@@ -734,11 +906,11 @@ class RBLNModelForAudioClassification(RBLNModel):
         cls,
         preprocessors: "AutoFeatureExtractor",
         model_config: "PretrainedConfig",
-        rbln_batch_size: Optional[int] = None,
-        rbln_max_length: Optional[int] = None,
-        rbln_num_mel_bins: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {}
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_max_length = rbln_kwargs.get("max_length", None)
+        rbln_num_mel_bins = rbln_kwargs.get("num_mel_bins", None)
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
@@ -764,11 +936,7 @@ class RBLNModelForAudioClassification(RBLNModel):
         if rbln_max_length is None:
             raise ValueError("`rbln_max_length` should be specified!")
 
-        meta["rbln_batch_size"] = rbln_batch_size
-        meta["rbln_max_length"] = rbln_max_length
-        meta["rbln_num_mel_bins"] = rbln_num_mel_bins
-
-        model_input_info = [
+        input_info = [
             (
                 "input_values",
                 [rbln_batch_size, rbln_max_length, rbln_num_mel_bins],
@@ -776,13 +944,19 @@ class RBLNModelForAudioClassification(RBLNModel):
             ),
         ]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=model_input_info, batch_size=rbln_batch_size)
-
-        rbln_config = RBLNConfig.from_rbln_runtime_configs(
-            [rbln_runtime_config],
-            _rbln_meta=meta,
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update(
+            {
+                "batch_size": rbln_batch_size,
+                "max_length": rbln_max_length,
+                "num_mel_bins": rbln_num_mel_bins,
+            }
         )
-
         return rbln_config
 
 
@@ -799,7 +973,6 @@ class RBLNModelForSequenceClassification(RBLNModel):
     Currently, this model class supports the 'XLMRoberta' and 'Roberta' model from the transformers library. Future updates may include support for additional model types.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForSequenceClassification
 
     @classmethod
@@ -807,10 +980,12 @@ class RBLNModelForSequenceClassification(RBLNModel):
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_model_input_names: Optional[List[str]] = None,
-        rbln_batch_size: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
         max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
             model_config, "max_position_embeddings", None
         )
@@ -829,25 +1004,39 @@ class RBLNModelForSequenceClassification(RBLNModel):
             raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
 
         if rbln_model_input_names is None:
-            # These are BERT's inputs
-            rbln_model_input_names = ["input_ids", "attention_mask"]
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_input_names"):
+                    rbln_model_input_names = tokenizer.model_input_names
+                    break
+            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
+                rbln_model_input_names = cls.rbln_model_input_names
+            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
+                original_model_class = getattr(transformers, model_config.architectures[0])
+                input_names_order = inspect.signature(original_model_class.forward).parameters.keys()
+                raise ValueError(
+                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
+                    f"and be sure to make the order of the inputs same as SequenceClassification forward() arguments like ({list(input_names_order)})"
+                )
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
+
         input_info = [
             (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
             for model_input_name in rbln_model_input_names
         ]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
-        rbln_runtime_config.batch_size = rbln_batch_size
-        meta = {"rbln_max_seq_len": rbln_max_seq_len}
-
-        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config
 
 
 class RBLNModelForMaskedLM(RBLNModel):
-    model_type = "rbln_model"
     auto_model_class = AutoModelForMaskedLM
 
     @classmethod
@@ -855,10 +1044,12 @@ class RBLNModelForMaskedLM(RBLNModel):
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_model_input_names: Optional[List[str]] = None,
-        rbln_batch_size: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
         max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
             model_config, "max_position_embeddings", None
         )
@@ -877,18 +1068,33 @@ class RBLNModelForMaskedLM(RBLNModel):
             raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
 
         if rbln_model_input_names is None:
-            # These are BERT's inputs
-            rbln_model_input_names = ["input_ids", "attention_mask"]
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_input_names"):
+                    rbln_model_input_names = tokenizer.model_input_names
+                    break
+            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
+                rbln_model_input_names = cls.rbln_model_input_names
+            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
+                original_model_class = getattr(transformers, model_config.architectures[0])
+                input_names_order = inspect.signature(original_model_class.forward).parameters.keys()
+                raise ValueError(
+                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
+                    f"and be sure to make the order of the inputs same as MaskedLM forward() arguments like ({list(input_names_order)})"
+                )
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
+
         input_info = [
             (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
             for model_input_name in rbln_model_input_names
         ]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
-        rbln_runtime_config.batch_size = rbln_batch_size
-        meta = {"rbln_max_seq_len": rbln_max_seq_len}
-
-        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config