[core] Some changes with prepare_model_for_training & few fixes (#105)

* changes

* apply to other notebooks

examples/int8_training/Finetune_flan_t5_large_bnb_peft.ipynb

@@ -11,6 +11,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "kBFPA3-aDT7H",
    "metadata": {
@@ -19,7 +20,9 @@
    "source": [
     "In this notebook we will see how to properly use `peft` , `transformers` & `bitsandbytes` to fine-tune `flan-t5-large` in a google colab!\n",
     "\n",
-    "We will finetune the model on [`financial_phrasebank`](https://huggingface.co/datasets/financial_phrasebank) dataset, that consists of pairs of text-labels to classify financial-related sentences, if they are either `positive`, `neutral` or `negative`."
+    "We will finetune the model on [`financial_phrasebank`](https://huggingface.co/datasets/financial_phrasebank) dataset, that consists of pairs of text-labels to classify financial-related sentences, if they are either `positive`, `neutral` or `negative`.\n",
+    "\n",
+    "Note that you could use the same notebook to fine-tune `flan-t5-xl` as well, but you would need to shard the models first to avoid CPU RAM issues on Google Colab, check [these weights](https://huggingface.co/ybelkada/flan-t5-xl-sharded-bf16)."
    ]
   },
   {
@@ -317,13 +320,14 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "4o3ePxrjEDzv",
    "metadata": {
     "id": "4o3ePxrjEDzv"
    },
    "source": [
-    "Some pre-processing needs to be done before training such a model using `peft`, therefore let's import an utiliy function `prepare_model_for_training` that will: \n",
+    "Some pre-processing needs to be done before training such an int8 model using `peft`, therefore let's import an utiliy function `prepare_model_for_int8_training` that will: \n",
     "- Cast the layer norm in `float32` for stability purposes\n",
     "- Add a `forward_hook` to the input embedding layer to enable gradient computation of the input hidden states\n",
     "- Enable gradient checkpointing for more memory-efficient training\n",
@@ -339,9 +343,9 @@
    },
    "outputs": [],
    "source": [
-    "from peft import prepare_model_for_training\n",
+    "from peft import prepare_model_for_int8_training\n",
     "\n",
-    "model = prepare_model_for_training(model)"
+    "model = prepare_model_for_int8_training(model)"
    ]
   },
   {
@@ -1275,7 +1279,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.9.16"
   },
   "vscode": {
    "interpreter": {

examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb

@@ -1158,9 +1158,9 @@
    },
    "outputs": [],
    "source": [
-    "from peft import prepare_model_for_training\n",
+    "from peft import prepare_model_for_int8_training\n",
     "\n",
-    "model = prepare_model_for_training(model, output_embedding_layer_name=\"proj_out\")"
+    "model = prepare_model_for_int8_training(model, output_embedding_layer_name=\"proj_out\")"
    ]
   },
   {

src/peft/__init__.py

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

src/peft/utils/__init__.py

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

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, output_embedding_layer_name="lm_head"):
+def prepare_model_for_int8_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