recipes/quickstart/finetuning/quickstart_peft_finetuning.ipynb

{
 "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-recipes/blob/main/recipes/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": "c7963d43806d432aaa3d00e2055e355c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "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",
    "\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.no_grad():\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": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/miniconda3/envs/llama/lib/python3.11/site-packages/datasets/load.py:1486: FutureWarning: The repository for samsum contains custom code which must be executed to correctly load the dataset. You can inspect the repository content at https://hf.co/datasets/samsum\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this dataset from the next major release of `datasets`.\n",
      "  warnings.warn(\n",
      "Preprocessing dataset: 100%|██████████| 14732/14732 [00:02<00:00, 6124.69it/s]\n"
     ]
    }
   ],
   "source": [
    "from llama_recipes.configs.datasets import samsum_dataset\n",
    "from llama_recipes.data.concatenator import ConcatDataset\n",
    "from llama_recipes.utils.config_utils import get_dataloader_kwargs\n",
    "from llama_recipes.utils.dataset_utils import get_preprocessed_dataset\n",
    "\n",
    "train_dataset = get_preprocessed_dataset(tokenizer, samsum_dataset, 'train')\n",
    "\n",
    "train_dl_kwargs = get_dataloader_kwargs(train_config, train_dataset, tokenizer, \"train\")\n",
    "\n",
    "if train_config.batching_strategy == \"packing\":\n",
    "        train_dataset = ConcatDataset(train_dataset, chunk_size=train_config.context_length)\n",
    "\n",
    "# Create DataLoaders for the training and validation dataset\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    train_dataset,\n",
    "    num_workers=train_config.num_workers_dataloader,\n",
    "    pin_memory=True,\n",
    "    **train_dl_kwargs,\n",
    ")"
   ]
  },
  {
   "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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/miniconda3/envs/llama/lib/python3.11/site-packages/torch/cuda/memory.py:330: FutureWarning: torch.cuda.reset_max_memory_allocated now calls torch.cuda.reset_peak_memory_stats, which resets /all/ peak memory stats.\n",
      "  warnings.warn(\n",
      "Training Epoch: 1:   0%|\u001b[34m          \u001b[0m| 0/319 [00:00<?, ?it/s]huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "/home/ubuntu/miniconda3/envs/llama/lib/python3.11/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/home/ubuntu/miniconda3/envs/llama/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:316: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n",
      "Training Epoch: 1/1, step 1278/1279 completed (loss: 0.27870458364486694): : 320it [2:07:09, 23.84s/it]                      3.94s/it]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Max CUDA memory allocated was 15 GB\n",
      "Max CUDA memory reserved was 16 GB\n",
      "Peak active CUDA memory was 15 GB\n",
      "CUDA Malloc retries : 0\n",
      "CPU Total Peak Memory consumed during the train (max): 2 GB\n",
      "Epoch 1: train_perplexity=1.3403, train_epoch_loss=0.2929, epoch time 7630.169942979002s\n"
     ]
    }
   ],
   "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",
    "    None,\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": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/miniconda3/envs/llama/lib/python3.11/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "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 her son. She will take him to the animal shelter tomorrow. B is not sure if he can go with her, but he's willing to.\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "with torch.no_grad():\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.10.14"
  },
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}