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] 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