Merge pull request #85 from younesbelkada/int8-wrapper

[`core`]  add `prepare_model_for_training`

src/peft/__init__.py

@@ -47,6 +47,7 @@ from .utils import (
     TaskType,
     bloom_model_postprocess_past_key_value,
     get_peft_model_state_dict,
+    prepare_model_for_training,
     set_peft_model_state_dict,
     shift_tokens_right,
 )

src/peft/utils/__init__.py

@@ -23,6 +23,7 @@ from .other import (
     TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING,
     _set_trainable,
     bloom_model_postprocess_past_key_value,
+    prepare_model_for_training,
     shift_tokens_right,
     transpose,
 )

src/peft/utils/other.py

@@ -30,6 +30,59 @@ def bloom_model_postprocess_past_key_value(past_key_values):
     return tuple(zip(keys, values))
 
 
+def prepare_model_for_training(model):
+    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
+        head to fp32
+
+    Args:
+        model, (`transformers.PreTrainedModel`):
+            The loaded model from `transformers`
+    """
+    loaded_in_8bit = getattr(model, "is_loaded_in_8bit", False)
+
+    for param in model.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:
+                param.data = param.data.to(torch.float32)
+
+    # 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)
+
+        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()
+
+    if hasattr(model, "lm_head"):
+        input_dtype = model.lm_head.weight.dtype
+
+        class CastOutputToFloat(torch.nn.Sequential):
+            r"""
+            Manually cast to the expected dtype of the lm_head as sometimes there is a final layer norm that is casted
+            in fp32
+
+            """
+
+            def forward(self, x):
+                return super().forward(x.to(input_dtype)).to(torch.float32)
+
+        model.lm_head = CastOutputToFloat(model.lm_head)
+
+    return model
+
+
 TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = {
     "bloom": bloom_model_postprocess_past_key_value,
 }

tests/test_peft_model.py

@@ -27,6 +27,7 @@ from peft import (
     PromptTuningConfig,
     get_peft_model,
     get_peft_model_state_dict,
+    prepare_model_for_training,
 )
 
 
@@ -85,6 +86,42 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
                 self.assertTrue(hasattr(model, "from_pretrained"))
                 self.assertTrue(hasattr(model, "push_to_hub"))
 
+    def test_prepare_for_training(self):
+        r"""
+        A test that checks if `prepare_for_training` behaves as expected
+        """
+        for model_id in self.checkpoints_to_test:
+            for i, config_cls in enumerate(self.config_classes):
+                model = AutoModelForCausalLM.from_pretrained(model_id)
+                config = config_cls(
+                    base_model_name_or_path=model_id,
+                    **self.config_kwargs[i],
+                )
+                model = get_peft_model(model, config)
+
+                dummy_input = torch.LongTensor([[1, 1, 1]])
+                dummy_output = model.get_input_embeddings()(dummy_input)
+
+                self.assertTrue(not dummy_output.requires_grad)
+
+                # load with `prepare_model_for_training`
+                model = AutoModelForCausalLM.from_pretrained(model_id)
+                model = prepare_model_for_training(model)
+
+                for param in model.parameters():
+                    self.assertTrue(not param.requires_grad)
+
+                config = config_cls(
+                    base_model_name_or_path=model_id,
+                    **self.config_kwargs[i],
+                )
+                model = get_peft_model(model, config)
+
+                dummy_input = torch.LongTensor([[1, 1, 1]])
+                dummy_output = model.get_input_embeddings()(dummy_input)
+
+                self.assertTrue(dummy_output.requires_grad)
+
     def test_save_pretrained(self):
         r"""
         A test to check if `save_pretrained` behaves as expected. This function should only save the state dict of the