optimum-rbln 0.9.3__py3-none-any.whl → 0.9.4a2__py3-none-any.whl

optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.3'
-__version_tuple__ = version_tuple = (0, 9, 3)
+__version__ = version = '0.9.4a2'
+__version_tuple__ = version_tuple = (0, 9, 4, 'a2')
 
 __commit_id__ = commit_id = None

optimum/rbln/configuration_utils.py

@@ -117,9 +117,14 @@ class RBLNCompileConfig:
         return self
 
     def get_dummy_inputs(
-        self, fill=0, static_tensors: Dict[str, torch.Tensor] = {}, meta_tensor_names: List[str] = []
+        self,
+        fill=0,
+        static_tensors: Optional[Dict[str, torch.Tensor]] = None,
+        meta_tensor_names: Optional[List[str]] = None,
     ):
         dummy = []
+        static_tensors = static_tensors if static_tensors is not None else {}
+        meta_tensor_names = meta_tensor_names if meta_tensor_names is not None else []
         for name, shape, dtype in self.input_info:
             if name in static_tensors:
                 tensor = static_tensors[name]
@@ -255,7 +260,7 @@ class RBLNAutoConfig:
     def load(
         path: str,
         passed_rbln_config: Optional["RBLNModelConfig"] = None,
-        kwargs: Optional[Dict[str, Any]] = {},
+        kwargs: Optional[Dict[str, Any]] = None,
         return_unused_kwargs: bool = False,
     ) -> Union["RBLNModelConfig", Tuple["RBLNModelConfig", Dict[str, Any]]]:
         """
@@ -269,6 +274,8 @@ class RBLNAutoConfig:
         Returns:
             RBLNModelConfig: The loaded RBLNModelConfig.
         """
+        if kwargs is None:
+            kwargs = {}
         cls, config_file = load_config(path)
 
         rbln_keys = [key for key in kwargs.keys() if key.startswith("rbln_")]
@@ -655,7 +662,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         timeout: Optional[int] = None,
         optimum_rbln_version: Optional[str] = None,
         _torch_dtype: Optional[str] = None,
-        _compile_cfgs: List[RBLNCompileConfig] = [],
+        _compile_cfgs: Optional[List[RBLNCompileConfig]] = None,
         *,
         optimize_host_memory: Optional[bool] = None,
         **kwargs: Any,
@@ -708,7 +715,8 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         if self.optimum_rbln_version is None:
             self.optimum_rbln_version = __version__
 
-        self._compile_cfgs: List[RBLNCompileConfig] = _compile_cfgs
+        compile_cfgs = _compile_cfgs if _compile_cfgs is not None else []
+        self._compile_cfgs: List[RBLNCompileConfig] = compile_cfgs
 
         if not isinstance(self._compile_cfgs, list):
             raise ValueError("`compile_cfgs` must be a list of `RBLNCompileConfig`.")

optimum/rbln/diffusers/modeling_diffusers.py

@@ -136,7 +136,7 @@ class RBLNDiffusionMixin:
         *,
         export: bool = None,
         model_save_dir: Optional[PathLike] = None,
-        rbln_config: Dict[str, Any] = {},
+        rbln_config: Optional[Dict[str, Any]] = None,
         lora_ids: Optional[Union[str, List[str]]] = None,
         lora_weights_names: Optional[Union[str, List[str]]] = None,
         lora_scales: Optional[Union[float, List[float]]] = None,

optimum/rbln/diffusers/models/controlnet.py

@@ -215,7 +215,7 @@ class RBLNControlNetModel(RBLNModel):
         encoder_hidden_states: torch.Tensor,
         controlnet_cond: torch.FloatTensor,
         conditioning_scale: torch.Tensor = 1.0,
-        added_cond_kwargs: Dict[str, torch.Tensor] = {},
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
         return_dict: bool = True,
         **kwargs,
     ):

optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -341,7 +341,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
         timestep_cond: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        added_cond_kwargs: Dict[str, torch.Tensor] = {},
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
         down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
         mid_block_additional_residual: Optional[torch.Tensor] = None,
         down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,

