From 0e8064801055a124dd7258dee3e25a0351ead4e7 Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Tue, 14 Feb 2023 11:12:37 +0000
Subject: [PATCH] add `prepare_model_for_training`

---
 src/peft/peft_model.py      | 30 ++++++++++++++++++++++++++++++
 tests/test_save_and_load.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index e2932d9..fc91a25 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -288,6 +288,36 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
         else:
             return self.base_model.model(*args, **kwargs)
 
+    def prepare_model_for_training(self):
+        r"""
+        This method wrapps the entire protocol for preparing a model before running a training. This includes:
+            1- Cast the layernorm in fp32 2- making output embedding layer require grads
+        """
+        loaded_in_8bit = getattr(self.base_model, "is_loaded_in_8bit", False)
+
+        for param in self.base_model.parameters():
+            # freeze base model's layers
+            param.requires_grad = False
+
+            if loaded_in_8bit:
+                # cast layer norm in fp32 for stability for 8bit models
+                if param.ndim == 1:
+                    param.data = param.data.to(torch.float32)
+
+        # For backward compatibility
+        if hasattr(self.base_model, "enable_input_require_grads"):
+            self.base_model.enable_input_require_grads()
+        else:
+
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+
+            self.base_model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+        if loaded_in_8bit:
+            # enable gradient checkpointing for memory efficiency
+            self.base_model.model.gradient_checkpointing_enable()
+
 
 class PeftModelForSequenceClassification(PeftModel):
     """
diff --git a/tests/test_save_and_load.py b/tests/test_save_and_load.py
index e17f558..b5d775e 100644
--- a/tests/test_save_and_load.py
+++ b/tests/test_save_and_load.py
@@ -85,6 +85,34 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
                 self.assertTrue(hasattr(model, "from_pretrained"))
                 self.assertTrue(hasattr(model, "push_to_hub"))
 
+    def test_prepare_for_training(self):
+        r"""
+        A test that checks if `prepare_for_training` behaves as expected
+        """
+        for model_id in self.checkpoints_to_test:
+            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)
+
+                dummy_input = torch.LongTensor([[1, 1, 1]])
+                dummy_output = model.get_input_embeddings()(dummy_input)
+
+                self.assertTrue(not dummy_output.requires_grad)
+
+                model.prepare_model_for_training()
+
+                for param in model.base_model.parameters():
+                    self.assertTrue(not param.requires_grad)
+
+                dummy_input = torch.LongTensor([[1, 1, 1]])
+                dummy_output = model.get_input_embeddings()(dummy_input)
+
+                self.assertTrue(dummy_output.requires_grad)
+
     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