{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/meta-llama/llama-cookbook/blob/main/getting-started/finetuning/quickstart_peft_finetuning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PEFT Finetuning Quick Start Notebook\n",
    "\n",
    "This notebook shows how to train a Meta Llama 3 model on a single GPU (e.g. A10 with 24GB) using int8 quantization and LoRA finetuning.\n",
    "\n",
    "**_Note:_** To run this notebook on a machine with less than 24GB VRAM (e.g. T4 with 16GB) the context length of the training dataset needs to be adapted.\n",
    "We do this based on the available VRAM during execution.\n",
    "If you run into OOM issues try to further lower the value of train_config.context_length."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 0: Install pre-requirements and convert checkpoint\n",
    "\n",
    "We need to have llama-recipes and its dependencies installed for this notebook. Additionally, we need to log in with the huggingface_cli and make sure that the account is able to to access the Meta Llama weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment if running from Colab T4\n",
    "# ! pip install llama-recipes ipywidgets\n",
    "\n",
    "# import huggingface_hub\n",
    "# huggingface_hub.login()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Load the model\n",
    "\n",
    "Setup training configuration and load the model and tokenizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "68838a4f42f84545912e95b339a31034",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import LlamaForCausalLM, AutoTokenizer\n",
    "from llama_recipes.configs import train_config as TRAIN_CONFIG\n",
    "\n",
    "train_config = TRAIN_CONFIG()\n",
    "train_config.model_name = \"meta-llama/Meta-Llama-3.1-8B\"\n",
    "train_config.num_epochs = 1\n",
    "train_config.run_validation = False\n",
    "train_config.gradient_accumulation_steps = 4\n",
    "train_config.batch_size_training = 1\n",
    "train_config.lr = 3e-4\n",
    "train_config.use_fast_kernels = True\n",
    "train_config.use_fp16 = True\n",
    "train_config.context_length = 1024 if torch.cuda.get_device_properties(0).total_memory < 16e9 else 2048 # T4 16GB or A10 24GB\n",
    "train_config.batching_strategy = \"packing\"\n",
    "train_config.output_dir = \"meta-llama-samsum\"\n",
    "train_config.use_peft = True\n",
    "\n",
    "from transformers import BitsAndBytesConfig\n",
    "config = BitsAndBytesConfig(\n",
    "    load_in_8bit=True,\n",
    ")\n",
    "\n",
    "model = LlamaForCausalLM.from_pretrained(\n",
    "            train_config.model_name,\n",
    "            device_map=\"auto\",\n",
    "            quantization_config=config,\n",
    "            use_cache=False,\n",
    "            attn_implementation=\"sdpa\" if train_config.use_fast_kernels else None,\n",
    "            torch_dtype=torch.float16,\n",
    "        )\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(train_config.model_name)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Check base model\n",
    "\n",
    "Run the base model on an example input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Summarize this dialog:\n",
      "A: Hi Tom, are you busy tomorrow’s afternoon?\n",
      "B: I’m pretty sure I am. What’s up?\n",
      "A: Can you go with me to the animal shelter?.\n",
      "B: What do you want to do?\n",
      "A: I want to get a puppy for my son.\n",
      "B: That will make him so happy.\n",
      "A: Yeah, we’ve discussed it many times. I think he’s ready now.\n",
      "B: That’s good. Raising a dog is a tough issue. Like having a baby ;-) \n",
      "A: I'll get him one of those little dogs.\n",
      "B: One that won't grow up too big;-)\n",
      "A: And eat too much;-))\n",
      "B: Do you know which one he would like?\n",
      "A: Oh, yes, I took him there last Monday. He showed me one that he really liked.\n",
      "B: I bet you had to drag him away.\n",
      "A: He wanted to take it home right away ;-).\n",
      "B: I wonder what he'll name it.\n",
      "A: He said he’d name it after his dead hamster – Lemmy  - he's  a great Motorhead fan :-)))\n",
      "---\n",
      "Summary:\n",
      "A: Hi Tom, are you busy tomorrow’s afternoon?\n",
      "B: I’m pretty sure I am. What’s up?\n",
      "A: Can you go with me to the animal shelter?.\n",
      "B: What do you want to do?\n",
      "A: I want to get a puppy for my son.\n",
      "B: That will make him so happy.\n",
      "A: Yeah, we’ve discussed it many times. I think he’s ready now.\n",
      "B: That’s good. Raising a dog is a tough issue\n"
     ]
    }
   ],
   "source": [
    "eval_prompt = \"\"\"\n",
    "Summarize this dialog:\n",
    "A: Hi Tom, are you busy tomorrow’s afternoon?\n",
    "B: I’m pretty sure I am. What’s up?\n",
    "A: Can you go with me to the animal shelter?.\n",
    "B: What do you want to do?\n",
    "A: I want to get a puppy for my son.\n",
    "B: That will make him so happy.\n",
    "A: Yeah, we’ve discussed it many times. I think he’s ready now.\n",
    "B: That’s good. Raising a dog is a tough issue. Like having a baby ;-) \n",
    "A: I'll get him one of those little dogs.\n",
    "B: One that won't grow up too big;-)\n",
    "A: And eat too much;-))\n",
    "B: Do you know which one he would like?\n",
    "A: Oh, yes, I took him there last Monday. He showed me one that he really liked.\n",
    "B: I bet you had to drag him away.\n",
    "A: He wanted to take it home right away ;-).\n",
    "B: I wonder what he'll name it.\n",
    "A: He said he’d name it after his dead hamster – Lemmy  - he's  a great Motorhead fan :-)))\n",
    "---\n",
    "Summary:\n",
    "\"\"\"\n",
    "\n",
    "model_input = tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "model.eval()\n",
    "with torch.inference_mode():\n",
    "    print(tokenizer.decode(model.generate(**model_input, max_new_tokens=100)[0], skip_special_tokens=True))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the base model only repeats the conversation.\n",
    "\n",
    "### Step 3: Load the preprocessed dataset\n",
    "\n",
    "We load and preprocess the samsum dataset which consists of curated pairs of dialogs and their summarization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_recipes.configs.datasets import samsum_dataset\n",
    "from llama_recipes.utils.dataset_utils import get_dataloader\n",
    "\n",
    "samsum_dataset.trust_remote_code = True\n",
    "\n",
    "train_dataloader = get_dataloader(tokenizer, samsum_dataset, train_config)\n",
    "eval_dataloader = get_dataloader(tokenizer, samsum_dataset, train_config, \"val\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Prepare model for PEFT\n",
    "\n",
    "Let's prepare the model for Parameter Efficient Fine Tuning (PEFT):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import get_peft_model, prepare_model_for_kbit_training, LoraConfig\n",
    "from dataclasses import asdict\n",
    "from llama_recipes.configs import lora_config as LORA_CONFIG\n",
    "\n",
    "lora_config = LORA_CONFIG()\n",
    "lora_config.r = 8\n",
    "lora_config.lora_alpha = 32\n",
    "lora_dropout: float=0.01\n",
    "\n",
    "peft_config = LoraConfig(**asdict(lora_config))\n",
    "\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "model = get_peft_model(model, peft_config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Fine tune the model\n",
    "\n",
    "Here, we fine tune the model for a single epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "from llama_recipes.utils.train_utils import train\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "\n",
    "model.train()\n",
    "\n",
    "optimizer = optim.AdamW(\n",
    "            model.parameters(),\n",
    "            lr=train_config.lr,\n",
    "            weight_decay=train_config.weight_decay,\n",
    "        )\n",
    "scheduler = StepLR(optimizer, step_size=1, gamma=train_config.gamma)\n",
    "\n",
    "# Start the training process\n",
    "results = train(\n",
    "    model,\n",
    "    train_dataloader,\n",
    "    eval_dataloader,\n",
    "    tokenizer,\n",
    "    optimizer,\n",
    "    scheduler,\n",
    "    train_config.gradient_accumulation_steps,\n",
    "    train_config,\n",
    "    None,\n",
    "    None,\n",
    "    None,\n",
    "    wandb_run=None,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6:\n",
    "Save model checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(train_config.output_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 7:\n",
    "Try the fine tuned model on the same example again to see the learning progress:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Summarize this dialog:\n",
      "A: Hi Tom, are you busy tomorrow’s afternoon?\n",
      "B: I’m pretty sure I am. What’s up?\n",
      "A: Can you go with me to the animal shelter?.\n",
      "B: What do you want to do?\n",
      "A: I want to get a puppy for my son.\n",
      "B: That will make him so happy.\n",
      "A: Yeah, we’ve discussed it many times. I think he’s ready now.\n",
      "B: That’s good. Raising a dog is a tough issue. Like having a baby ;-) \n",
      "A: I'll get him one of those little dogs.\n",
      "B: One that won't grow up too big;-)\n",
      "A: And eat too much;-))\n",
      "B: Do you know which one he would like?\n",
      "A: Oh, yes, I took him there last Monday. He showed me one that he really liked.\n",
      "B: I bet you had to drag him away.\n",
      "A: He wanted to take it home right away ;-).\n",
      "B: I wonder what he'll name it.\n",
      "A: He said he’d name it after his dead hamster – Lemmy  - he's  a great Motorhead fan :-)))\n",
      "---\n",
      "Summary:\n",
      "A wants to get a puppy for his son. A took him to the animal shelter last Monday and he showed A one he really liked. A wants to get him one of those little dogs. A and B agree that raising a dog is a tough issue.\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "with torch.inference_mode():\n",
    "    print(tokenizer.decode(model.generate(**model_input, max_new_tokens=100)[0], skip_special_tokens=True))\n"
   ]
  }
 ],
 "metadata": {
  "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.11.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}