optimum/rbln/diffusers/pipelines/auto_pipeline.py

@@ -15,7 +15,7 @@
 
 import importlib
 from pathlib import Path
-from typing import Any, Dict, Type, Union
+from typing import Any, Dict, Optional, Type, Union
 
 from diffusers.models.controlnets import ControlNetUnionModel
 from diffusers.pipelines.auto_pipeline import (
@@ -174,7 +174,7 @@ class RBLNAutoPipelineBase:
         model_id: Union[str, Path],
         *,
         export: bool = None,
-        rbln_config: Union[Dict[str, Any], RBLNModelConfig] = {},
+        rbln_config: Optional[Union[Dict[str, Any], RBLNModelConfig]] = None,
         **kwargs: Any,
     ):
         """

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py

@@ -151,7 +151,9 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
             for image_ in image:
                 self.check_image(image_, prompt, prompt_embeds)
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # Check `controlnet_conditioning_scale`
         if (
@@ -180,7 +182,9 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
                     " the same length as the number of controlnets"
                 )
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         if not isinstance(control_guidance_start, (tuple, list)):
             control_guidance_start = [control_guidance_start]
@@ -254,7 +258,7 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
         control_guidance_end: Union[float, List[float]] = 1.0,
         clip_skip: Optional[int] = None,
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
     ):
         r"""
@@ -393,6 +397,9 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
             )
 
         # 1. Check inputs. Raise error if not correct
+        if callback_on_step_end_tensor_inputs is None:
+            callback_on_step_end_tensor_inputs = ["latents"]
+
         self.check_inputs(
             prompt,
             image,
@@ -503,7 +510,9 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
             image = images
             height, width = image[0].shape[-2:]
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # 5. Prepare timesteps
         timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py

@@ -152,7 +152,9 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
             for image_ in image:
                 self.check_image(image_, prompt, prompt_embeds)
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # Check `controlnet_conditioning_scale`
         if (
@@ -178,7 +180,9 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
                     " the same length as the number of controlnets"
                 )
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         if len(control_guidance_start) != len(control_guidance_end):
             raise ValueError(
@@ -247,7 +251,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
         control_guidance_end: Union[float, List[float]] = 1.0,
         clip_skip: Optional[int] = None,
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
     ):
         r"""
@@ -384,6 +388,9 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
             )
 
         # 1. Check inputs. Raise error if not correct
+        if callback_on_step_end_tensor_inputs is None:
+            callback_on_step_end_tensor_inputs = ["latents"]
+
         self.check_inputs(
             prompt,
             control_image,
@@ -490,7 +497,9 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
 
             control_image = control_images
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # 5. Prepare timesteps
         self.scheduler.set_timesteps(num_inference_steps, device=device)

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py

@@ -178,7 +178,9 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
             for image_ in image:
                 self.check_image(image_, prompt, prompt_embeds)
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # Check `controlnet_conditioning_scale`
         if (
@@ -204,7 +206,9 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
                     " the same length as the number of controlnets"
                 )
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         if not isinstance(control_guidance_start, (tuple, list)):
             control_guidance_start = [control_guidance_start]
@@ -288,7 +292,7 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
         negative_target_size: Optional[Tuple[int, int]] = None,
         clip_skip: Optional[int] = None,
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
     ):
         r"""
@@ -466,6 +470,9 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
             )
 
         # 1. Check inputs. Raise error if not correct
