training/instruction-tune-llama-2-int4.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extended Guide: Instruction-tune Llama 2\n",
    "\n",
    "This blog post is an extended guide on instruction-tuning Llama 2 from Meta AI. The idea of the blog post is to focus on creating the instruction dataset, which we can then use to fine-tune the base model of Llama 2 to follow our instructions. \n",
    "\n",
    "The goal is to create a model which can create instructions based on input. The idea behind this is that this can then be used for others to create instruction data from inputs. That's especially helpful if you want to personalize models for, e.g., tweeting, email writing, etc, which means that you would be able to generate an instruction dataset from your emails to then train a model to mimic your email writing. \n",
    "\n",
    "Okay, so can we get started on this? In the blog, we are going to:\n",
    "\n",
    "1. Define our use case in detail and create a prompt template for our instructions\n",
    "2. Create an instruction dataset\n",
    "3. Instruction-tune Llama 2 using `trl` and the `SFTTrainer` \n",
    "4. Test the Model and run Inference\n",
    "\n",
    "## 1. Define our use case in detail and create a template for our instructions\n",
    "\n",
    "Before we describe our use case, we need to better understand what even is an instruction. \n",
    "\n",
    "> An instruction is a piece of text or prompt that is provided to an LLM, like Llama, GPT-4, or Claude, to guide it to generate a response. Instructions allow humans to steer the conversation and constrain the language model's output to be more natural, useful, and aligned with the user's goals. Crafting clear, well-formulated instructions is key to productive conversations.\n",
    "> \n",
    "\n",
    "Examples of instructions are listed below in the table.\n",
    "\n",
    "| Capability | Example Instruction |\n",
    "| --- | --- |\n",
    "| Brainstorming | Provide a diverse set of creative ideas for new flavors of ice cream. |\n",
    "| Classification | Categorize these movies as either comedy, drama, or horror based on the plot summary. |\n",
    "| Closed QA | Answer the question 'What is the capital of France?' with a single word. |\n",
    "| Generation | Write a poem in the style of Robert Frost about nature and the changing seasons. |\n",
    "| Information Extraction | Extract the names of the main characters from this short story. |\n",
    "| Open QA | Why do leaves change color in autumn? Explain the scientific reasons. |\n",
    "| Summarization | Summarize this article on recent advancements in renewable energy in 2-3 sentences. |\n",
    "\n",
    "As described in the beginning, we want to fine-tune a model to be able to generate instructions based on input. (output). We want to use this as a way to create synthetic datasets to personalize LLMs and Agents. \n",
    "\n",
    "Converting the idea into a basic prompt template following the [Alpaca format](https://github.com/tatsu-lab/stanford_alpaca#data-release) we get. \n",
    "\n",
    "```python\n",
    "### Instruction:\n",
    "Use the Input below to create an instruction, which could have been used to generate the input using an LLM. \n",
    "\n",
    "### Input:\n",
    "Dear [boss name],\n",
    "\n",
    "I'm writing to request next week, August 1st through August 4th,\n",
    "off as paid time off.\n",
    "\n",
    "I have some personal matters to attend to that week that require \n",
    "me to be out of the office. I wanted to give you as much advance \n",
    "notice as possible so you can plan accordingly while I am away.\n",
    "\n",
    "Please let me know if you need any additional information from me \n",
    "or have any concerns with me taking next week off. I appreciate you \n",
    "considering this request.\n",
    "\n",
    "Thank you, [Your name]\n",
    "\n",
    "### Response:\n",
    "Write an email to my boss that I need next week 08/01 - 08/04 off.\n",
    "```\n",
    "\n",
    "## 2. Create an instruction dataset\n",
    "\n",
    "After we defined our use case and prompt template, we need to create our instruction dataset. Creating a high-quality instruction dataset is key for a good-performing model. Research shows that [“Less Is More for Alignment”](https://arxiv.org/abs/2305.11206) shows that creating a high-quality, low-quantity (~1000 samples) dataset can achieve the same performance as less-quality and high-quantity datasets. \n",
    "\n",
    "There are several ways to create an instruction dataset, including: \n",
    "\n",
    "1. Using an existing dataset and converting it into an instruction dataset, e.g., [FLAN](https://huggingface.co/datasets/SirNeural/flan_v2)\n",
    "2. Use existing LLMs to create synthetically instruction datasets, e.g., [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca)\n",
    "3. Use Humans to create instructions datasets, e.g., [Dolly](https://huggingface.co/datasets/databricks/databricks-dolly-15k). \n",
    "\n",
    "Each of the methods has its own advantages and disadvantages and depends on the budget, time, and quality requirements. For example, using an existing dataset is the easiest but might not be tailored to your specific use case, while using humans might be the most accurate but can be time-consuming and expensive. It is also possible to combine several methods to create an instruction dataset, as shown in [Orca: Progressive Learning from Complex Explanation Traces of GPT-4.](https://arxiv.org/abs/2306.02707)\n",
    "\n",
    "To keep it simple, we are going to use **[Dolly](https://huggingface.co/datasets/databricks/databricks-dolly-15k)** an open-source dataset of instruction-following records generated by thousands of Databricks employees in several of the behavioral categories outlined in the **[InstructGPT paper](https://arxiv.org/abs/2203.02155)**, including brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization.\n",
    "\n",
    "Let's start coding, but first, let's install our dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"transformers==4.34.0\" \"datasets==2.13.0\" \"peft==0.4.0\" \"accelerate==0.23.0\" \"bitsandbytes==0.41.1\" \"trl==0.4.7\" \"safetensors>=0.3.1\" --upgrade"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To load the **`databricks/databricks-dolly-15k`** dataset, we use the **`load_dataset()`** method from the 🤗 Datasets library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset json (/home/ubuntu/.cache/huggingface/datasets/databricks___json/databricks--databricks-dolly-15k-7427aa6e57c34282/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset size: 15011\n",
      "{'instruction': 'On what month and day was Antwan Deon Odom born?', 'context': 'Antwan Deon Odom (born September 24, 1981) is a former American football defensive end. He was drafted by the Tennessee Titans in the second round of the 2004 NFL Draft. He played college football at Alabama. He has also played for the Cincinnati Bengals.', 'response': 'September 24', 'category': 'closed_qa'}\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from random import randrange\n",
    "\n",
    "# Load dataset from the hub\n",
    "dataset = load_dataset(\"databricks/databricks-dolly-15k\", split=\"train\")\n",
    "\n",
    "print(f\"dataset size: {len(dataset)}\")\n",
    "print(dataset[randrange(len(dataset))])\n",
    "# dataset size: 15011"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To instruct tune our model, we need to convert our structured examples into a collection of tasks described via instructions. We define a **`formatting_function`** that takes a sample and returns a string with our format instruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_instruction(sample):\n",
    "\treturn f\"\"\"### Instruction:\n",
    "Use the Input below to create an instruction, which could have been used to generate the input using an LLM. \n",
    "\n",
    "### Input:\n",
    "{sample['response']}\n",
    "\n",
    "### Response:\n",
    "{sample['instruction']}\n",
    "\"\"\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test our formatting function on a random example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Instruction:\n",
      "Use the Input below to create an instruction, which could have been used to generate the input using an LLM. \n",
      "\n",
      "### Input:\n",
      "Sir Dorabji Tata and Allied Trusts and Sir Ratan Tata Trust\n",
      "\n",
      "### Response:\n",
      "What are the names of Tata trusts which Ratan Tata heads?\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from random import randrange\n",
    "\n",
    "print(format_instruction(dataset[randrange(len(dataset))]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Instruction-tune Llama 2 using `trl` and the `SFTTrainer`\n",
    "\n",
    " We will use the recently introduced method in the paper \"[QLoRA: Quantization-aware Low-Rank Adapter Tuning for Language Generation](https://arxiv.org/abs/2305.14314)\" by Tim Dettmers et al. QLoRA is a new technique to reduce the memory footprint of large language models during finetuning, without sacrificing performance. The TL;DR; of how QLoRA works is:\n",
    "\n",
    "- Quantize the pre-trained model to 4 bits and freeze it.\n",
    "- Attach small, trainable adapter layers. (LoRA)\n",
    "- Finetune only the adapter layers while using the frozen quantized model for context.\n",
    "\n",
    "If you want to learn more about QLoRA and how it works, I recommend you to read the [Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA](https://huggingface.co/blog/4bit-transformers-bitsandbytes) blog post.\n",
    "\n",
    "### Flash Attention\n",
    "\n",
    "Flash Attention is a an method that reorders the attention computation and leverages classical techniques (tiling, recomputation) to significantly speed it up and reduce memory usage from quadratic to linear in sequence length. It is based on the paper \"[FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135)\".\n",
    "The TL;DR; accelerates training up to 3x. Learn more at [FlashAttention](https://github.com/Dao-AILab/flash-attention/tree/main). Flash Attention is currently only available for Ampere (A10, A40, A100, ...) & Hopper (H100, ...) GPUs. You can check if your GPU is supported and install it using the following command:\n",
    "\n",
    "_Note: If your machine has less than 96GB of RAM and lots of CPU cores, reduce the number of `MAX_JOBS`. On the `g5.2xlarge` we used `4`._\n",
    "\n",
    "```bash\n",
    "python -c \"import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'\"\n",
    "pip install ninja packaging\n",
    "MAX_JOBS=4 pip install flash-attn --no-build-isolation\n",
    "```\n",
    "\n",
    "_Installing flash attention can take quite a bit of time (10-45 minutes)._\n",
    "\n",
    "The example supports the use of Flash Attention for all Llama checkpoints, but is not enabled by default. To use Flash Attention change the value of `use_flash_attentin` to `True`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.006969451904296875,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Loading checkpoint shards",
       "rate": null,
       "total": 2,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f3b7caee136b47528f5606c5f95c20e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "use_flash_attention = False\n",
    "\n",
    "# Hugging Face model id\n",
    "model_id = \"NousResearch/Llama-2-7b-hf\"  # non-gated\n",
    "# model_id = \"meta-llama/Llama-2-7b-hf\" # gated\n",
    "\n",
    "\n",
    "# BitsAndBytesConfig int-4 config\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type=\"nf4\", bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "# Load model and tokenizer\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    quantization_config=bnb_config,\n",
    "    use_cache=False,\n",
    "    use_flash_attention_2=use_flash_attention,\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "model.config.pretraining_tp = 1\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = \"right\"\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `SFTTrainer`  supports a native integration with `peft`, which makes it super easy to efficiently instruction tune LLMs. We only need to create our `LoRAConfig` and provide it to the trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model\n",
    "\n",
    "# LoRA config based on QLoRA paper\n",
    "peft_config = LoraConfig(\n",
    "        lora_alpha=16,\n",
    "        lora_dropout=0.1,\n",
    "        r=64,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\", \n",
    ")\n",
    "\n",
    "\n",
    "# prepare model for training\n",
    "model = prepare_model_for_kbit_training(model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we can start our training we need to define the hyperparameters (`TrainingArguments`) we want to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"llama-7-int4-dolly\",\n",
    "    num_train_epochs=3,\n",
    "    per_device_train_batch_size=6 if use_flash_attention else 4,\n",
    "    gradient_accumulation_steps=2,\n",
    "    gradient_checkpointing=True,\n",
    "    optim=\"paged_adamw_32bit\",\n",
    "    logging_steps=10,\n",
    "    save_strategy=\"epoch\",\n",
    "    learning_rate=2e-4,\n",
    "    bf16=True,\n",
    "    fp16=False,\n",
    "    tf32=True,\n",
    "    max_grad_norm=0.3,\n",
    "    warmup_ratio=0.03,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    disable_tqdm=False,  # disable tqdm since with packing values are in correct\n",
    ")\n",
    "\n",
    "\n",
    "# Upcast layer for flash attnetion\n",
    "if use_flash_attention:\n",
    "    from utils.llama_patch import upcast_layer_for_flash_attention\n",
    "    torch_dtype = torch.bfloat16 if args.bf16 else torch.float16 if args.fp16 else torch.float32\n",
    "    model = upcast_layer_for_flash_attention(model, torch_dtype)\n",
    "\n",
    "model = get_peft_model(model, peft_config)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have every building block we need to create our `SFTTrainer` to start then training our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTTrainer\n",
    "\n",
    "max_seq_length = 2048 # max sequence length for model and packing of the dataset\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    train_dataset=dataset,\n",
    "    peft_config=peft_config,\n",
    "    max_seq_length=max_seq_length,\n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,\n",
    "    formatting_func=format_instruction, \n",
    "    args=args,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start training our model by calling the `train()` method on our `Trainer` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    }
   ],
   "source": [
    "# train\n",
    "trainer.train() # there will not be a progress bar since tqdm is disabled\n",
    "\n",
    "# save model\n",
    "trainer.save_model()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training without Flash Attention enabled took 03:08:00 on a `g5.2xlarge`. The instance costs `1,212$/h` which brings us to a total cost of `3.7$`. \n",
    "The training with Flash Attention enabled took 02:08:00 on a `g5.2xlarge`. The instance costs `1,212$/h` which brings us to a total cost of `2.6$`.\n",
    "\n",
    "The results using Flash Attention are mind blowing and impressive, 1.5x faster and 30% cheaper."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Test Model and run Inference\n",
    "\n",
    "After the training is done we want to run and test our model. We will use `peft` and `transformers` to load our LoRA adapter into our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_flash_attention:\n",
    "    # unpatch flash attention\n",
    "    from utils.llama_patch import unplace_flash_attn_with_attn\n",
    "    unplace_flash_attn_with_attn()\n",
    "    \n",
    "import torch\n",
    "from peft import AutoPeftModelForCausalLM\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "\n",
    "args.output_dir = \"llama-7-int4-dolly\"\n",
    "\n",
    "# load base LLM model and tokenizer\n",
    "model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "    args.output_dir,\n",
    "    low_cpu_mem_usage=True,\n",
    "    torch_dtype=torch.float16,\n",
    "    load_in_4bit=True,\n",
    ") \n",
    "tokenizer = AutoTokenizer.from_pretrained(args.output_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s load the dataset again with a random sample to try to generate an instruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset \n",
    "from random import randrange\n",
    "\n",
    "\n",
    "# Load dataset from the hub and get a sample\n",
    "dataset = load_dataset(\"databricks/databricks-dolly-15k\", split=\"train\")\n",
    "sample = dataset[randrange(len(dataset))]\n",
    "\n",
    "prompt = f\"\"\"### Instruction:\n",
    "Use the Input below to create an instruction, which could have been used to generate the input using an LLM. \n",
    "\n",
    "### Input:\n",
    "{sample['response']}\n",
    "\n",
    "### Response:\n",
    "\"\"\"\n",
    "\n",
    "input_ids = tokenizer(prompt, return_tensors=\"pt\", truncation=True).input_ids.cuda()\n",
    "# with torch.inference_mode():\n",
    "outputs = model.generate(input_ids=input_ids, max_new_tokens=100, do_sample=True, top_p=0.9,temperature=0.9)\n",
    "\n",
    "print(f\"Prompt:\\n{sample['response']}\\n\")\n",
    "print(f\"Generated instruction:\\n{tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}\")\n",
    "print(f\"Ground truth:\\n{sample['instruction']}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice! our model works! If want to accelerate our model we can deploy it with [Text Generation Inference](https://github.com/huggingface/text-generation-inference). Therefore we would need to merge our adapter weights into the base model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "    args.output_dir,\n",
    "    low_cpu_mem_usage=True,\n",
    ") \n",
    "\n",
    "# Merge LoRA and base model\n",
    "merged_model = model.merge_and_unload()\n",
    "\n",
    "# Save the merged model\n",
    "merged_model.save_pretrained(\"merged_model\",safe_serialization=True)\n",
    "tokenizer.save_pretrained(\"merged_model\")\n",
    "\n",
    "# push merged model to the hub\n",
    "# merged_model.push_to_hub(\"user/repo\")\n",
    "# tokenizer.push_to_hub(\"user/repo\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}