From 8f1a43cd91cb22b65f1f840f6bca0e156e5e8495 Mon Sep 17 00:00:00 2001
From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
Date: Fri, 3 Nov 2023 12:23:02 +0100
Subject: [PATCH] [`PEFT` / `Tests` ] Fix peft integration failing tests
 (#27258)

fix peft integration issues
---
 .../peft_integration/test_peft_integration.py | 50 +++++++++++++++++--
 1 file changed, 46 insertions(+), 4 deletions(-)

diff --git a/tests/peft_integration/test_peft_integration.py b/tests/peft_integration/test_peft_integration.py
index 809282c770..95f5f84477 100644
--- a/tests/peft_integration/test_peft_integration.py
+++ b/tests/peft_integration/test_peft_integration.py
@@ -98,17 +98,18 @@ class PeftIntegrationTester(unittest.TestCase, PeftTesterMixin):
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     peft_model.save_pretrained(tmpdirname)
 
-                    self.assertTrue("adapter_model.bin" in os.listdir(tmpdirname))
+                    self.assertTrue("adapter_model.safetensors" in os.listdir(tmpdirname))
                     self.assertTrue("adapter_config.json" in os.listdir(tmpdirname))
 
                     self.assertTrue("config.json" not in os.listdir(tmpdirname))
                     self.assertTrue("pytorch_model.bin" not in os.listdir(tmpdirname))
+                    self.assertTrue("model.safetensors" not in os.listdir(tmpdirname))
 
                     peft_model = transformers_class.from_pretrained(tmpdirname).to(torch_device)
                     self.assertTrue(self._check_lora_correctly_converted(peft_model))
 
-                    peft_model.save_pretrained(tmpdirname, safe_serialization=True)
-                    self.assertTrue("adapter_model.safetensors" in os.listdir(tmpdirname))
+                    peft_model.save_pretrained(tmpdirname, safe_serialization=False)
+                    self.assertTrue("adapter_model.bin" in os.listdir(tmpdirname))
                     self.assertTrue("adapter_config.json" in os.listdir(tmpdirname))
 
                     peft_model = transformers_class.from_pretrained(tmpdirname).to(torch_device)
@@ -328,9 +329,10 @@ class PeftIntegrationTester(unittest.TestCase, PeftTesterMixin):
 
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     peft_model.save_pretrained(tmpdirname)
-                    self.assertTrue("adapter_model.bin" in os.listdir(tmpdirname))
+                    self.assertTrue("adapter_model.safetensors" in os.listdir(tmpdirname))
                     self.assertTrue("adapter_config.json" in os.listdir(tmpdirname))
                     self.assertTrue("pytorch_model.bin" not in os.listdir(tmpdirname))
+                    self.assertTrue("model.safetensors" not in os.listdir(tmpdirname))
 
         # 8-bit
         for model_id in self.peft_test_model_ids:
@@ -344,9 +346,49 @@ class PeftIntegrationTester(unittest.TestCase, PeftTesterMixin):
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     peft_model.save_pretrained(tmpdirname)
 
+                    self.assertTrue("adapter_model.safetensors" in os.listdir(tmpdirname))
+                    self.assertTrue("adapter_config.json" in os.listdir(tmpdirname))
+                    self.assertTrue("pytorch_model.bin" not in os.listdir(tmpdirname))
+                    self.assertTrue("model.safetensors" not in os.listdir(tmpdirname))
+
+    @require_torch_gpu
+    def test_peft_save_quantized_regression(self):
+        """
+        Simple test that tests the basic usage of PEFT model save_pretrained with quantized base models
+        Regression test to make sure everything works as expected before the safetensors integration.
+        """
+        # 4bit
+        for model_id in self.peft_test_model_ids:
+            for transformers_class in self.transformers_test_model_classes:
+                peft_model = transformers_class.from_pretrained(model_id, load_in_4bit=True, device_map="auto")
+
+                module = peft_model.model.decoder.layers[0].self_attn.v_proj
+                self.assertTrue(module.__class__.__name__ == "Linear4bit")
+                self.assertTrue(peft_model.hf_device_map is not None)
+
+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    peft_model.save_pretrained(tmpdirname, safe_serialization=False)
                     self.assertTrue("adapter_model.bin" in os.listdir(tmpdirname))
                     self.assertTrue("adapter_config.json" in os.listdir(tmpdirname))
                     self.assertTrue("pytorch_model.bin" not in os.listdir(tmpdirname))
+                    self.assertTrue("model.safetensors" not in os.listdir(tmpdirname))
+
+        # 8-bit
+        for model_id in self.peft_test_model_ids:
+            for transformers_class in self.transformers_test_model_classes:
+                peft_model = transformers_class.from_pretrained(model_id, load_in_8bit=True, device_map="auto")
+
+                module = peft_model.model.decoder.layers[0].self_attn.v_proj
+                self.assertTrue(module.__class__.__name__ == "Linear8bitLt")
+                self.assertTrue(peft_model.hf_device_map is not None)
+
+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    peft_model.save_pretrained(tmpdirname, safe_serialization=False)
+
+                    self.assertTrue("adapter_model.bin" in os.listdir(tmpdirname))
+                    self.assertTrue("adapter_config.json" in os.listdir(tmpdirname))
+                    self.assertTrue("pytorch_model.bin" not in os.listdir(tmpdirname))
+                    self.assertTrue("model.safetensors" not in os.listdir(tmpdirname))
 
     def test_peft_pipeline(self):
         """