+        if callback_on_step_end_tensor_inputs is None:
+            callback_on_step_end_tensor_inputs = ["latents"]
+
         self.check_inputs(
             prompt,
             prompt_2,
@@ -581,7 +588,9 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
             image = images
             height, width = image[0].shape[-2:]
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # 5. Prepare timesteps
         self.scheduler.set_timesteps(num_inference_steps, device=device)

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py

@@ -190,7 +190,9 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
             for image_ in image:
                 self.check_image(image_, prompt, prompt_embeds)
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # Check `controlnet_conditioning_scale`
         if (
@@ -216,7 +218,9 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
                     " the same length as the number of controlnets"
                 )
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         if not isinstance(control_guidance_start, (tuple, list)):
             control_guidance_start = [control_guidance_start]
@@ -303,7 +307,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
         negative_aesthetic_score: float = 2.5,
         clip_skip: Optional[int] = None,
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
     ):
         r"""
@@ -500,6 +504,9 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
             )
 
         # 1. Check inputs. Raise error if not correct
+        if callback_on_step_end_tensor_inputs is None:
+            callback_on_step_end_tensor_inputs = ["latents"]
+
         self.check_inputs(
             prompt,
             prompt_2,
@@ -618,7 +625,9 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
             control_image = control_images
             height, width = control_image[0].shape[-2:]
         else:
-            assert False
+            raise TypeError(
+                "Unsupported controlnet type. Expected `RBLNControlNetModel` or `RBLNMultiControlNetModel`."
+            )
 
         # 5. Prepare timesteps
         self.scheduler.set_timesteps(num_inference_steps, device=device)

optimum/rbln/diffusers/pipelines/cosmos/cosmos_guardrail.py

@@ -81,7 +81,7 @@ def is_compiled_dir(dir: str) -> bool:
     if not os.path.exists(dir):
         return False
 
-    for root, dirs, files in os.walk(dir):
+    for _, _, files in os.walk(dir):
         for file in files:
             if file.endswith(".rbln"):
                 return True

optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py

@@ -86,7 +86,7 @@ class RBLNCosmosTextToWorldPipeline(RBLNDiffusionMixin, CosmosTextToWorldPipelin
         *,
         export: bool = False,
         safety_checker: Optional[RBLNCosmosSafetyChecker] = None,
-        rbln_config: Dict[str, Any] = {},
+        rbln_config: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
     ):
         """

optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py

@@ -86,7 +86,7 @@ class RBLNCosmosVideoToWorldPipeline(RBLNDiffusionMixin, CosmosVideoToWorldPipel
         *,
         export: bool = False,
         safety_checker: Optional[RBLNCosmosSafetyChecker] = None,
-        rbln_config: Dict[str, Any] = {},
+        rbln_config: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
     ):
         """
@@ -118,7 +118,6 @@ class RBLNCosmosVideoToWorldPipeline(RBLNDiffusionMixin, CosmosVideoToWorldPipel
                 RBLN compilation process. These may include parameters specific to individual submodules
                 or the particular diffusion pipeline being used.
         """
-
         rbln_config, kwargs = cls.get_rbln_config_class().initialize_from_kwargs(rbln_config, **kwargs)
         if safety_checker is None and export:
             safety_checker = RBLNCosmosSafetyChecker(rbln_config=rbln_config.safety_checker)

optimum/rbln/modeling_base.py

@@ -15,7 +15,6 @@
 import importlib
 import os
 import shutil
-from abc import ABC
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
@@ -39,7 +38,7 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
-class PreTrainedModel(ABC):  # noqa: F811
+class PreTrainedModel:  # noqa: F811
     pass
 
 
@@ -63,7 +62,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
         rbln_compiled_models: Optional[rebel.RBLNCompiledModel] = None,
-        rbln_submodules: List["RBLNBaseModel"] = [],
+        rbln_submodules: Optional[List["RBLNBaseModel"]] = None,
         **kwargs,
     ):
         self.model = models
@@ -106,6 +105,8 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             self.model_save_dir = model_save_dir
         self.subfolder = subfolder
 
+        if rbln_submodules is None:
+            rbln_submodules = []
         self.rbln_submodules = rbln_submodules
         self.__post_init__(**kwargs)
 
@@ -181,7 +182,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         # passed from compile function
         rbln_config: Optional[RBLNModelConfig] = None,
         rbln_compiled_models: Optional[Dict[str, rebel.RBLNCompiledModel]] = None,
