model: fix pytorch model (#328)

Author: rebel-kblee

src/optimum/rbln/modeling.py

@@ -249,37 +249,11 @@ def get_pytorch_model(
         trust_remote_code: bool = False,
         # Some rbln-config should be applied before loading torch module (i.e. quantized llm)
         rbln_config: Optional[RBLNModelConfig] = None,
-        dtype: Optional[Union[str, torch.dtype, dict]] = None,
         **kwargs,
     ) -> "PreTrainedModel":
         kwargs = cls.update_kwargs(kwargs)
 
-        hf_class = cls.get_hf_class()
-
-        if dtype is not None:
-            config = hf_class.config_class.from_pretrained(
-                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,
-            )
-
-            config, processed_dtype, dtype_orig = _get_dtype(
-                cls=hf_class,
-                dtype=dtype,
-                config=config,
-            )
-
-            kwargs["torch_dtype"] = processed_dtype
-
-            if dtype_orig is not None:
-                hf_class._set_default_dtype(dtype_orig)
-
-        return hf_class.from_pretrained(
+        return cls.get_hf_class().from_pretrained(
             model_id,
             subfolder=subfolder,
             revision=revision,

src/optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -322,22 +322,30 @@ def get_pytorch_model(
         *args,
         rbln_config: Optional[RBLNDecoderOnlyModelConfig] = None,
         num_hidden_layers: Optional[int] = None,
+        trust_remote_code: Optional[bool] = None,
+        torch_dtype: Optional[torch.dtype] = None,
+        dtype: Optional[torch.dtype] = None,
         **kwargs,
     ) -> PreTrainedModel:
         if rbln_config and rbln_config.quantization:
             model = cls.get_quantized_model(model_id, *args, rbln_config=rbln_config, **kwargs)
         else:
+            # TODO : resolve how to control PreTrainedConfig for hf_kwargs
             if num_hidden_layers is not None:
-                trust_remote_code = kwargs.get("trust_remote_code", None)
                 config, kwargs = AutoConfig.from_pretrained(
-                    model_id, return_unused_kwargs=True, num_hidden_layers=num_hidden_layers, **kwargs
+                    model_id,
+                    return_unused_kwargs=True,
+                    trust_remote_code=trust_remote_code,
+                    num_hidden_layers=num_hidden_layers,
+                    **kwargs,
                 )
                 if hasattr(config, "layer_types"):
                     config.layer_types = config.layer_types[:num_hidden_layers]
                 kwargs["config"] = config
-                kwargs["trust_remote_code"] = trust_remote_code
 
-            model = super().get_pytorch_model(model_id, *args, **kwargs)
+            model = super().get_pytorch_model(
+                model_id, *args, trust_remote_code=trust_remote_code, torch_dtype=torch_dtype, dtype=dtype, **kwargs
+            )
 
         return model