Llama_2_Fine_Tune_With_Your_Dataset.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMRpOGQjpoxurAw6GSGiv5U",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/amitsangani/Llama/blob/main/Llama_2_Fine_Tune_With_Your_Dataset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Practical Introduction to Llama 2 Fine-Tuning**"
      ],
      "metadata": {
        "id": "lJ4soaSMhgUo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **How to fine-tune Llama 2 🦙**\n",
        "In this section, we will fine-tune a 7B parameter Llama 2 model on a T4 GPU with high RAM using Google Colab. Note that a T4 only has 16 GB of VRAM, which is barely enough to store Llama 2–7b’s weights (7b × 2 bytes = 14 GB in FP16). In addition, we need to consider the overhead due to optimizer states, gradients, and forward activations (see this excellent article for more information). This means that a full fine-tuning is not possible here: we need parameter-efficient fine-tuning (PEFT) techniques like LoRA or QLoRA.\n",
        "\n",
        "To drastically reduce the VRAM usage, we must fine-tune the model in 4-bit precision, which is why we’ll use QLoRA here. The good thing is that we can leverage the Hugging Face ecosystem with the transformers, accelerate, peft, trl, and bitsandbytes libraries. First, we install and load these libraries."
      ],
      "metadata": {
        "id": "9sFY3weAiG0X"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.31.0 trl==0.4.7"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JJANcmlUhpk-",
        "outputId": "7d51f366-8025-4980-f791-cb50ba5c6f70"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m244.2/244.2 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m72.9/72.9 kB\u001b[0m \u001b[31m7.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m92.5/92.5 MB\u001b[0m \u001b[31m5.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.4/7.4 MB\u001b[0m \u001b[31m50.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.4/77.4 kB\u001b[0m \u001b[31m7.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m25.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m22.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m84.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m519.3/519.3 kB\u001b[0m \u001b[31m35.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.1/194.1 kB\u001b[0m \u001b[31m16.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m11.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import torch\n",
        "from datasets import load_dataset\n",
        "from transformers import (\n",
        "    AutoModelForCausalLM,\n",
        "    AutoTokenizer,\n",
        "    BitsAndBytesConfig,\n",
        "    HfArgumentParser,\n",
        "    TrainingArguments,\n",
        "    pipeline,\n",
        "    logging,\n",
        ")\n",
        "from peft import LoraConfig, PeftModel\n",
        "from trl import SFTTrainer"
      ],
      "metadata": {
        "id": "jP-0sDCvhppv"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, we want to load a llama-2-7b-chat-hf model (chat model) and train it on the mlabonne/guanaco-llama2-1k (1,000 samples), which will produce our fine-tuned model llama-2-7b-miniguanaco. If you’re interested in how this dataset was created, you can check this notebook. Feel free to change it: there are many good datasets on the Hugging Face Hub, like databricks/databricks-dolly-15k.\n",
        "\n",
        "QLoRA will use a rank of 64 with a scaling parameter of 16 (see this article for more information about LoRA parameters). We’ll load the Llama 2 model directly in 4-bit precision using the NF4 type and train it for one epoch. To get more information about the other parameters, check the TrainingArguments, PeftModel, and SFTTrainer documentation.\n"
      ],
      "metadata": {
        "id": "dy79lphLjFZo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# The model that you want to train from the Hugging Face hub\n",
        "model_name = \"NousResearch/Llama-2-7b-chat-hf\"\n",
        "\n",
        "# The instruction dataset to use\n",
        "dataset_name = \"mlabonne/guanaco-llama2-1k\"\n",
        "\n",
        "# Fine-tuned model name\n",
        "new_model = \"llama-2-7b-miniguanaco\""
      ],
      "metadata": {
        "id": "WZ9KcGFqjVj-"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "################################################################################\n",
        "# QLoRA parameters\n",
        "################################################################################\n",
        "\n",
        "# LoRA attention dimension\n",
        "lora_r = 64\n",
        "\n",
        "# Alpha parameter for LoRA scaling\n",
        "lora_alpha = 16\n",
        "\n",
        "# Dropout probability for LoRA layers\n",
        "lora_dropout = 0.1"
      ],
      "metadata": {
        "id": "7Oxu2QXTjb38"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "################################################################################\n",
        "# bitsandbytes parameters\n",
        "################################################################################\n",
        "\n",
        "# Activate 4-bit precision base model loading\n",
        "use_4bit = True\n",
        "\n",
        "# Compute dtype for 4-bit base models\n",
        "bnb_4bit_compute_dtype = \"float16\"\n",
        "\n",
        "# Quantization type (fp4 or nf4)\n",
        "bnb_4bit_quant_type = \"nf4\"\n",
        "\n",
        "# Activate nested quantization for 4-bit base models (double quantization)\n",
        "use_nested_quant = False"
      ],
      "metadata": {
        "id": "tdXSjZOLjf_-"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "################################################################################\n",
        "# TrainingArguments parameters\n",
        "################################################################################\n",
        "\n",
        "# Output directory where the model predictions and checkpoints will be stored\n",
        "output_dir = \"./results\"\n",
        "\n",
        "# Number of training epochs\n",
        "num_train_epochs = 1\n",
        "\n",
        "# Enable fp16/bf16 training (set bf16 to True with an A100)\n",
        "fp16 = False\n",
        "bf16 = False\n",
        "\n",
        "# Batch size per GPU for training\n",
        "per_device_train_batch_size = 4\n",
        "\n",
        "# Batch size per GPU for evaluation\n",
        "per_device_eval_batch_size = 4\n",
        "\n",
        "# Number of update steps to accumulate the gradients for\n",
        "gradient_accumulation_steps = 1\n",
        "\n",
        "# Enable gradient checkpointing\n",
        "gradient_checkpointing = True\n",
        "\n",
        "# Maximum gradient normal (gradient clipping)\n",
        "max_grad_norm = 0.3\n",
        "\n",
        "# Initial learning rate (AdamW optimizer)\n",
        "learning_rate = 2e-4\n",
        "\n",
        "# Weight decay to apply to all layers except bias/LayerNorm weights\n",
        "weight_decay = 0.001\n",
        "\n",
        "# Optimizer to use\n",
        "optim = \"paged_adamw_32bit\"\n",
        "\n",
        "# Learning rate schedule (constant a bit better than cosine)\n",
        "lr_scheduler_type = \"constant\"\n",
        "\n",
        "# Number of training steps (overrides num_train_epochs)\n",
        "max_steps = -1\n",
        "\n",
        "# Ratio of steps for a linear warmup (from 0 to learning rate)\n",
        "warmup_ratio = 0.03\n",
        "\n",
        "# Group sequences into batches with same length\n",
        "# Saves memory and speeds up training considerably\n",
        "group_by_length = True\n",
        "\n",
        "# Save checkpoint every X updates steps\n",
        "save_steps = 25\n",
        "\n",
        "# Log every X updates steps\n",
        "logging_steps = 25"
      ],
      "metadata": {
        "id": "GOQhnX68h_6H"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "################################################################################\n",
        "# SFT parameters\n",
        "################################################################################\n",
        "\n",
        "# Maximum sequence length to use\n",
        "max_seq_length = None\n",
        "\n",
        "# Pack multiple short examples in the same input sequence to increase efficiency\n",
        "packing = False\n",
        "\n",
        "# Load the entire model on the GPU 0\n",
        "device_map = {\"\": 0}"
      ],
      "metadata": {
        "id": "vlaRC-1ejoAg"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **We can now start a fine-tuning process**\n",
        "\n",
        "1. Load the pre-processed dataset.\n",
        "2. Configure bitsandbytes for 4-bit quantization\n",
        "3. Load Llama 2 model in 4-bit precision on a GPU with tokenizer\n",
        "4. Load configurations for QLoRA, regular training params using SFTTrainer"
      ],
      "metadata": {
        "id": "Yj2SRFXfjy19"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 1: Load dataset (you can process it here)\n",
        "dataset = load_dataset(dataset_name, split=\"train\")"
      ],
      "metadata": {
        "id": "S5EGK7ZokhiW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 2: configure bitsandbytes for 4-bit quantization\n",
        "\n",
        "# Load tokenizer and model with QLoRA configuration\n",
        "compute_dtype = getattr(torch, bnb_4bit_compute_dtype)\n",
        "\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=use_4bit,\n",
        "    bnb_4bit_quant_type=bnb_4bit_quant_type,\n",
        "    bnb_4bit_compute_dtype=compute_dtype,\n",
        "    bnb_4bit_use_double_quant=use_nested_quant,\n",
        ")"
      ],
      "metadata": {
        "id": "kJN8qgbpk2Qq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 3: Load Llama 2 model in 4-bit precision on a GPU with tokenizer\n",
        "\n",
        "# Check GPU compatibility with bfloat16\n",
        "if compute_dtype == torch.float16 and use_4bit:\n",
        "    major, _ = torch.cuda.get_device_capability()\n",
        "    if major >= 8:\n",
        "        print(\"=\" * 80)\n",
        "        print(\"Your GPU supports bfloat16: accelerate training with bf16=True\")\n",
        "        print(\"=\" * 80)\n",
        "\n",
        "# Load Llama 2 base model\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_name,\n",
        "    quantization_config=bnb_config,\n",
        "    device_map=device_map\n",
        ")\n",
        "\n",
        "model.config.use_cache = False\n",
        "model.config.pretraining_tp = 1\n",
        "\n",
        "# Load LLaMA tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_side = \"right\" # Fix weird overflow issue with fp16 training"
      ],
      "metadata": {
        "id": "V7BgWp2BlM8N"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 4: Load configurations for QLoRA\n",
        "\n",
        "# Load LoRA configuration\n",
        "peft_config = LoraConfig(\n",
        "    lora_alpha=lora_alpha,\n",
        "    lora_dropout=lora_dropout,\n",
        "    r=lora_r,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")"
      ],
      "metadata": {
        "id": "kusnpXbUlkqi"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#**Start training the model with given dataset**"
      ],
      "metadata": {
        "id": "dlsWiGROl6Eo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Set training parameters\n",
        "training_arguments = TrainingArguments(\n",
        "    output_dir=output_dir,\n",
        "    num_train_epochs=num_train_epochs,\n",
        "    per_device_train_batch_size=per_device_train_batch_size,\n",
        "    gradient_accumulation_steps=gradient_accumulation_steps,\n",
        "    optim=optim,\n",
        "    save_steps=save_steps,\n",
        "    logging_steps=logging_steps,\n",
        "    learning_rate=learning_rate,\n",
        "    weight_decay=weight_decay,\n",
        "    fp16=fp16,\n",
        "    bf16=bf16,\n",
        "    max_grad_norm=max_grad_norm,\n",
        "    max_steps=max_steps,\n",
        "    warmup_ratio=warmup_ratio,\n",
        "    group_by_length=group_by_length,\n",
        "    lr_scheduler_type=lr_scheduler_type,\n",
        "    report_to=\"tensorboard\"\n",
        ")\n",
        "\n",
        "# Set supervised fine-tuning parameters\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    train_dataset=dataset,\n",
        "    peft_config=peft_config,\n",
        "    dataset_text_field=\"text\",\n",
        "    max_seq_length=max_seq_length,\n",
        "    tokenizer=tokenizer,\n",
        "    args=training_arguments,\n",
        "    packing=packing,\n",
        ")\n",
        "\n",
        "# Train model\n",
        "trainer.train()\n",
        "\n",
        "# Save trained model\n",
        "trainer.model.save_pretrained(new_model)"
      ],
      "metadata": {
        "id": "Z1gEgaGhmL_k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The training can be very long, depending on the size of your dataset. Here, it took less than an hour on a T4 GPU. We can check the plots on tensorboard, as follows:"
      ],
      "metadata": {
        "id": "ag7Goqb1maby"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir results/runs"
      ],
      "metadata": {
        "id": "Bmbk71vkm7zg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let’s make sure that the model is behaving correctly. It would require a more exhaustive evaluation, but we can use the text generation pipeline to ask questions like “What is a large language model?” Note that I’m formatting the input to match Llama 2’s prompt template."
      ],
      "metadata": {
        "id": "JK_EeD6ymyix"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Ignore warnings\n",
        "logging.set_verbosity(logging.CRITICAL)\n",
        "\n",
        "# Run text generation pipeline with our next model\n",
        "prompt = \"Explain a large language model and its architecture?\"\n",
        "pipe = pipeline(task=\"text-generation\", model=model, tokenizer=tokenizer, max_length=200)\n",
        "result = pipe(f\"<s>[INST] {prompt} [/INST]\")\n",
        "print(result[0]['generated_text'])"
      ],
      "metadata": {
        "id": "ZE_SnhI0m4iF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "From experience, it is very coherent for a model with only 7 billion parameters. You can play with it and ask harder questions from evaluation datasets like BigBench-Hard. Guanaco is an excellent dataset that has produced high-quality models in the past. You can train a Llama 2 model on the entire dataset using mlabonne/guanaco-llama2.\n",
        "\n",
        "How can we store our new llama-2-7b-miniguanaco model now? We need to merge the weights from LoRA with the base model. Unfortunately, as far as I know, there is no straightforward way to do it: we need to reload the base model in FP16 precision and use the peft library to merge everything. Alas, it also creates a problem with the VRAM (despite emptying it), so I recommend restarting the notebook, re-executing the three first cells, and then executing the next one. Please contact me if you know a fix!\n"
      ],
      "metadata": {
        "id": "iqcSsD3EnTvc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Reload model in FP16 and merge it with LoRA weights\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_name,\n",
        "    low_cpu_mem_usage=True,\n",
        "    return_dict=True,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=device_map,\n",
        ")\n",
        "model = PeftModel.from_pretrained(base_model, new_model)\n",
        "model = model.merge_and_unload()\n",
        "\n",
        "# Reload tokenizer to save it\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_side = \"right\""
      ],
      "metadata": {
        "id": "jFkEgDzJnZ_D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Our weights are merged and we reloaded the tokenizer. You can now push everything to the Hugging Face Hub to save our model."
      ],
      "metadata": {
        "id": "z__SI47mncdY"
      }
    }
  ]
}