-        rbln_submodules: List["RBLNBaseModel"] = [],
+        rbln_submodules: Optional[List["RBLNBaseModel"]] = None,
         **kwargs,
     ) -> "RBLNBaseModel":
         if rbln_compiled_models is None:
@@ -217,8 +218,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 )
 
             if len(cls._rbln_submodules) > 0:
-                rbln_submodules = cls._load_submodules(model_save_dir=model_id, rbln_config=rbln_config, **kwargs)
-            else:
+                if rbln_submodules is None:
+                    rbln_submodules = cls._load_submodules(model_save_dir=model_id, rbln_config=rbln_config, **kwargs)
+            elif rbln_submodules is None:
                 rbln_submodules = []
 
             rbln_config.freeze()
@@ -279,9 +281,12 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         config: "PretrainedConfig",
         model_save_dir: Union[Path, str],
         subfolder: Union[Path, str],
-        rbln_submodules: List["RBLNBaseModel"] = [],
+        rbln_submodules: Optional[List["RBLNBaseModel"]] = None,
         **kwargs,
     ):
+        if rbln_submodules is None:
+            rbln_submodules = []
+
         if isinstance(model_save_dir, str):
             model_save_dir = Path(model_save_dir)
 

optimum/rbln/transformers/modeling_attention_utils.py

@@ -195,7 +195,7 @@ class RBLNDecoderOnlyFlashAttentionMixin:
     ) -> int:
         tensor_parallel_size = rbln_config.tensor_parallel_size or 1
         vocab_size = model_config.vocab_size
-        hidden_size = getattr(model_config, "n_embd", None) or getattr(model_config, "hidden_size")
+        hidden_size = getattr(model_config, "n_embd", None) or model_config.hidden_size
         lm_head_params = align(vocab_size, 64) * hidden_size
 
         nbytes_per_param = 2  # Assuming lm_head is always not quantized
@@ -214,7 +214,7 @@ class RBLNDecoderOnlyFlashAttentionMixin:
     ) -> int:
         # This is an *APPROXIMATE* calculation based on the number of parameters
         tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        num_hidden_layers = getattr(model_config, "n_layer", None) or getattr(model_config, "num_hidden_layers")
+        num_hidden_layers = getattr(model_config, "n_layer", None) or model_config.num_hidden_layers
 
         n_model_params = sum(p.numel() for p in model.parameters())
         embed_token_params = sum(p.numel() for p in model.get_input_embeddings().parameters())
@@ -307,9 +307,9 @@ class RBLNDecoderOnlyFlashAttentionMixin:
 
             return dram_per_block
 
-        num_attention_heads = getattr(model_config, "n_head", None) or getattr(model_config, "num_attention_heads")
+        num_attention_heads = getattr(model_config, "n_head", None) or model_config.num_attention_heads
         head_dim = getattr(model_config, "head_dim", None) or model_config.hidden_size // num_attention_heads
-        num_hidden_layers = getattr(model_config, "n_layer", None) or getattr(model_config, "num_hidden_layers")
+        num_hidden_layers = getattr(model_config, "n_layer", None) or model_config.num_hidden_layers
         num_key_value_heads = getattr(model_config, "num_key_value_heads", None) or num_attention_heads
         tensor_parallel_size = rbln_config.tensor_parallel_size or 1
 

optimum/rbln/transformers/modeling_outputs.py

@@ -24,6 +24,7 @@ class RBLNDecoderOnlyOutput(ModelOutput):
     logits: torch.FloatTensor = None
     generate_idx: torch.Tensor = None
     padded_cache_lengths: int = None
+    hidden_states: Tuple[torch.FloatTensor] = None
 
 
 @dataclass

optimum/rbln/transformers/models/auto/auto_factory.py

@@ -150,6 +150,7 @@ class _BaseAutoModelClass:
                 f"from the checkpoint, leading to potential unintended behavior. If this is not intentional, consider calling the "
                 f"`from_pretrained()` method directly from the `RBLN{config.architectures[0]}` class instead.",
                 UserWarning,
