From 7dfb47242405b2e57cdc739ab8afa47d40b9935d Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Fri, 24 Feb 2023 13:02:40 +0530
Subject: [PATCH 1/2] make gradient checkpointing optional when using PEFT+INT8

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

diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py
index 3f4d2d5..4c7f19c 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_int8_training(model, output_embedding_layer_name="lm_head"):
+def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head", use_gradient_checkpointing=True):
     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
@@ -51,17 +51,17 @@ def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head"
             if param.ndim == 1 and "layer_norm" in name:
                 param.data = param.data.to(torch.float32)
 
-    # For backward compatibility
-    if hasattr(model, "enable_input_require_grads"):
-        model.enable_input_require_grads()
-    else:
+    if loaded_in_8bit and use_gradient_checkpointing:
+        # For backward compatibility
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
 
-        def make_inputs_require_grad(module, input, output):
-            output.requires_grad_(True)
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
 
-        model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
 
-    if loaded_in_8bit:
         # enable gradient checkpointing for memory efficiency
         model.gradient_checkpointing_enable()
 

From 83d6d55d4be292d1633e9af3f81468f3757c3476 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Fri, 24 Feb 2023 18:24:27 +0530
Subject: [PATCH 2/2] address `layernorm` issue

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

diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py
index 4c7f19c..132b033 100644
--- a/src/peft/utils/other.py
+++ b/src/peft/utils/other.py
@@ -30,7 +30,9 @@ def bloom_model_postprocess_past_key_value(past_key_values):
     return tuple(zip(keys, values))
 
 
-def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head", use_gradient_checkpointing=True):
+def prepare_model_for_int8_training(
+    model, output_embedding_layer_name="lm_head", use_gradient_checkpointing=True, layer_norm_names=["layer_norm"]
+):
     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
@@ -48,7 +50,7 @@ def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head"
 
         if loaded_in_8bit:
             # cast layer norm in fp32 for stability for 8bit models
-            if param.ndim == 1 and "layer_norm" in name:
+            if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
                 param.data = param.data.to(torch.float32)
 
     if loaded_in_8bit and use_gradient_checkpointing: