making prepare_model_for_training flexible

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.

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