+                stacklevel=2,
             )
 
         return model_class

optimum/rbln/transformers/models/colpali/colpali_architecture.py

@@ -156,8 +156,8 @@ class ColPaliAttention(nn.Module):
     def __init__(self, self_attn):
         super().__init__()
         self._original_mod = self_attn
-        self.num_heads = getattr(self._original_mod, "num_heads", None) or getattr(
-            self._original_mod.config, "num_attention_heads"
+        self.num_heads = (
+            getattr(self._original_mod, "num_heads", None) or self._original_mod.config.num_attention_heads
         )
         self.head_dim = self._original_mod.head_dim
         self.scaling = self.head_dim**-0.5

optimum/rbln/transformers/models/colpali/modeling_colpali.py

@@ -332,7 +332,7 @@ class RBLNColPaliForRetrieval(RBLNModel):
         ]
         outputs.append(torch.empty(size=language_model_out_size, dtype=torch.float32, device="cpu"))
         if self.rbln_config.output_hidden_states:
-            for i in range(self.config.vlm_config.text_config.num_hidden_layers + 1):
+            for _ in range(self.config.vlm_config.text_config.num_hidden_layers + 1):
                 outputs.append(torch.empty(size=language_model_hidden_states_size, dtype=torch.float32, device="cpu"))
 
         # Embedding_proj_layer is fused on the bottom of the language model.

optimum/rbln/transformers/models/colqwen2/configuration_colqwen2.py

@@ -58,7 +58,6 @@ class RBLNColQwen2ForRetrievalConfig(RBLNDecoderOnlyModelConfig):
         visual: Optional[RBLNModelConfig] = None,
         batch_size: Optional[int] = None,
         use_inputs_embeds: bool = True,
-        output_hidden_states: Optional[bool] = False,
         **kwargs,
     ):
         super().__init__(use_inputs_embeds=use_inputs_embeds, **kwargs)
@@ -71,4 +70,3 @@ class RBLNColQwen2ForRetrievalConfig(RBLNDecoderOnlyModelConfig):
             raise ValueError("batch_size is not supported for RBLNColQwen2ForRetrievalConfig")
 
         self.visual = visual
-        self.output_hidden_states = output_hidden_states

optimum/rbln/transformers/models/decoderonly/configuration_decoderonly.py

@@ -58,6 +58,7 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
         sliding_window_layers: Optional[List[int]] = None,
         phases: Optional[List[PhaseType]] = None,
         logits_to_keep: Optional[int] = None,
+        output_hidden_states: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -112,6 +113,7 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
                 ["prefill", "decode"] if DecoderOnlyModelForCausalLM is used.
             logits_to_keep (Optional[int]): The number of logits to keep for the decoder.  If set to 0, the decoder will keep all logits.
                 Defaults to 0 if DecoderOnlyModel is used, 1 if DecoderOnlyModelForCausalLM is used.
+            output_hidden_states (Optional[bool]): Whether to output the hidden states of the decoder. Defaults to False.
             kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -232,6 +234,8 @@ class RBLNDecoderOnlyModelConfig(RBLNModelConfig):
         if self.logits_to_keep is not None and self.logits_to_keep > 1:
             raise NotImplementedError("`logits_to_keep` > 1 is currently not supported for RBLN models.")
 
+        self.output_hidden_states = output_hidden_states or False
+
         self.decoder_batch_sizes = None
         if "decode" in self.phases:
             self.decoder_batch_sizes = decoder_batch_sizes

optimum/rbln/transformers/models/decoderonly/configuration_lora.py

@@ -183,7 +183,7 @@ class RBLNLoRAAdapterConfig(RBLNSerializableConfigProtocol):
                 f"Failed to download LoRA adapter '{path.as_posix()}' from HuggingFace Hub. "
                 f"Please check if the model ID is correct or provide a valid local path. "
                 f"Error: {e}"
-            )
+            ) from e
 
     def _load_adapter_config(self) -> Dict[str, Any]:
         """