Train_eluwa.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WE5GJ6s7y0Xo"
      },
      "source": [
        "## Eluwa training notebook\n",
        "\n",
        "This is a straightforward mash-up of two sources - [a tutorial notebook on retraining OPT](https://colab.research.google.com/drive/1jCkpikz0J2o20FBQmYmAGdiKmJGOMo-o) and settings + data from the [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca) github."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TfBzP8gWzkpv"
      },
      "source": [
        "### Install requirements\n",
        "\n",
        "First, run the cells below to install the requirements:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "otj46qRbtpnd"
      },
      "outputs": [],
      "source": [
        "!pip install -q bitsandbytes datasets accelerate loralib transformers peft\n",
        "import os\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import bitsandbytes as bnb\n",
        "from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FOtwYRI3zzXI"
      },
      "source": [
        "### Model loading\n",
        "\n",
        "Here let's load the `opt` model and tokenizer from their huggingface link. The model is loaded in [8bit](https://) mode. This drastically reduces the amount of memory required to run the model. Without it, any attempt to train models above a certain size (say, something like `pythia1b`) will max out the available RAM/VRAM and get your nowhere. To understand what's going on with that `8bit` flag, it might be useful to read [this first](https://https://huggingface.co/blog/hf-bitsandbytes-integration)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cg3fiQOvmI3Q"
      },
      "outputs": [],
      "source": [
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    \"facebook/opt-1.3b\", \n",
        "    load_in_8bit=True, \n",
        "    device_map='auto',\n",
        ")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-1.3b\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QdjWif4CVXR6"
      },
      "source": [
        "### Training\n",
        "Here we load the alpaca dataset. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AQ_HCYruWIHU"
      },
      "outputs": [],
      "source": [
        "import transformers\n",
        "from datasets import load_dataset\n",
        "data = load_dataset(\"tatsu-lab/alpaca\")\n",
        "data = data.map(lambda samples: tokenizer(samples['instruction']), batched=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9fTSZntA1iUG"
      },
      "source": [
        "### Post-processing on the model\n",
        "\n",
        "We need to apply some post-processing on the 8-bit model to enable training, let's freeze all our layers, and cast the layer-norm in `float32` for stability. We also cast the output of the last layer in `float32` for the same reasons."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T-gy-LxM0yAi"
      },
      "outputs": [],
      "source": [
        "for param in model.parameters():\n",
        "  param.requires_grad = False  # freeze the model - train adapters later\n",
        "  if param.ndim == 1:\n",
        "    # cast the small parameters (e.g. layernorm) to fp32 for stability\n",
        "    param.data = param.data.to(torch.float32)\n",
        "\n",
        "model.gradient_checkpointing_enable()  # reduce number of stored activations\n",
        "model.enable_input_require_grads()\n",
        "\n",
        "class CastOutputToFloat(nn.Sequential):\n",
        "  def forward(self, x): return super().forward(x).to(torch.float32)\n",
        "model.lm_head = CastOutputToFloat(model.lm_head)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KwOTr7B3NlM3"
      },
      "source": [
        "### Apply LoRA\n",
        "\n",
        "A LoRA (Low-Rank Adapter) is a way of training a portion of a model, instead of training the entire model. In inference, the LoRA then fits 'on top' of the existing model to modify its outputs.  It takes only a fraction of the memory required to store. I highly recommend [reading the paper](https://arxiv.org/pdf/2106.09685.pdf): without it we'd be stuck retraining entire models from scratch every time. \n",
        "\n",
        "This magic happens with `peft`! Let's load a `PeftModel` and specify that we are going to use low-rank adapters (LoRA) using `get_peft_model` utility function from `peft`. This code will output how many parameters we can actually train here with this method. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4W1j6lxaNnxC"
      },
      "outputs": [],
      "source": [
        "def print_trainable_parameters(model):\n",
        "    \"\"\"\n",
        "    Prints the number of trainable parameters in the model.\n",
        "    \"\"\"\n",
        "    trainable_params = 0\n",
        "    all_param = 0\n",
        "    for _, param in model.named_parameters():\n",
        "        all_param += param.numel()\n",
        "        if param.requires_grad:\n",
        "            trainable_params += param.numel()\n",
        "    print(\n",
        "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
        "    )\n",
        "\n",
        "from peft import LoraConfig, get_peft_model \n",
        "\n",
        "config = LoraConfig(\n",
        "    r=16,\n",
        "    lora_alpha=32,\n",
        "    target_modules=[\"q_proj\", \"v_proj\"],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\"\n",
        ")\n",
        "\n",
        "model = get_peft_model(model, config)\n",
        "print_trainable_parameters(model)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Defining the training arguments\n",
        "The `trainer` function here specifies important things. [The documentation](https://huggingface.co/transformers/v3.0.2/main_classes/trainer.html#transformers.TrainingArguments) covers all these parameters and then some. \n",
        "Things you typically want to pay attention to:\n",
        " is the \n",
        "\n",
        "*   `learning rate` (3e-4 is commonly seen, although that may have been a joke by Karpathy) \n",
        "*   `num_train_epochs`: a measure of how long you want to train your model for. Each time a dataset passes through an algorithm, it is said to have completed an epoch. `max_steps ` is an alternate way of controlling how long you're going to train for. Helps to comment out one if you're using the other.\n",
        "\n"
      ],
      "metadata": {
        "id": "JdEnTEr-_yWN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "trainer = transformers.Trainer(\n",
        "    model=model, \n",
        "    train_dataset=data['train'],\n",
        "    args=transformers.TrainingArguments(\n",
        "        per_device_train_batch_size=8, \n",
        "        gradient_accumulation_steps=4,\n",
        "        warmup_steps=100, \n",
        "        num_train_epochs=1,\n",
        "        #max_steps=1000, \n",
        "        learning_rate=2e-4, \n",
        "        fp16=True,\n",
        "        logging_steps=10, \n",
        "        output_dir='outputs'\n",
        "    ),\n",
        "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)\n",
        ")\n",
        "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!"
      ],
      "metadata": {
        "id": "5FmaIP5T_xqW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "###The actual training"
      ],
      "metadata": {
        "id": "Mmzb7aNRBN6V"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "I keep `resume_from_checkpoint = False` because the function can't seem to handle quantized models very well. This does mean that you're starting from scratch every time. Colab has an annoying habit of restarting runtimes out of nowhere, so if you're training for a long time, say a prayer before you press the button."
      ],
      "metadata": {
        "id": "AgQ_A39WBRgN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "trainer.train(resume_from_checkpoint = False)"
      ],
      "metadata": {
        "id": "_5GaD7dMfMR3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Saving your LoRA\n",
        "\n",
        "Important! Remember to run this so you can save and download your LoRa. The training process will generate .bin files, but they aren't the models you're looking for.\n",
        "\n"
      ],
      "metadata": {
        "id": "GXMR53MrBgAn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.save_pretrained(\"lora-eluwa-opt\")"
      ],
      "metadata": {
        "id": "iQDSGcfQehDc"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}