redpajama_chat3b_lora.ipynb

{
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3",
      "language": "python"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.10",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Finetune RedPajama using LoRA\n"
      ],
      "metadata": {
        "id": "IvAyQkXPsdXl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "raw_dataset_link = \"https://huggingface.co/datasets/Fredithefish/Instruction-Tuning-with-GPT-4-RedPajama-Chat/resolve/main/alpaca_gpt4_data.jsonl\" #@param {type: \"string\"}\n",
        "\n",
        "\n",
        "base_model_hf = \"togethercomputer/RedPajama-INCITE-Chat-3B-v1\" #@param {type: \"string\"}\n",
        "finetune_epochs = 1 #@param {type: \"integer\"}\n",
        "\n",
        "save_lora_adapters_to_google_drive = \"yes\"  #@param [\"yes\", \"no\"]\n",
        "ADAPTERS_NAME='RedPajama-LoRA' #@param {type: \"string\"}\n",
        "#@markdown  *Name under which the adapters will be saved (not needed, if save_lora_adapters_to_google_drive is \"no\")*  \n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "XkyQv_0CARhp"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade transformers datasets accelerate peft"
      ],
      "metadata": {
        "id": "ANgw2FayFRie",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:18:25.021121Z",
          "iopub.execute_input": "2023-05-17T13:18:25.021548Z",
          "iopub.status.idle": "2023-05-17T13:18:27.097537Z",
          "shell.execute_reply.started": "2023-05-17T13:18:25.021513Z",
          "shell.execute_reply": "2023-05-17T13:18:27.096334Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch \n",
        "import torch.nn as nn \n",
        "import json\n",
        "import transformers \n",
        "from datasets import Dataset \n",
        "from peft import LoraConfig, get_peft_model\n",
        "from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM"
      ],
      "metadata": {
        "id": "Py2I-GnXzEst",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:18:27.100239Z",
          "iopub.execute_input": "2023-05-17T13:18:27.100719Z",
          "iopub.status.idle": "2023-05-17T13:18:40.941097Z",
          "shell.execute_reply.started": "2023-05-17T13:18:27.100680Z",
          "shell.execute_reply": "2023-05-17T13:18:40.940090Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import subprocess\n",
        "\n",
        "# Define the wget command and the URL of the file to download\n",
        "wget_cmd = [\"wget\", str(raw_dataset_link)]\n",
        "\n",
        "# Run the wget command\n",
        "subprocess.run(wget_cmd)"
      ],
      "metadata": {
        "id": "4S_FShyGj2Qx",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:18:40.942344Z",
          "iopub.execute_input": "2023-05-17T13:18:40.943619Z",
          "iopub.status.idle": "2023-05-17T13:18:42.481657Z",
          "shell.execute_reply.started": "2023-05-17T13:18:40.943585Z",
          "shell.execute_reply": "2023-05-17T13:18:42.480426Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "raw_splitted_link = raw_dataset_link.split(\"/\")\n",
        "\n",
        "# read datasets\n",
        "with open(f'./{raw_splitted_link[-1]}', 'r') as fp:\n",
        "    data = [json.loads(x) for x in fp.readlines()]"
      ],
      "metadata": {
        "id": "bT8L5ndmjtzB",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:18:42.485932Z",
          "iopub.execute_input": "2023-05-17T13:18:42.488039Z",
          "iopub.status.idle": "2023-05-17T13:18:42.882218Z",
          "shell.execute_reply.started": "2023-05-17T13:18:42.488006Z",
          "shell.execute_reply": "2023-05-17T13:18:42.881036Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    \"togethercomputer/RedPajama-INCITE-Chat-3B-v1\", \n",
        "    device_map='auto',\n",
        ")"
      ],
      "metadata": {
        "id": "k2em5lH0KpXm",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:18:42.883679Z",
          "iopub.execute_input": "2023-05-17T13:18:42.884716Z",
          "iopub.status.idle": "2023-05-17T13:20:23.600271Z",
          "shell.execute_reply.started": "2023-05-17T13:18:42.884680Z",
          "shell.execute_reply": "2023-05-17T13:20:23.598879Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(\"togethercomputer/RedPajama-INCITE-Chat-3B-v1\")\n",
        "tokenizer.pad_token = tokenizer.eos_token"
      ],
      "metadata": {
        "id": "yxwx1oE-SBRE",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:20:23.609081Z",
          "iopub.execute_input": "2023-05-17T13:20:23.612525Z",
          "iopub.status.idle": "2023-05-17T13:20:24.348518Z",
          "shell.execute_reply.started": "2023-05-17T13:20:23.612455Z",
          "shell.execute_reply": "2023-05-17T13:20:24.347494Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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)"
      ],
      "metadata": {
        "id": "_9kwEzEYjecL",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:20:24.350136Z",
          "iopub.execute_input": "2023-05-17T13:20:24.350650Z",
          "iopub.status.idle": "2023-05-17T13:20:24.364505Z",
          "shell.execute_reply.started": "2023-05-17T13:20:24.350609Z",
          "shell.execute_reply": "2023-05-17T13:20:24.363224Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.gradient_checkpointing_enable()  # reduce number of stored activations\n",
        "model.enable_input_require_grads()"
      ],
      "metadata": {
        "id": "FfdOMZVwjit5",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:20:24.365971Z",
          "iopub.execute_input": "2023-05-17T13:20:24.366890Z",
          "iopub.status.idle": "2023-05-17T13:20:24.373742Z",
          "shell.execute_reply.started": "2023-05-17T13:20:24.366854Z",
          "shell.execute_reply": "2023-05-17T13:20:24.372734Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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",
        "    )"
      ],
      "metadata": {
        "id": "4o-kiH9yjkv_",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:20:24.374979Z",
          "iopub.execute_input": "2023-05-17T13:20:24.375337Z",
          "iopub.status.idle": "2023-05-17T13:20:24.385245Z",
          "shell.execute_reply.started": "2023-05-17T13:20:24.375304Z",
          "shell.execute_reply": "2023-05-17T13:20:24.384330Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = LoraConfig(\n",
        "    r=16,\n",
        "    lora_alpha=32,\n",
        "    target_modules=[\"query_key_value\", \"xxx\"],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\"\n",
        ")"
      ],
      "metadata": {
        "id": "TQPhiv2-jnr6",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:20:24.389820Z",
          "iopub.execute_input": "2023-05-17T13:20:24.390078Z",
          "iopub.status.idle": "2023-05-17T13:20:24.396786Z",
          "shell.execute_reply.started": "2023-05-17T13:20:24.390056Z",
          "shell.execute_reply": "2023-05-17T13:20:24.395940Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = get_peft_model(model, config)\n",
        "print_trainable_parameters(model)"
      ],
      "metadata": {
        "id": "0w8efTdTkChr",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:20:24.398228Z",
          "iopub.execute_input": "2023-05-17T13:20:24.398796Z",
          "iopub.status.idle": "2023-05-17T13:20:35.668456Z",
          "shell.execute_reply.started": "2023-05-17T13:20:24.398760Z",
          "shell.execute_reply": "2023-05-17T13:20:35.667388Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data = Dataset.from_list(data)\n",
        "data = data.map(lambda samples: tokenizer(samples['text']), batched=True)"
      ],
      "metadata": {
        "id": "2JoG-E0ckVtz",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:20:35.669985Z",
          "iopub.execute_input": "2023-05-17T13:20:35.671079Z",
          "iopub.status.idle": "2023-05-17T13:21:02.391459Z",
          "shell.execute_reply.started": "2023-05-17T13:20:35.671024Z",
          "shell.execute_reply": "2023-05-17T13:21:02.390412Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "trainer = transformers.Trainer(\n",
        "    model=model, \n",
        "    train_dataset=data,\n",
        "    args=transformers.TrainingArguments(\n",
        "        per_device_train_batch_size=4, \n",
        "        gradient_accumulation_steps=4,\n",
        "        warmup_steps=80, \n",
        "        learning_rate=2e-4, \n",
        "        fp16=True,\n",
        "        num_train_epochs=finetune_epochs,\n",
        "        logging_steps=1, \n",
        "        output_dir='outputs',\n",
        "        report_to=\"none\"\n",
        "    ),\n",
        "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)\n",
        ")"
      ],
      "metadata": {
        "id": "iUm75M-gkdNT",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:21:02.393239Z",
          "iopub.execute_input": "2023-05-17T13:21:02.394008Z",
          "iopub.status.idle": "2023-05-17T13:21:03.309538Z",
          "shell.execute_reply.started": "2023-05-17T13:21:02.393969Z",
          "shell.execute_reply": "2023-05-17T13:21:03.308558Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.config.use_cache = False"
      ],
      "metadata": {
        "id": "Tv3XwCZ5lHOY",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:21:03.310787Z",
          "iopub.execute_input": "2023-05-17T13:21:03.311113Z",
          "iopub.status.idle": "2023-05-17T13:21:03.316291Z",
          "shell.execute_reply.started": "2023-05-17T13:21:03.311082Z",
          "shell.execute_reply": "2023-05-17T13:21:03.315335Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "dKycPIy_lOe6",
        "execution": {
          "iopub.status.busy": "2023-05-17T13:21:03.317498Z",
          "iopub.execute_input": "2023-05-17T13:21:03.318415Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if save_lora_adapters_to_google_drive == \"yes\":\n",
        "  from google.colab import drive\n",
        "  drive.mount('/gdrive')\n",
        "  # save the trained adapter to disk\n",
        "  model.save_pretrained(f\"/gdrive/My Drive/Colab Notebooks/Models/{ADAPTERS_NAME}\")\n",
        "else:\n",
        "  pass"
      ],
      "metadata": {
        "id": "lM0kbCPJCaun"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}