huggingface · younesbelkada · Jun 15, 2023 · Jun 14, 2023 · Jun 14, 2023 · Jun 15, 2023
diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
@@ -177,7 +177,12 @@ def from_pretrained(cls, model, model_id, adapter_name="default", is_trainable=F
 
         # load the config
         config = PEFT_TYPE_TO_CONFIG_MAPPING[
-            PeftConfig.from_pretrained(model_id, subfolder=kwargs.get("subfolder", None), **kwargs).peft_type
+            PeftConfig._get_peft_type(
+                model_id,
+                subfolder=kwargs.get("subfolder", None),
+                revision=kwargs.get("revision", None),
+                cache_dir=kwargs.get("cache_dir", None),
+            )
         ].from_pretrained(model_id, subfolder=kwargs.get("subfolder", None), **kwargs)
 
         if (getattr(model, "hf_device_map", None) is not None) and len(
@@ -380,14 +385,39 @@ def set_additional_trainable_modules(self, peft_config, adapter_name):
                 self.modules_to_save.update(peft_config.modules_to_save)
             _set_trainable(self, adapter_name)
 
+    @classmethod
+    def _split_kwargs(cls, kwargs):
+        hf_hub_download_kwargs = {}
+        other_kwargs = {}
+
+        for key, value in kwargs.items():
+            if key in inspect.signature(hf_hub_download).parameters:
+                hf_hub_download_kwargs[key] = value
+            else:
+                other_kwargs[key] = value
+
+        return hf_hub_download_kwargs, other_kwargs
+
     def load_adapter(self, model_id, adapter_name, is_trainable=False, **kwargs):
         from .mapping import PEFT_TYPE_TO_CONFIG_MAPPING
 
+        hf_hub_download_kwargs, kwargs = self._split_kwargs(kwargs)
+
         if adapter_name not in self.peft_config:
             # load the config
             peft_config = PEFT_TYPE_TO_CONFIG_MAPPING[
-                PeftConfig.from_pretrained(model_id, subfolder=kwargs.get("subfolder", None)).peft_type
-            ].from_pretrained(model_id, subfolder=kwargs.get("subfolder", None))
+                PeftConfig._get_peft_type(
+                    model_id,
+                    subfolder=kwargs.get("subfolder", None),
+                    revision=kwargs.get("revision", None),
+                    cache_dir=kwargs.get("cache_dir", None),
+                )
+            ].from_pretrained(
+                model_id,
+                subfolder=kwargs.get("subfolder", None),
+                revision=kwargs.get("revision", None),
+                cache_dir=kwargs.get("cache_dir", None),
+            )
             if isinstance(peft_config, PromptLearningConfig) and is_trainable:
                 raise ValueError("Cannot set a prompt learning adapter to trainable when loading pretrained adapter.")
             else:
@@ -412,12 +442,15 @@ def load_adapter(self, model_id, adapter_name, is_trainable=False, **kwargs):
             if has_remote_safetensors_file:
                 # Priority 1: load safetensors weights
                 filename = hf_hub_download(
-                    model_id, SAFETENSORS_WEIGHTS_NAME, subfolder=kwargs.get("subfolder", None), **kwargs
+                    model_id,
+                    SAFETENSORS_WEIGHTS_NAME,
+                    subfolder=kwargs.get("subfolder", None),
+                    **hf_hub_download_kwargs,
                 )
             else:
                 try:
                     filename = hf_hub_download(
-                        model_id, WEIGHTS_NAME, subfolder=kwargs.get("subfolder", None), **kwargs
+                        model_id, WEIGHTS_NAME, subfolder=kwargs.get("subfolder", None), **hf_hub_download_kwargs
                     )
                 except EntryNotFoundError:
                     raise ValueError(

diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py
@@ -101,7 +101,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None, **kwargs
             else pretrained_model_name_or_path
         )
 
-        hf_hub_download_kwargs, class_kwargs = cls._split_kwargs(kwargs)
+        hf_hub_download_kwargs, class_kwargs, other_kwargs = cls._split_kwargs(kwargs)
 
         if os.path.isfile(os.path.join(path, CONFIG_NAME)):
             config_file = os.path.join(path, CONFIG_NAME)
@@ -141,14 +141,40 @@ def from_json_file(cls, path_json_file, **kwargs):
     def _split_kwargs(cls, kwargs):
         hf_hub_download_kwargs = {}
         class_kwargs = {}
+        other_kwargs = {}
 
         for key, value in kwargs.items():
             if key in inspect.signature(hf_hub_download).parameters:
                 hf_hub_download_kwargs[key] = value
-            else:
+            elif key in list(cls.__annotations__):
                 class_kwargs[key] = value
+            else:
+                other_kwargs[key] = value
+
+        return hf_hub_download_kwargs, class_kwargs, other_kwargs
+
+    @classmethod
+    def _get_peft_type(
+        cls,
+        model_id,
+        subfolder: Optional[str] = None,
+        revision: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+    ):
+        path = os.path.join(model_id, subfolder) if subfolder is not None else model_id
+
+        if os.path.isfile(os.path.join(path, CONFIG_NAME)):
+            config_file = os.path.join(path, CONFIG_NAME)
+        else:
+            try:
+                config_file = hf_hub_download(
+                    model_id, CONFIG_NAME, subfolder=subfolder, revision=revision, cache_dir=cache_dir
+                )
+            except Exception:
+                raise ValueError(f"Can't find '{CONFIG_NAME}' at '{model_id}'")
 
-        return hf_hub_download_kwargs, class_kwargs
+        loaded_attributes = cls.from_json_file(config_file)
+        return loaded_attributes["peft_type"]
 
 
 @dataclass

diff --git a/tests/test_config.py b/tests/test_config.py
@@ -16,7 +16,14 @@
 import tempfile
 import unittest
 
-from peft import AdaptionPromptConfig, LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
+from peft import (
+    AdaptionPromptConfig,
+    LoraConfig,
+    PeftConfig,
+    PrefixTuningConfig,
+    PromptEncoderConfig,
+    PromptTuningConfig,
+)
 
 
 PEFT_MODELS_TO_TEST = [("lewtun/tiny-random-OPTForCausalLM-delta", "v1")]
@@ -97,6 +104,24 @@ def test_to_dict(self):
             self.assertEqual(config.to_dict(), config.__dict__)
             self.assertTrue(isinstance(config.to_dict(), dict))
 
+    def test_from_pretrained_cache_dir(self):
+        r"""
+        Test if the config is correctly loaded with extra kwargs
+        """
+        with tempfile.TemporaryDirectory() as tmp_dirname:
+            for config_class in self.all_config_classes:
+                for model_name, revision in PEFT_MODELS_TO_TEST:
+                    # Test we can load config from delta
+                    _ = config_class.from_pretrained(model_name, revision=revision, cache_dir=tmp_dirname)
+
+    def test_from_pretrained_cache_dir_remote(self):
+        r"""
+        Test if the config is correctly loaded with a checkpoint from the hub
+        """
+        with tempfile.TemporaryDirectory() as tmp_dirname:
+            _ = PeftConfig.from_pretrained("ybelkada/test-st-lora", cache_dir=tmp_dirname)
+            self.assertTrue("models--ybelkada--test-st-lora" in os.listdir(tmp_dirname))
+
     def test_set_attributes(self):
         # manually set attributes and check if they are correctly written
         for config_class in self.all_config_classes:

diff --git a/tests/test_decoder_models.py b/tests/test_decoder_models.py
@@ -103,3 +103,7 @@ def test_training_decoders_gradient_checkpointing(self, test_name, model_id, con
     @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
     def test_inference_safetensors(self, test_name, model_id, config_cls, config_kwargs):
         self._test_inference_safetensors(model_id, config_cls, config_kwargs)
+
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
+    def test_peft_model_device_map(self, test_name, model_id, config_cls, config_kwargs):
+        self._test_peft_model_device_map(model_id, config_cls, config_kwargs)
diff --git a/tests/test_encoder_decoder_models.py b/tests/test_encoder_decoder_models.py
@@ -106,3 +106,7 @@ def test_training_encoder_decoders_gradient_checkpointing(self, test_name, model
     @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
     def test_inference_safetensors(self, test_name, model_id, config_cls, config_kwargs):
         self._test_inference_safetensors(model_id, config_cls, config_kwargs)
+
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
+    def test_peft_model_device_map(self, test_name, model_id, config_cls, config_kwargs):
+        self._test_peft_model_device_map(model_id, config_cls, config_kwargs)
diff --git a/tests/testing_common.py b/tests/testing_common.py
@@ -463,3 +463,25 @@ def _test_training_gradient_checkpointing(self, model_id, config_cls, config_kwa
                 self.assertIsNotNone(param.grad)
             else:
                 self.assertIsNone(param.grad)
+
+    def _test_peft_model_device_map(self, model_id, config_cls, config_kwargs):
+        if config_cls not in (LoraConfig,):
+            return
+
+        config = config_cls(
+            base_model_name_or_path=model_id,
+            **config_kwargs,
+        )
+
+        model = self.transformers_class.from_pretrained(model_id)
+
+        model = get_peft_model(model, config)
+        model = model.to(self.torch_device)
+
+        with tempfile.TemporaryDirectory() as tmp_dirname:
+            model.save_pretrained(tmp_dirname)
+
+            model_from_pretrained = self.transformers_class.from_pretrained(model_id)
+            _ = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname, device_map={"": "cpu"}).to(
+                self.torch_device
+            )