fixes

Makefile

@@ -1,6 +1,6 @@
 .PHONY: quality style test docs
 
-check_dirs := src examples
+check_dirs := src tests examples
 
 # Check that source code meets quality standards
 
@@ -9,11 +9,11 @@ quality:
 	black --check $(check_dirs)
 	isort --check-only $(check_dirs)
 	flake8 $(check_dirs)
-	doc-builder style src --max_len 119 --check_only
+	doc-builder style src tests --max_len 119 --check_only
 
 # Format source code automatically and check is there are any problems left that need manual fixing
 style:
 	black $(check_dirs)
 	isort $(check_dirs)
-	doc-builder style src --max_len 119
+	doc-builder style src tests --max_len 119
 	

src/peft/utils/save_and_load.py

@@ -50,7 +50,10 @@ def get_peft_model_state_dict(model, state_dict=None):
             raise NotImplementedError
     else:
         to_return = {}
-        prompt_embeddings = model.get_prompt_embedding_to_save()
+        if model.peft_config.inference_mode:
+            prompt_embeddings = model.prompt_encoder.embedding.weight
+        else:
+            prompt_embeddings = model.get_prompt_embedding_to_save()
         to_return["prompt_embeddings"] = prompt_embeddings
     if model.modules_to_save is not None:
         for key, value in state_dict.items():

tests/test_config.py

@@ -12,11 +12,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import unittest
-import tempfile
 import os
+import tempfile
+import unittest
+
+from peft import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
 
-from peft import LoraConfig, PromptEncoderConfig, PrefixTuningConfig, PromptTuningConfig
 
 class PeftConfigTestMixin:
     all_config_classes = (
@@ -43,16 +44,15 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin):
             self.assertTrue(hasattr(config, "save_pretrained"))
             self.assertTrue(hasattr(config, "from_pretrained"))
             self.assertTrue(hasattr(config, "from_json_file"))
-    
+
     def test_task_type(self):
         for config_class in self.all_config_classes:
             # assert this will not fail
             _ = config_class(task_type="test")
 
-
     def test_save_pretrained(self):
         r"""
-        Test if the config is correctly saved and loaded using 
+        Test if the config is correctly saved and loaded using
         - save_pretrained
         """
         for config_class in self.all_config_classes:
@@ -62,7 +62,7 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin):
 
                 config_from_pretrained = config_class.from_pretrained(tmp_dirname)
                 self.assertEqual(config.to_dict(), config_from_pretrained.to_dict())
-    
+
     def test_from_json_file(self):
         for config_class in self.all_config_classes:
             config = config_class()
@@ -71,12 +71,11 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin):
 
                 config_from_json = config_class.from_json_file(os.path.join(tmp_dirname, "adapter_config.json"))
                 self.assertEqual(config.to_dict(), config_from_json)
-    
 
     def test_to_dict(self):
         r"""
         Test if the config can be correctly converted to a dict using:
-        - to_dict 
+        - to_dict
         - __dict__
         """
         for config_class in self.all_config_classes:
@@ -84,7 +83,6 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin):
             self.assertEqual(config.to_dict(), config.__dict__)
             self.assertTrue(isinstance(config.to_dict(), dict))
 
-    
     def test_set_attributes(self):
         # manually set attributes and check if they are correctly written
         for config_class in self.all_config_classes:
@@ -95,4 +93,4 @@ class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin):
                 config.save_pretrained(tmp_dirname)
 
                 config_from_pretrained = config_class.from_pretrained(tmp_dirname)
-                self.assertEqual(config.to_dict(), config_from_pretrained.to_dict())
+                self.assertEqual(config.to_dict(), config_from_pretrained.to_dict())

tests/test_save_and_load.py

@@ -13,23 +13,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-import torch
 import tempfile
 import unittest
 
+import torch
+from transformers import AutoModelForCausalLM
+
 from peft import (
-    PeftConfig,
-    PeftModel,
     LoraConfig,
-    get_peft_model_state_dict,
+    PeftModel,
     PrefixTuningConfig,
     PromptEncoderConfig,
     PromptTuningConfig,
     get_peft_model,
+    get_peft_model_state_dict,
 )
 
-from transformers import AutoModelForCausalLM
-
 
 class PeftTestMixin:
     checkpoints_to_test = [
@@ -72,11 +71,10 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
     - test if the model has the expected methods
     """
 
-    def test_attributes_lora_model(self):
+    def test_attributes_model(self):
         for model_id in self.checkpoints_to_test:
-            model = AutoModelForCausalLM.from_pretrained(model_id)
-
             for i, config_cls in enumerate(self.config_classes):
+                model = AutoModelForCausalLM.from_pretrained(model_id)
                 config = config_cls(
                     base_model_name_or_path=model_id,
                     **self.config_kwargs[i],
@@ -89,9 +87,8 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
 
     def test_save_pretrained(self):
         r"""
-        A test to check if `save_pretrained` behaves as expected. This function
-        should only save the state dict of the adapter model and not the state
-        dict of the base model. Hence inside each saved directory you should have:
+        A test to check if `save_pretrained` behaves as expected. This function should only save the state dict of the
+        adapter model and not the state dict of the base model. Hence inside each saved directory you should have:
 
         - README.md (that contains an entry `base_model`)
         - adapter_config.json
@@ -99,20 +96,21 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
 
         """
         for model_id in self.checkpoints_to_test:
-            model = AutoModelForCausalLM.from_pretrained(model_id)
-
             for i, config_cls in enumerate(self.config_classes):
+                model = AutoModelForCausalLM.from_pretrained(model_id)
                 config = config_cls(
                     base_model_name_or_path=model_id,
                     **self.config_kwargs[i],
                 )
                 model = get_peft_model(model, config)
+                model.to(model.device)
 
                 with tempfile.TemporaryDirectory() as tmp_dirname:
                     model.save_pretrained(tmp_dirname)
 
                     model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id)
                     model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname)
+                    model_from_pretrained.to(model.device)
 
                     # check if the state dicts are equal
                     state_dict = get_peft_model_state_dict(model)