implementation/qlora_for_ft_falcon_7b.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing QLORA for Fine-Tuning Falcon-7b LLM\n",
    "\n",
    "Here we implement QLORA for fine-tuning the Falcon-7b LLM on a custom dataset. Note the environment requirement of GPU for this notebook is 40GB Nvidia A100 GPU. Thanks for [Amod's](https://medium.com/@amodwrites) excellent article and I have credited it at the bottom."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the required libraries\n",
    "\n",
    "The first step in the process is to load the necessary libraries:\n",
    "\n",
    "* `bitsandbytes`: A lightweight wrapper by Hugging Face (🤗) around CUDA custom functions, particularly 8-bit optimizers and quantization functions. It’s used to handle the quantization process in QLoRA.\n",
    "* `peft`: A library by 🤗 that enables parameter efficient fine tuning.\n",
    "* `transformers`: A library by 🤗 that provides pre-trained models and training utilities for various natural language processing tasks.\n",
    "* `datasets`: Another library by 🤗 that provides easy access to a wide range of datasets.\n",
    "* `accelerate`: A library that by 🤗 abstracts the boilerplate code related to multi-GPUs/TPU/fp16, making it easier to write the training loop of PyTorch models.\n",
    "* `loralib`: A PyTorch implementation of Low-Rank Adaptation (LoRA), a parameter-efficient approach to adapt a large pre-trained deep learning model.\n",
    "* `einops`: A library that simplifies tensor operations.\n",
    "* `xformers`: A collection of composable Transformer building blocks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "!pip install -qU bitsandbytes transformers datasets accelerate loralib einops xformers\n",
    "!pip install -q -U git+https://github.com/huggingface/peft.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import bitsandbytes as bnb\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import transformers\n",
    "from datasets import load_dataset\n",
    "from peft import (\n",
    "    LoraConfig,\n",
    "    PeftConfig,\n",
    "    get_peft_model,\n",
    "    prepare_model_for_kbit_training,\n",
    ")\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForCausalLM,\n",
    "    AutoTokenizer,\n",
    "    BitsAndBytesConfig,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the Pre-Trained Model\n",
    "\n",
    "Here we load the pre-trained model. In this case, \n",
    "* the Falcon 7b model is loaded using the `AutoModelForCausalLM.from_pretrained()` function from the 🤗 transformers library.\n",
    "* the model is loaded in 4-bit using the `BitsAndBytesConfig` from the bitsandbytes library. \n",
    "\n",
    "This is part of the QLoRA process, which involves quantizing the pre-trained weights of the model to 4-bit and keeping them fixed during fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"tiiuae/falcon-7b\"\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    load_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
    ")\n",
    "\n",
    "model =AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    device_map=\"auto\",\n",
    "    trust_remote_code=True,\n",
    "    quantization_config=bnb_config,\n",
    ")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the Model for QLoRA\n",
    "\n",
    "The model is then prepared for QLoRA using the `prepare_model_for_kbit_training()` function. This function initializes the model for QLoRA by setting up the necessary configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = prepare_model_for_kbit_training(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuring LoRA\n",
    "\n",
    "The LoRA configuration is set up using the `LoraConfig` class. The parameters in this configuration include:\n",
    "\n",
    "* `r`: The rank of the update matrices. Lower rank results in smaller update matrices with fewer trainable parameters.\n",
    "* `lora_alpha`: LoRA scaling factor.\n",
    "* `target_modules`: The modules (for example, attention blocks) to apply the LoRA update matrices.\n",
    "* `lora_dropout`: Dropout probability of the LoRA layers.\n",
    "* `bias`: Specifies if the bias parameters should be trained. Can be ‘none’, ‘all’ or ‘lora_only’.\n",
    "\n",
    "The model is then updated with the LoRA configuration using the `get_peft_model()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"query_key_value\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\"\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading and Preparing the Dataset\n",
    "\n",
    "The dataset is loaded using the `load_dataset()` function from the 🤗 datasets library. The dataset is then shuffled and mapped to the `generate_and_tokenize_prompt()` function, which generates and tokenizes each data point in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_prompt(data_point):\n",
    "  return f\"\"\"\n",
    "<Human>: {data_point[\"Context\"]}\n",
    "<AI>: {data_point[\"Response\"]}\n",
    "  \"\"\".strip()\n",
    "\n",
    "def generate_and_tokenize_prompt(data_point):\n",
    "  full_prompt = generate_prompt(data_point)\n",
    "  tokenized_full_prompt = tokenizer(full_prompt, padding=True, truncation=True)\n",
    "  return tokenized_full_prompt\n",
    "\n",
    "from datasets import load_dataset\n",
    "dataset_name = 'Amod/mental_health_counseling_conversations'\n",
    "dataset = load_dataset(dataset_name, split=\"train\")\n",
    "\n",
    "dataset = dataset.shuffle().map(generate_and_tokenize_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up the Training Arguments\n",
    "\n",
    "The training arguments are set up using the `TrainingArguments` class from the transformers library. These arguments include:\n",
    "\n",
    "* `auto_find_batch_size`: If set to True, the trainer will automatically find the largest batch size that fits into memory.\n",
    "* `num_train_epochs`: The number of training epochs.\n",
    "* `learning_rate`: The learning rate for the optimizer.\n",
    "* `bf16`: If set to True, the trainer will use bf16 precision for training.\n",
    "* `save_total_limit`: The total number of checkpoints that can be saved.\n",
    "* `logging_steps`: The number of steps between each logging.\n",
    "* `output_dir`: The directory where the model checkpoints will be saved.\n",
    "* `save_strategy`: The strategy to use for saving checkpoints. In this case, a checkpoint is saved after each epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = transformers.TrainingArguments(\n",
    "    auto_find_batch_size=True,\n",
    "    num_train_epochs=4,\n",
    "    learning_rate=2e-4,\n",
    "    bf16=True,\n",
    "    save_total_limit=4,\n",
    "    logging_steps=10,\n",
    "    output_dir=./output,\n",
    "    save_strategy='epoch',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the Model\n",
    "\n",
    "Finally, the model is trained using the `Trainer` class from the transformers library. The trainer takes in:\n",
    "* the model\n",
    "* the dataset\n",
    "* the training arguments\n",
    "* a data collator for language modeling\n",
    "\n",
    "The training process is then started using the `train()` method of the trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = transformers.Trainer(\n",
    "    model=model,\n",
    "    train_dataset=dataset,\n",
    "    args=training_args,\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "model.config.use_cache = False\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary\n",
    "\n",
    "This is a high-level overview of the code. Each step in the process is crucial for fine-tuning the Falcon 7b model using QLoRA. The combination of these steps allows for efficient and effective fine-tuning of the model on a custom dataset. After the training has completed, you can save the updated LoRAs locally or upload them to Hugging Face to be used with 🤗 PEFT. Or you can choose to merge the LoRAs with the corresponding foundation LLM using ‘model.merge_and_unload()’ function from the PEFT library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Credit\n",
    "\n",
    "https://medium.com/@amodwrites/a-definitive-guide-to-qlora-fine-tuning-falcon-7b-with-peft-78f500a1f337"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}