From f1ee1e4c0f037eeafd413ce74aa575ad29a05b94 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 15 Feb 2023 12:51:23 +0530
Subject: [PATCH 1/4] making `prepare_model_for_training` flexible

---
 src/peft/tuners/lora.py | 1 +
 src/peft/utils/other.py | 9 +++++----
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index a2667fa..f7fa9d4 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -189,6 +189,7 @@ class LoraModel(torch.nn.Module):
 # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
 # and modified to work with PyTorch FSDP
 
+
 #  ------------------------------------------------------------------------------------------
 #  Copyright (c) Microsoft Corporation. All rights reserved.
 #  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py
index c062534..878cb03 100644
--- a/src/peft/utils/other.py
+++ b/src/peft/utils/other.py
@@ -30,7 +30,7 @@ def bloom_model_postprocess_past_key_value(past_key_values):
     return tuple(zip(keys, values))
 
 
-def prepare_model_for_training(model):
+def prepare_model_for_training(model, output_embedding_layer_name="lm_head"):
     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 3- Add the upcasting of the lm
@@ -65,8 +65,9 @@ def prepare_model_for_training(model):
         # enable gradient checkpointing for memory efficiency
         model.gradient_checkpointing_enable()
 
-    if hasattr(model, "lm_head"):
-        input_dtype = model.lm_head.weight.dtype
+    if hasattr(model, output_embedding_layer_name):
+        output_embedding_layer = getattr(model, output_embedding_layer_name)
+        input_dtype = output_embedding_layer.weight.dtype
 
         class CastOutputToFloat(torch.nn.Sequential):
             r"""
@@ -78,7 +79,7 @@ def prepare_model_for_training(model):
             def forward(self, x):
                 return super().forward(x.to(input_dtype)).to(torch.float32)
 
-        model.lm_head = CastOutputToFloat(model.lm_head)
+        setattr(model, output_embedding_layer_name, CastOutputToFloat(output_embedding_layer))
 
     return model
 

From 6d6149cf81e64e92e31be324c105134e040110eb Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 15 Feb 2023 14:03:47 +0530
Subject: [PATCH 2/4] preventing other 1D layers to be casted in FP32

---
 src/peft/utils/other.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py
index 878cb03..ac93a1e 100644
--- a/src/peft/utils/other.py
+++ b/src/peft/utils/other.py
@@ -42,13 +42,13 @@ def prepare_model_for_training(model, output_embedding_layer_name="lm_head"):
     """
     loaded_in_8bit = getattr(model, "is_loaded_in_8bit", False)
 
-    for param in model.parameters():
+    for name,param in model.named_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:
+            if param.ndim == 1 and "layer_norm" in name:
                 param.data = param.data.to(torch.float32)
 
     # For backward compatibility

From af7414a67d5aecc68b59e6cee87f55dd05b9b854 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 15 Feb 2023 15:44:26 +0530
Subject: [PATCH 3/4] fix forward signature

---
 src/peft/peft_model.py  | 1 +
 src/peft/utils/other.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index e2932d9..cca1c02 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -81,6 +81,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
             self.modules_to_save = self.peft_config.modules_to_save
             _set_trainable(self)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.forward.__func__.__signature__ = inspect.signature(self.base_model.forward)
 
     def save_pretrained(self, save_directory, **kwargs):
         r"""
diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py
index ac93a1e..82a4a45 100644
--- a/src/peft/utils/other.py
+++ b/src/peft/utils/other.py
@@ -42,7 +42,7 @@ def prepare_model_for_training(model, output_embedding_layer_name="lm_head"):
     """
     loaded_in_8bit = getattr(model, "is_loaded_in_8bit", False)
 
-    for name,param in model.named_parameters():
+    for name, param in model.named_parameters():
         # freeze base model's layers
         param.requires_grad = False
 

From c9b225d2573cbb3cc5725f8f93b0bb1bc1ccfb3e Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 15 Feb 2023 16:58:33 +0530
Subject: [PATCH 4/4] revert

---
 src/peft/peft_model.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index cca1c02..e2932d9 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -81,7 +81,6 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
             self.modules_to_save = self.peft_config.modules_to_save
             _set_trainable(self)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.forward.__func__.__signature__ = inspect.signature(self.base_model.forward)
 
     def save_pretrained(self, save_directory, **kwargs):
         r"""