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