tutorial/notebook/tutorial_qlora.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tutorial on PEFT methods + Quantization strategies, especially QLoRA, to do llms fine-tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial is based on the previous one [here](https://medium.com/@newhardwarefound/qlora-with-llama-2-ca1b4bcf26f0), which is also built upon this [github repo](https://github.com/taprosoft/llm_finetuning/tree/efa6df245fee4faf27206d84802d8f58d4b6e77d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial runs successfully with **torch==2.1.2**, **transformers==4.36.2**, **bitsandbytes==0.41.3**, **peft==0.7.1**, **accelerate==0.21.0**,  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step1.1 load the bfp16 model for full-power inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33m1452019841\u001b[0m (\u001b[33mstrivin\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "# 1. set the account token as an environment variable \"HUGGING_FACE_HUB_TOKEN\", \n",
    "#   which should be accessible to meta-llama/ in the huggingface hub inside the .env\n",
    "# 2. since we have downloaded the model ckpt already, so here we put the local meta-llama models root in .env too\n",
    "load_dotenv(find_dotenv())\n",
    "\n",
    "# set the visible gpus\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3' \n",
    "\n",
    "# set the wandb related environment variables\n",
    "os.environ['WANDB_PROJECT'] = 'tutorial on qlora'\n",
    "os.environ['WANDB_LOG_MODEL'] = 'checkpoint'\n",
    "os.environ['WANDB_NOTEBOOK_NAME'] = './tutorial_qlora.ipynb'\n",
    "\n",
    "# login to wandb\n",
    "import wandb\n",
    "# here you need to login to your wandb account at https://wandb.ai/site, \n",
    "# and get your api-key at https://wandb.ai/authorize\n",
    "wandb.login() \n",
    "\n",
    "from transformers import logging\n",
    "logging.set_verbosity_error()\n",
    "\n",
    "import torch\n",
    "import transformers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "## set the pretrained model root and model name\n",
    "\n",
    "# Llama2-7b\n",
    "model_root = os.getenv(\"LOCAL_LLAMA_MODEL_ROOT\")\n",
    "model_name = \"Llama-2-7b-chat\"\n",
    "\n",
    "# Mistral-7b\n",
    "# model_root = os.getenv(\"LOCAL_MISTRAL_MODEL_ROOT\")\n",
    "# model_name = \"Mistral-7B-v0.1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def get_model_params(model, trainable=False):\n",
    "    model_params = sum([p.numel() for p in model.parameters() if p.requires_grad == True or not trainable])\n",
    "    \n",
    "    b_size = model_params // 1000**3\n",
    "    m_size= (model_params % 1000**3) // 1000**2\n",
    "    \n",
    "    return f\"{b_size}B {m_size}M\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def get_mem_occupancy(model=None):\n",
    "    import torch.cuda as cuda\n",
    "    if model == None: # all of the memory occupancy\n",
    "        mem_occ = sum([\n",
    "            cuda.memory_allocated(i)\n",
    "            for i in range(cuda.device_count())\n",
    "        ])\n",
    "    else: # the memory occupancy for one particular model\n",
    "        mem_occ = model.get_memory_footprint()\n",
    "    \n",
    "    gb_size = mem_occ // 1024**3\n",
    "    mb_size = (mem_occ % 1024**3) // 1024**2\n",
    "    \n",
    "    return f\"{gb_size}GB {mb_size}MB\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def load_quant_model(model_path, mode=8, device_map='auto'):\n",
    "    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "    \n",
    "    kwargs = dict(device_map=device_map)\n",
    "    \n",
    "    # config the quantization with bitsandbytes\n",
    "    if mode == 8: # int 8bit quantization\n",
    "        kwargs['quantization_config'] = BitsAndBytesConfig(\n",
    "            load_in_8bit=True,\n",
    "            llm_int8_threshold=3.0 # Any hidden states value that is above this threshold will be considered an 'outlier' and the operation on those values will be done in fp16\n",
    "        )\n",
    "    elif mode == 4: # int 4bit quantization\n",
    "        kwargs['quantization_config'] = BitsAndBytesConfig(\n",
    "            load_in_4bit=True,\n",
    "            bnb_4bit_use_double_quant=True,\n",
    "            bnb_4bit_quant_type=\"nf4\",\n",
    "            bnb_4bit_compute_dtype=torch.float16\n",
    "        )\n",
    "    elif mode == 16: # no quantization\n",
    "        kwargs['torch_dtype'] = torch.bfloat16\n",
    "        \n",
    "    # load the tokenizer\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n",
    "    # set padding tokens to allow batch inference\n",
    "    tokenizer.pad_token_id = tokenizer.unk_token_id # to be different from the eos token\n",
    "    tokenizer.padding_side = \"left\" # for text generation task\n",
    "    \n",
    "    # load the model\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_path, trust_remote_code=True, **kwargs\n",
    "    )\n",
    "    model.eval()\n",
    "    \n",
    "    return model, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def show_model_tokenizer_info(model, tokenizer):\n",
    "    print(f\"Some meta information about the tokenizer:\\n{tokenizer}\\n\")\n",
    "    \n",
    "    print(f\"The model structure is as follows:\\n{model}\\n\")\n",
    "    print(f\"And the number of model's parameters is: {get_model_params(model)}\")\n",
    "    print(f\"And the memory occupied by the model is: {get_mem_occupancy(model)}, where the total footprint is {get_mem_occupancy()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-12-30 15:01:28.264852: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-12-30 15:01:29.099967: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9bd2698f020f4f81a4e8d8c97740818e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# the bf16 model has the biggest size of parameters and memory footprint\n",
    "model, tokenizer = load_quant_model(\n",
    "    os.path.join(model_root, model_name),\n",
    "    mode=16,\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Some meta information about the tokenizer:\n",
      "LlamaTokenizerFast(name_or_path='/data1/model/llama2/meta-llama/Llama-2-7b-chat', vocab_size=32000, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '<unk>'}, clean_up_tokenization_spaces=False),  added_tokens_decoder={\n",
      "\t0: AddedToken(\"<unk>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t1: AddedToken(\"<s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t2: AddedToken(\"</s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "}\n",
      "\n",
      "The model structure is as follows:\n",
      "LlamaForCausalLM(\n",
      "  (model): LlamaModel(\n",
      "    (embed_tokens): Embedding(32000, 4096)\n",
      "    (layers): ModuleList(\n",
      "      (0-31): 32 x LlamaDecoderLayer(\n",
      "        (self_attn): LlamaSdpaAttention(\n",
      "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
      "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
      "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
      "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
      "          (rotary_emb): LlamaRotaryEmbedding()\n",
      "        )\n",
      "        (mlp): LlamaMLP(\n",
      "          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)\n",
      "          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)\n",
      "          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)\n",
      "          (act_fn): SiLU()\n",
      "        )\n",
      "        (input_layernorm): LlamaRMSNorm()\n",
      "        (post_attention_layernorm): LlamaRMSNorm()\n",
      "      )\n",
      "    )\n",
      "    (norm): LlamaRMSNorm()\n",
      "  )\n",
      "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
      ")\n",
      "\n",
      "And the number of model's parameters is: 6B 738M\n",
      "And the memory occupied by the model is: 12GB 596MB, where the total footprint is 12GB 596MB\n"
     ]
    }
   ],
   "source": [
    "# show bfp16 model info\n",
    "show_model_tokenizer_info(model, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def text_completion(model, tokenizer, input_text, stream=False):\n",
    "    from transformers import TextStreamer\n",
    "    \n",
    "    inputs = tokenizer(input_text, return_tensors=\"pt\", padding=True).to('cuda')\n",
    "    outputs = model.generate(**inputs, \n",
    "                    do_sample=True,\n",
    "                    top_p=0.9,\n",
    "                    temperature=1e-4,\n",
    "                    max_new_tokens=100,\n",
    "                    streamer=TextStreamer(tokenizer, skip_prompt=True) if stream else None\n",
    "                )\n",
    "    \n",
    "    output_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
    "    \n",
    "    return output_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def stream_chat(model, tokenizer):\n",
    "    default_text = \"<s>[INST] <<SYS>> You are a helpful assistant. <</SYS>>\\\n",
    "    Extract the place names from the given sentence. [\\INST]\\n\\\n",
    "    The capital of the United States is Washington D.C.\"\n",
    "\n",
    "    while True:\n",
    "        input_text = input(\"Enter your prompt ('quit' to exit): \")\n",
    "        if input_text == \"quit\": break\n",
    "        elif input_text == \"\": input_text = default_text\n",
    "        \n",
    "        print(\"=\"*50)\n",
    "        print(f\"=> Prompt from the user:\\n{input_text}\")\n",
    "        \n",
    "        print(f\"=> Generated response:\")\n",
    "        text_completion(model, tokenizer, input_text, stream=True)\n",
    "        \n",
    "        print(\"=\"*50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "=> Prompt from the user:\n",
      "<s>[INST] <<SYS>> You are a helpful assistant. <</SYS>>    Extract the place names from the given sentence. [\\INST]\n",
      "    The capital of the United States is Washington D.C.\n",
      "=> Generated response:\n",
      "and the capital of France is Paris.\n",
      "\n",
      "Can you please extract the place names from the sentence?\n",
      "\n",
      "I need the place names: Washington, D.C. and Paris.\n",
      "\n",
      "Thank you!\n",
      "\n",
      "]]  Sure, I'd be happy to help! The place names in the sentence are:\n",
      "\n",
      "* Washington\n",
      "* D.C.\n",
      "* Paris</s>\n",
      "==================================================\n",
      "==================================================\n",
      "=> Prompt from the user:\n",
      "could you explain the world war II ?\n",
      "=> Generated response:\n",
      "\n",
      "\n",
      "World War II, also known as the Second World War, was a global conflict that lasted from 1939 to 1945. It was the deadliest war in history, with an estimated 50 to 80 million fatalities, including military personnel, civilians, and prisoners of war. The war was fought between two main alliances: the Allies, which consisted of the United States, the United Kingdom, and the Soviet Union\n",
      "==================================================\n",
      "==================================================\n",
      "=> Prompt from the user:\n",
      "can you list a few of research fileds for machine learning?\n",
      "=> Generated response:\n",
      "\n",
      "\n",
      "Machine learning is a subfield of artificial intelligence (AI) that involves the use of algorithms and statistical models to enable machines to learn from data, make decisions, and improve their performance on a specific task over time. Here are some of the most popular research fields in machine learning:\n",
      "1. Supervised Learning: This involves training a machine learning model on labeled data to make predictions on new, unseen data. Research areas include image classification, speech recognition, and natural language processing\n",
      "==================================================\n",
      "==================================================\n",
      "=> Prompt from the user:\n",
      "thank you\n",
      "=> Generated response:\n",
      "for your interest in our company. We are a leading provider of innovative solutions for the automotive industry, and we are committed to delivering high-quality products and services to our customers.\n",
      "However, we are unable to provide you with the information you have requested at this time. Our company policy is to not disclose any information that could potentially harm our business or compromise our relationships with our customers.\n",
      "We understand that this may be disappointing, but we hope you can understand\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "# as you can see, the bf16 model runs good for text completion tasks\n",
    "stream_chat(model, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step1.2 load the int8 quantized model to save the memory footprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2c4b4876c07401190ef094e257c3b04",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model_int8, tokenizer = load_quant_model(\n",
    "    os.path.join(model_root, model_name),\n",
    "    mode=8,\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Some meta information about the tokenizer:\n",
      "LlamaTokenizerFast(name_or_path='/data1/model/llama2/meta-llama/Llama-2-7b-chat', vocab_size=32000, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '<unk>'}, clean_up_tokenization_spaces=False),  added_tokens_decoder={\n",
      "\t0: AddedToken(\"<unk>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t1: AddedToken(\"<s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t2: AddedToken(\"</s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "}\n",
      "\n",
      "The model structure is as follows:\n",
      "LlamaForCausalLM(\n",
      "  (model): LlamaModel(\n",
      "    (embed_tokens): Embedding(32000, 4096)\n",
      "    (layers): ModuleList(\n",
      "      (0-31): 32 x LlamaDecoderLayer(\n",
      "        (self_attn): LlamaSdpaAttention(\n",
      "          (q_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
      "          (k_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
      "          (v_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
      "          (o_proj): Linear8bitLt(in_features=4096, out_features=4096, bias=False)\n",
      "          (rotary_emb): LlamaRotaryEmbedding()\n",
      "        )\n",
      "        (mlp): LlamaMLP(\n",
      "          (gate_proj): Linear8bitLt(in_features=4096, out_features=11008, bias=False)\n",
      "          (up_proj): Linear8bitLt(in_features=4096, out_features=11008, bias=False)\n",
      "          (down_proj): Linear8bitLt(in_features=11008, out_features=4096, bias=False)\n",
      "          (act_fn): SiLU()\n",
      "        )\n",
      "        (input_layernorm): LlamaRMSNorm()\n",
      "        (post_attention_layernorm): LlamaRMSNorm()\n",
      "      )\n",
      "    )\n",
      "    (norm): LlamaRMSNorm()\n",
      "  )\n",
      "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
      ")\n",
      "\n",
      "And the number of model's parameters is: 6B 738M\n",
      "And the memory occupied by the model is: 6GB 564MB, where the total footprint is 19GB 238MB\n"
     ]
    }
   ],
   "source": [
    "# show int8 model info\n",
    "# as we can see, compared to 12GB footprint for the bfp16 model, \n",
    "# int8 model only consumes half of them\n",
    "show_model_tokenizer_info(model_int8, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "=> Prompt from the user:\n",
      "<s>[INST] <<SYS>> You are a helpful assistant. <</SYS>>    Extract the place names from the given sentence. [\\INST]\n",
      "    The capital of the United States is Washington D.C.\n",
      "=> Generated response:\n",
      "and the capital of France is Paris.\n",
      "\n",
      "Can you please extract the place names from the sentence?\n",
      "\n",
      "I need the place names: Washington, D.C. and Paris.\n",
      "\n",
      "Thank you!\n",
      "\n",
      "]]  Sure, I'd be happy to help! The place names in the sentence are:\n",
      "\n",
      "* Washington\n",
      "* D.C.\n",
      "* Paris\n",
      "\n",
      "I hope this helps! Let me know if you have any other questions.</s>\n",
      "==================================================\n",
      "==================================================\n",
      "=> Prompt from the user:\n",
      "could you please explain the World War II?\n",
      "=> Generated response:\n",
      "\n",
      "\n",
      "World War II was a global conflict that lasted from 1939 to 1945. It was the deadliest war in history, with an estimated 50 to 80 million fatalities, most of whom were civilians. The war was fought between two main alliances: the Allies, which consisted of the United States, the United Kingdom, and the Soviet Union, among others; and the Axis powers, which included Germany,\n",
      "==================================================\n",
      "==================================================\n",
      "=> Prompt from the user:\n",
      "could you list a few research fields for machine learning?\n",
      "=> Generated response:\n",
      "\n",
      "\n",
      "Machine learning is a broad field that encompasses a wide range of research areas, including:\n",
      "\n",
      "1. Supervised learning: This involves training a machine learning algorithm to predict a target variable based on input features. Examples of supervised learning tasks include image classification, speech recognition, and sentiment analysis.\n",
      "2. Unsupervised learning: This involves training a machine learning algorithm to discover patterns or structure in data without any prior knowledge of the target variable. Examples of un\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "# as you can see, the int8 model runs almost the same as the bfp16 one\n",
    "# no clear drop in performance\n",
    "stream_chat(model_int8, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step1.3 load the int4 quantized model to further save the memory footprint and be fine-tuned later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-12-31 02:15:17.777843: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-12-31 02:15:18.654959: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2cb3cb7354f94afb982c458f7764904b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# the bf16 model has the biggest size of parameters and memory footprint\n",
    "model_int4, tokenizer = load_quant_model(\n",
    "    os.path.join(model_root, model_name),\n",
    "    mode=4,\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Some meta information about the tokenizer:\n",
      "LlamaTokenizerFast(name_or_path='/data1/model/llama2/meta-llama/Llama-2-7b-chat', vocab_size=32000, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '<unk>'}, clean_up_tokenization_spaces=False),  added_tokens_decoder={\n",
      "\t0: AddedToken(\"<unk>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t1: AddedToken(\"<s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "\t2: AddedToken(\"</s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
      "}\n",
      "\n",
      "The model structure is as follows:\n",
      "LlamaForCausalLM(\n",
      "  (model): LlamaModel(\n",
      "    (embed_tokens): Embedding(32000, 4096)\n",
      "    (layers): ModuleList(\n",
      "      (0-31): 32 x LlamaDecoderLayer(\n",
      "        (self_attn): LlamaSdpaAttention(\n",
      "          (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
      "          (k_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
      "          (v_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
      "          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
      "          (rotary_emb): LlamaRotaryEmbedding()\n",
      "        )\n",
      "        (mlp): LlamaMLP(\n",
      "          (gate_proj): Linear4bit(in_features=4096, out_features=11008, bias=False)\n",
      "          (up_proj): Linear4bit(in_features=4096, out_features=11008, bias=False)\n",
      "          (down_proj): Linear4bit(in_features=11008, out_features=4096, bias=False)\n",
      "          (act_fn): SiLU()\n",
      "        )\n",
      "        (input_layernorm): LlamaRMSNorm()\n",
      "        (post_attention_layernorm): LlamaRMSNorm()\n",
      "      )\n",
      "    )\n",
      "    (norm): LlamaRMSNorm()\n",
      "  )\n",
      "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
      ")\n",
      "\n",
      "And the number of model's parameters is: 3B 500M\n",
      "And the memory occupied by the model is: 3GB 548MB, where the total footprint is 22GB 906MB\n"
     ]
    }
   ],
   "source": [
    "# we further reduce the footprint to a quarter of the bfp16 model,\n",
    "# which only occupy 3GB of gpu memory \n",
    "show_model_tokenizer_info(model_int4, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "=> Prompt from the user:\n",
      "<s>[INST] <<SYS>> You are a helpful assistant. <</SYS>>    Extract the place names from the given sentence. [\\INST]\n",
      "    The capital of the United States is Washington D.C.\n",
      "=> Generated response:\n",
      "and the capital of France is Paris.\n",
      "\n",
      "Can you please extract the place names from the sentence?\n",
      "\n",
      "I can certainly try! Here are the place names mentioned in the sentence:\n",
      "\n",
      "* Washington D.C.\n",
      "* Paris</s>\n",
      "==================================================\n",
      "==================================================\n",
      "=> Prompt from the user:\n",
      "could me explain the world war II ? \n",
      "=> Generated response:\n",
      "\n",
      "Unterscheidung of the world war II?\n",
      "\n",
      "World War II was a global conflict that lasted from 1939 to 1945. It was the deadliest war in history, with an estimated 50 to 80 million fatalities, including military personnel, civilians, and prisoners of war. The war was fought between two main alliances: the Allies and the Axis.\n",
      "The Allies were a group of countries that included the\n",
      "==================================================\n",
      "==================================================\n",
      "=> Prompt from the user:\n",
      "can you list a few of research fileds for machine learning?\n",
      "=> Generated response:\n",
      "\n",
      "nobody knows the answer to this question.\n",
      "\n",
      "Some of the most popular research fields in machine learning include:\n",
      "\n",
      "1. Natural Language Processing (NLP): This field focuses on developing algorithms and models that can process, analyze, and generate human language data, such as text, speech, and sentiment analysis.\n",
      "2. Computer Vision: This field involves developing algorithms and models that can interpret and analyze visual data from images and videos, such as object detection, image segmentation\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "# as you can see, the output does not bad with a little bit out of control, \n",
    "# as a result of trading memory footprint for performance\n",
    "stream_chat(model_int4, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step2.1 load the raw instruction tuning dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The source IFT dataset we use is a cleaned version of Alpaca, which can be found [here](https://huggingface.co/datasets/yahma/alpaca-cleaned) in the huggingface hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['instruction', 'input', 'output'],\n",
       "        num_rows: 51760\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "data_path = './data/alpaca_data_cleaned.json' # 51k instructions\n",
    "\n",
    "dataset = load_dataset(\"json\", data_files=data_path) # only train split\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'instruction': 'Give three tips for staying healthy.',\n",
       " 'input': '',\n",
       " 'output': '1. Eat a balanced and nutritious diet: Make sure your meals are inclusive of a variety of fruits and vegetables, lean protein, whole grains, and healthy fats. This helps to provide your body with the essential nutrients to function at its best and can help prevent chronic diseases.\\n\\n2. Engage in regular physical activity: Exercise is crucial for maintaining strong bones, muscles, and cardiovascular health. Aim for at least 150 minutes of moderate aerobic exercise or 75 minutes of vigorous exercise each week.\\n\\n3. Get enough sleep: Getting enough quality sleep is crucial for physical and mental well-being. It helps to regulate mood, improve cognitive function, and supports healthy growth and immune function. Aim for 7-9 hours of sleep each night.'}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "dataset['train'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step2.2 split the dataset into train / val splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['instruction', 'input', 'output'],\n",
       "        num_rows: 46584\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['instruction', 'input', 'output'],\n",
       "        num_rows: 5176\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_ratio = 0.1\n",
    "\n",
    "splitted_dataset = dataset[\"train\"].train_test_split(\n",
    "                test_size=val_ratio, shuffle=True, seed=42\n",
    "            )\n",
    "splitted_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step2.3 define the prompt templates to transfer sample to prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [],
   "source": [
    "alpaca_templates = {\n",
    "    \"description\": \"Template used by Alpaca-LoRA.\",\n",
    "    \"prompt_input\": \"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\\n\\n### Instruction:\\n{instruction}\\n\\n### Input:\\n{input}\\n\\n### Response:\\n\",\n",
    "    \"prompt_no_input\": \"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n### Instruction:\\n{instruction}\\n\\n### Response:\\n\",\n",
    "    \"response_split\": \"### Response:\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n",
      "\n",
      "### Instruction:\n",
      "{instruction}\n",
      "\n",
      "### Input:\n",
      "{input}\n",
      "\n",
      "### Response:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "alpaca_template_with_input = alpaca_templates[\"prompt_input\"]\n",
    "print(alpaca_template_with_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
      "\n",
      "### Instruction:\n",
      "{instruction}\n",
      "\n",
      "### Response:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "alpaca_template_without_input = alpaca_templates[\"prompt_no_input\"]\n",
    "print(alpaca_template_without_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def get_response_from_prompt(prompt: str) -> str:\n",
    "    \"\"\"retrieve the response part in one full prompt string\"\"\"\n",
    "    try:\n",
    "        response = prompt.split(alpaca_templates[\"response_split\"])[1].strip()\n",
    "    except (KeyError, IndexError):\n",
    "        response = prompt\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step2.4 generate and tokenize the final prompts with labels to build the actual model-input dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import PreTrainedTokenizer\n",
    "\n",
    "def generate_prompt(sample: dict) -> str:\n",
    "    \"\"\"generate the prompt from the dict sample with the template, and tokenize it with the tokenizer\n",
    "    \"\"\"\n",
    "    instruction, input, output = sample['instruction'], sample['input'], sample['output']\n",
    "    \n",
    "    # two different prompt templates\n",
    "    if input: \n",
    "        prompt = alpaca_template_with_input.format(instruction=instruction, input=input)\n",
    "    else:\n",
    "        prompt = alpaca_template_without_input.format(instruction=instruction)\n",
    "        \n",
    "    if output: # append the label\n",
    "        prompt = f\"{prompt}{output}\"\n",
    "        \n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def tokenize_prompt(prompt: str, \n",
    "                    tokenizer: PreTrainedTokenizer, \n",
    "                    max_length=512,\n",
    "                    add_eos_token=True,\n",
    "                    ):\n",
    "    inputs = tokenizer(prompt, \n",
    "                max_length=max_length, \n",
    "                truncation=True, \n",
    "                padding=False, \n",
    "                return_tensors=None # just return a list-like \"input_ids\", \"attention_mask\", ..\n",
    "            )\n",
    "    if (\n",
    "        inputs['input_ids'][-1] != tokenizer.eos_token_id and\n",
    "        len(inputs['input_ids']) < max_length and\n",
    "        add_eos_token\n",
    "    ):\n",
    "        inputs['input_ids'].append(tokenizer.eos_token_id)\n",
    "        inputs['attention_mask'].append(1)\n",
    "    \n",
    "    inputs['labels'] = inputs['input_ids'].copy() # unsupervised learning\n",
    "    \n",
    "    return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def gen_and_tokenize_prompt(sample: dict, tokenizer: PreTrainedTokenizer):\n",
    "    prompt = generate_prompt(sample)\n",
    "    inputs = tokenize_prompt(prompt, tokenizer)\n",
    "    return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c2a945e8510429ca18b90746dcadae9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/5176 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from functools import partial\n",
    "\n",
    "# map each sample to get the actual model-input dataset for both train / val splits\n",
    "prompt_map_func = partial(gen_and_tokenize_prompt, tokenizer=tokenizer)\n",
    "\n",
    "train_dataset = splitted_dataset['train'].shuffle(seed=42).map(prompt_map_func)\n",
    "val_dataset = splitted_dataset['test'].map(prompt_map_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['instruction', 'input', 'output', 'input_ids', 'attention_mask', 'labels'],\n",
      "    num_rows: 46584\n",
      "})\n",
      "Dataset({\n",
      "    features: ['instruction', 'input', 'output', 'input_ids', 'attention_mask', 'labels'],\n",
      "    num_rows: 5176\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "print(train_dataset)\n",
    "print(val_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'instruction': 'Rearrange the following sentence to make the sentence more interesting.',\n",
       " 'input': 'She left the party early',\n",
       " 'output': 'Early, she left the party.',\n",
       " 'input_ids': [1,\n",
       "  13866,\n",
       "  338,\n",
       "  385,\n",
       "  15278,\n",
       "  393,\n",
       "  16612,\n",
       "  263,\n",
       "  3414,\n",
       "  29892,\n",
       "  3300,\n",
       "  2859,\n",
       "  411,\n",
       "  385,\n",
       "  1881,\n",
       "  393,\n",
       "  8128,\n",
       "  4340,\n",
       "  3030,\n",
       "  29889,\n",
       "  14350,\n",
       "  263,\n",
       "  2933,\n",
       "  393,\n",
       "  7128,\n",
       "  2486,\n",
       "  1614,\n",
       "  2167,\n",
       "  278,\n",
       "  2009,\n",
       "  29889,\n",
       "  13,\n",
       "  13,\n",
       "  2277,\n",
       "  29937,\n",
       "  2799,\n",
       "  4080,\n",
       "  29901,\n",
       "  13,\n",
       "  29934,\n",
       "  799,\n",
       "  3881,\n",
       "  278,\n",
       "  1494,\n",
       "  10541,\n",
       "  304,\n",
       "  1207,\n",
       "  278,\n",
       "  10541,\n",
       "  901,\n",
       "  8031,\n",
       "  29889,\n",
       "  13,\n",
       "  13,\n",
       "  2277,\n",
       "  29937,\n",
       "  10567,\n",
       "  29901,\n",
       "  13,\n",
       "  13468,\n",
       "  2175,\n",
       "  278,\n",
       "  6263,\n",
       "  4688,\n",
       "  13,\n",
       "  13,\n",
       "  2277,\n",
       "  29937,\n",
       "  13291,\n",
       "  29901,\n",
       "  13,\n",
       "  29923,\n",
       "  279,\n",
       "  368,\n",
       "  29892,\n",
       "  1183,\n",
       "  2175,\n",
       "  278,\n",
       "  6263,\n",
       "  29889,\n",
       "  2],\n",
       " 'attention_mask': [1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1],\n",
       " 'labels': [1,\n",
       "  13866,\n",
       "  338,\n",
       "  385,\n",
       "  15278,\n",
       "  393,\n",
       "  16612,\n",
       "  263,\n",
       "  3414,\n",
       "  29892,\n",
       "  3300,\n",
       "  2859,\n",
       "  411,\n",
       "  385,\n",
       "  1881,\n",
       "  393,\n",
       "  8128,\n",
       "  4340,\n",
       "  3030,\n",
       "  29889,\n",
       "  14350,\n",
       "  263,\n",
       "  2933,\n",
       "  393,\n",
       "  7128,\n",
       "  2486,\n",
       "  1614,\n",
       "  2167,\n",
       "  278,\n",
       "  2009,\n",
       "  29889,\n",
       "  13,\n",
       "  13,\n",
       "  2277,\n",
       "  29937,\n",
       "  2799,\n",
       "  4080,\n",
       "  29901,\n",
       "  13,\n",
       "  29934,\n",
       "  799,\n",
       "  3881,\n",
       "  278,\n",
       "  1494,\n",
       "  10541,\n",
       "  304,\n",
       "  1207,\n",
       "  278,\n",
       "  10541,\n",
       "  901,\n",
       "  8031,\n",
       "  29889,\n",
       "  13,\n",
       "  13,\n",
       "  2277,\n",
       "  29937,\n",
       "  10567,\n",
       "  29901,\n",
       "  13,\n",
       "  13468,\n",
       "  2175,\n",
       "  278,\n",
       "  6263,\n",
       "  4688,\n",
       "  13,\n",
       "  13,\n",
       "  2277,\n",
       "  29937,\n",
       "  13291,\n",
       "  29901,\n",
       "  13,\n",
       "  29923,\n",
       "  279,\n",
       "  368,\n",
       "  29892,\n",
       "  1183,\n",
       "  2175,\n",
       "  278,\n",
       "  6263,\n",
       "  29889,\n",
       "  2]}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step2.5 define the data collator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "2"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForSeq2Seq\n",
    "data_collator = DataCollatorForSeq2Seq(\n",
    "    tokenizer,\n",
    "    pad_to_multiple_of=8,\n",
    "    return_tensors=\"pt\", \n",
    "    padding=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step3.1 prepare some helper functions before training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the helper functions below are all based on the previous ones defined [here](https://github.com/artidoro/qlora/blob/main/qlora.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def print_trainable_parameters(model, use_4bit=True):\n",
    "    \"\"\"\n",
    "    Prints the number of trainable parameters in the model.\n",
    "    \"\"\"\n",
    "    trainable_params = 0\n",
    "    all_param = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        num_params = param.numel()\n",
    "        # if using DS Zero 3 and the weights are initialized empty\n",
    "        if num_params == 0 and hasattr(param, \"ds_numel\"):\n",
    "            num_params = param.ds_numel\n",
    "\n",
    "        all_param += num_params\n",
    "        if param.requires_grad:\n",
    "            trainable_params += num_params\n",
    "    if use_4bit: # a 4bit int counts a half byte\n",
    "        trainable_params /= 2\n",
    "        \n",
    "        \n",
    "    print(\n",
    "        f\"all params: {all_param // 1000**3:.1f}B {(all_param % 1000**3) // 1000**2:.1f}M\\n\" + \n",
    "        f\"trainable params: {trainable_params // 1000**3:.1f}B {(trainable_params % 1000**3) // 1000**2:.1f}M\\n\" + \n",
    "        f\"trainable percent: {100 * trainable_params / all_param:.2f}%\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def find_all_linear_names(model):\n",
    "    from bitsandbytes.nn import Linear4bit\n",
    "    \n",
    "    lora_module_names = set()\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, Linear4bit):\n",
    "            names = name.split(\".\")\n",
    "            lora_module_names.add(names[0] if len(names) == 1 else names[-1])\n",
    "\n",
    "    if \"lm_head\" in lora_module_names:  # needed for 16-bit\n",
    "        lora_module_names.remove(\"lm_head\")\n",
    "        \n",
    "    return list(lora_module_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def preprocess_modules_dtype(model, bf16=True):\n",
    "    from peft.tuners.lora import LoraLayer\n",
    "    \n",
    "    for name, module in model.named_modules():\n",
    "        # upcast lora layer dtype to bfloat16\n",
    "        if isinstance(module, LoraLayer) and bf16:\n",
    "            module = module.to(torch.bfloat16)\n",
    "        # upcast layer norm dtype to float32\n",
    "        if 'norm' in name:\n",
    "            module = module.to(torch.float32)\n",
    "        # upcast lm_head and word_embed dtype from float32 to bfloat16\n",
    "        if (\n",
    "            (\"lm_head\" in name or \"embed_tokens\" in name) and\n",
    "            hasattr(module, \"weight\") and\n",
    "            bf16 and module.weight.dtype == torch.float32\n",
    "        ): module = module.to(torch.bfloat16)\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step3.2 create the qlora model with peft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def create_qlora_model(model, grad_ckpt=True, bf16=True):\n",
    "    from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training\n",
    "    \n",
    "    # prepare int4 model for training\n",
    "    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=grad_ckpt)\n",
    "    if grad_ckpt: model.gradient_checkpointing_enable()\n",
    "    \n",
    "    # find lora target modules\n",
    "    target_modules = find_all_linear_names(model)\n",
    "    print(f\"Found {len(target_modules)} modules to quantize:\\n {target_modules}\")\n",
    "    \n",
    "    # define lora config\n",
    "    lora_config = LoraConfig(\n",
    "        r=64, # rank\n",
    "        lora_alpha=16,\n",
    "        target_modules=target_modules,\n",
    "        lora_dropout=0.1,\n",
    "        bias='none',\n",
    "        task_type=TaskType.CAUSAL_LM\n",
    "    )\n",
    "    \n",
    "    # get the peft model from the lora config\n",
    "    model = get_peft_model(model, lora_config)\n",
    "    \n",
    "    # pre-process with some of the modules' dtype\n",
    "    model = preprocess_modules_dtype(model, bf16)\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 7 modules to quantize:\n",
      " ['gate_proj', 'k_proj', 'v_proj', 'o_proj', 'down_proj', 'q_proj', 'up_proj']\n"
     ]
    }
   ],
   "source": [
    "qlora_model = create_qlora_model(model_int4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "all params: 3.0B 660.0M\n",
      "trainable params: 0.0B 79.0M\n",
      "trainable percent: 2.18%\n"
     ]
    }
   ],
   "source": [
    "# print the model's (trainable) parameters \n",
    "print_trainable_parameters(qlora_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): LlamaForCausalLM(\n",
       "      (model): LlamaModel(\n",
       "        (embed_tokens): Embedding(32000, 4096)\n",
       "        (layers): ModuleList(\n",
       "          (0-31): 32 x LlamaDecoderLayer(\n",
       "            (self_attn): LlamaSdpaAttention(\n",
       "              (q_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (v_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (rotary_emb): LlamaRotaryEmbedding()\n",
       "            )\n",
       "            (mlp): LlamaMLP(\n",
       "              (gate_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=11008, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=11008, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (up_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=11008, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=11008, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (down_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=11008, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=11008, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (act_fn): SiLU()\n",
       "            )\n",
       "            (input_layernorm): LlamaRMSNorm()\n",
       "            (post_attention_layernorm): LlamaRMSNorm()\n",
       "          )\n",
       "        )\n",
       "        (norm): LlamaRMSNorm()\n",
       "      )\n",
       "      (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qlora_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step3.3 define the training arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainingArguments(\n",
       "_n_gpu=4,\n",
       "adafactor=False,\n",
       "adam_beta1=0.9,\n",
       "adam_beta2=0.999,\n",
       "adam_epsilon=1e-08,\n",
       "auto_find_batch_size=False,\n",
       "bf16=False,\n",
       "bf16_full_eval=False,\n",
       "data_seed=None,\n",
       "dataloader_drop_last=False,\n",
       "dataloader_num_workers=0,\n",
       "dataloader_persistent_workers=False,\n",
       "dataloader_pin_memory=True,\n",
       "ddp_backend=None,\n",
       "ddp_broadcast_buffers=None,\n",
       "ddp_bucket_cap_mb=None,\n",
       "ddp_find_unused_parameters=None,\n",
       "ddp_timeout=1800,\n",
       "debug=[],\n",
       "deepspeed=None,\n",
       "disable_tqdm=True,\n",
       "dispatch_batches=None,\n",
       "do_eval=True,\n",
       "do_predict=False,\n",
       "do_train=True,\n",
       "eval_accumulation_steps=None,\n",
       "eval_delay=0,\n",
       "eval_steps=100,\n",
       "evaluation_strategy=steps,\n",
       "fp16=True,\n",
       "fp16_backend=auto,\n",
       "fp16_full_eval=False,\n",
       "fp16_opt_level=O1,\n",
       "fsdp=[],\n",
       "fsdp_config={'min_num_params': 0, 'xla': False, 'xla_fsdp_grad_ckpt': False},\n",
       "fsdp_min_num_params=0,\n",
       "fsdp_transformer_layer_cls_to_wrap=None,\n",
       "full_determinism=False,\n",
       "gradient_accumulation_steps=1,\n",
       "gradient_checkpointing=True,\n",
       "gradient_checkpointing_kwargs=None,\n",
       "greater_is_better=False,\n",
       "group_by_length=False,\n",
       "half_precision_backend=auto,\n",
       "hub_always_push=False,\n",
       "hub_model_id=None,\n",
       "hub_private_repo=False,\n",
       "hub_strategy=every_save,\n",
       "hub_token=<HUB_TOKEN>,\n",
       "ignore_data_skip=False,\n",
       "include_inputs_for_metrics=False,\n",
       "include_num_input_tokens_seen=False,\n",
       "include_tokens_per_second=False,\n",
       "jit_mode_eval=False,\n",
       "label_names=None,\n",
       "label_smoothing_factor=0.0,\n",
       "learning_rate=0.0002,\n",
       "length_column_name=length,\n",
       "load_best_model_at_end=True,\n",
       "local_rank=0,\n",
       "log_level=passive,\n",
       "log_level_replica=warning,\n",
       "log_on_each_node=True,\n",
       "logging_dir=./logs,\n",
       "logging_first_step=False,\n",
       "logging_nan_inf_filter=True,\n",
       "logging_steps=10,\n",
       "logging_strategy=steps,\n",
       "lr_scheduler_kwargs={},\n",
       "lr_scheduler_type=linear,\n",
       "max_grad_norm=1.0,\n",
       "max_steps=-1,\n",
       "metric_for_best_model=loss,\n",
       "mp_parameters=,\n",
       "neftune_noise_alpha=None,\n",
       "no_cuda=False,\n",
       "num_train_epochs=1,\n",
       "optim=paged_adamw_8bit,\n",
       "optim_args=None,\n",
       "output_dir=.,\n",
       "overwrite_output_dir=False,\n",
       "past_index=-1,\n",
       "per_device_eval_batch_size=1,\n",
       "per_device_train_batch_size=1,\n",
       "prediction_loss_only=False,\n",
       "push_to_hub=False,\n",
       "push_to_hub_model_id=None,\n",
       "push_to_hub_organization=None,\n",
       "push_to_hub_token=<PUSH_TO_HUB_TOKEN>,\n",
       "ray_scope=last,\n",
       "remove_unused_columns=False,\n",
       "report_to=['wandb'],\n",
       "resume_from_checkpoint=None,\n",
       "run_name=.,\n",
       "save_on_each_node=False,\n",
       "save_only_model=False,\n",
       "save_safetensors=True,\n",
       "save_steps=500,\n",
       "save_strategy=steps,\n",
       "save_total_limit=2,\n",
       "seed=42,\n",
       "skip_memory_metrics=True,\n",
       "split_batches=False,\n",
       "tf32=None,\n",
       "torch_compile=False,\n",
       "torch_compile_backend=None,\n",
       "torch_compile_mode=None,\n",
       "torchdynamo=None,\n",
       "tpu_metrics_debug=False,\n",
       "tpu_num_cores=None,\n",
       "use_cpu=False,\n",
       "use_ipex=False,\n",
       "use_legacy_prediction_loop=False,\n",
       "use_mps_device=False,\n",
       "warmup_ratio=0.0,\n",
       "warmup_steps=100,\n",
       "weight_decay=0.0,\n",
       ")"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "# define fine-tuning arguments\n",
    "output_dir = '.'\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    do_train=True,\n",
    "    # optimizer\n",
    "    learning_rate=2e-4,\n",
    "    optim=\"paged_adamw_8bit\", # use \"adamw_torch\" if not mode = 4,8\n",
    "    # this is critical since it determines the amount of memory consumption. \n",
    "    # the bigger per_device_train_batch_size is, the faster the training will finish\n",
    "    # the smaller the lower probability of an OOM error\n",
    "    # per_device_train_batch_size=1,\n",
    "    \n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    warmup_steps=100,\n",
    "    \n",
    "    num_train_epochs=1, # for fine-tuning\n",
    "    eval_steps=100, # evaluate on the validation set every 100 steps\n",
    "    save_steps=500, # save every 500 steps (default)\n",
    "    # to further save memory footprint\n",
    "    fp16=True,\n",
    "    gradient_checkpointing=True,\n",
    "    gradient_accumulation_steps=1, # default\n",
    "    # logging\n",
    "    logging_dir=os.path.join(output_dir, 'logs'),\n",
    "    logging_strategy='steps', # choose \"steps\" or \"epoch\"\n",
    "    logging_steps=10,\n",
    "    remove_unused_columns=False,\n",
    "    # others\n",
    "    report_to='wandb',\n",
    "    seed=42,\n",
    "    load_best_model_at_end=True,\n",
    "    save_total_limit=2,\n",
    "    evaluation_strategy=\"steps\",\n",
    ")\n",
    "training_args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step3.4 instantiate the trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<transformers.trainer.Trainer at 0x7fe4710d23a0>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "qlora_model.config.use_cache = False # before training, make sure to set `use_cache` to `False`\n",
    "\n",
    "# for multi-gpus env, the model's attribute `is_model_parallel` has to set to `True` \n",
    "# to avoid any unexpected behavior such as device placement mismatching\n",
    "trainer = Trainer(\n",
    "    model=qlora_model,\n",
    "    args=training_args,\n",
    "    \n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=val_dataset,\n",
    "    # data_collator=data_collator,\n",
    ")\n",
    "trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step3.5 train the model and save it when finished"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to save the memory,\n",
    "# we might delete the bfp16 model and the int8 model during training\n",
    "# and load them back after finishing\n",
    "# del model, model_int8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.16.1 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.8"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/data/hyp/llms/tutorials/wandb/run-20231231_021655-811f9g19</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/strivin/tutorial%20on%20qlora/runs/811f9g19' target=\"_blank\">autumn-voice-2</a></strong> to <a href='https://wandb.ai/strivin/tutorial%20on%20qlora' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/strivin/tutorial%20on%20qlora' target=\"_blank\">https://wandb.ai/strivin/tutorial%20on%20qlora</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/strivin/tutorial%20on%20qlora/runs/811f9g19' target=\"_blank\">https://wandb.ai/strivin/tutorial%20on%20qlora/runs/811f9g19</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "ValueError",
     "evalue": "expected sequence of length 426 at dim 1 (got 326)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[33], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# start to train\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/transformers/trainer.py:1537\u001b[0m, in \u001b[0;36mTrainer.train\u001b[0;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[1;32m   1535\u001b[0m         hf_hub_utils\u001b[38;5;241m.\u001b[39menable_progress_bars()\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1537\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1538\u001b[0m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1539\u001b[0m \u001b[43m        \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1540\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtrial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1541\u001b[0m \u001b[43m        \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1542\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/transformers/trainer.py:1821\u001b[0m, in \u001b[0;36mTrainer._inner_training_loop\u001b[0;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[0m\n\u001b[1;32m   1818\u001b[0m     rng_to_sync \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m   1820\u001b[0m step \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m\n\u001b[0;32m-> 1821\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m step, inputs \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(epoch_iterator):\n\u001b[1;32m   1822\u001b[0m     total_batched_samples \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m   1824\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39minclude_num_input_tokens_seen:\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/accelerate/data_loader.py:384\u001b[0m, in \u001b[0;36mDataLoaderShard.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    382\u001b[0m \u001b[38;5;66;03m# We iterate one batch ahead to check when we are at the end\u001b[39;00m\n\u001b[1;32m    383\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 384\u001b[0m     current_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mnext\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mdataloader_iter\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    385\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m:\n\u001b[1;32m    386\u001b[0m     \u001b[38;5;28;01myield\u001b[39;00m\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/torch/utils/data/dataloader.py:630\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    627\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    628\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[1;32m    629\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[0;32m--> 630\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    631\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    632\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    633\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    634\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/torch/utils/data/dataloader.py:674\u001b[0m, in \u001b[0;36m_SingleProcessDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    672\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_next_data\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    673\u001b[0m     index \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_next_index()  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[0;32m--> 674\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_dataset_fetcher\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfetch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindex\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[1;32m    675\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory:\n\u001b[1;32m    676\u001b[0m         data \u001b[38;5;241m=\u001b[39m _utils\u001b[38;5;241m.\u001b[39mpin_memory\u001b[38;5;241m.\u001b[39mpin_memory(data, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory_device)\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:54\u001b[0m, in \u001b[0;36m_MapDatasetFetcher.fetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     53\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n\u001b[0;32m---> 54\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollate_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/transformers/data/data_collator.py:70\u001b[0m, in \u001b[0;36mdefault_data_collator\u001b[0;34m(features, return_tensors)\u001b[0m\n\u001b[1;32m     64\u001b[0m \u001b[38;5;66;03m# In this function we'll make the assumption that all `features` in the batch\u001b[39;00m\n\u001b[1;32m     65\u001b[0m \u001b[38;5;66;03m# have the same attributes.\u001b[39;00m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;66;03m# So we will look at the first element as a proxy for what attributes exist\u001b[39;00m\n\u001b[1;32m     67\u001b[0m \u001b[38;5;66;03m# on the whole batch.\u001b[39;00m\n\u001b[1;32m     69\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m return_tensors \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpt\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 70\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch_default_data_collator\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfeatures\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     71\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m return_tensors \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtf\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m     72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m tf_default_data_collator(features)\n",
      "File \u001b[0;32m/data/hyp/anaconda3/envs/llms/lib/python3.8/site-packages/transformers/data/data_collator.py:136\u001b[0m, in \u001b[0;36mtorch_default_data_collator\u001b[0;34m(features)\u001b[0m\n\u001b[1;32m    134\u001b[0m             batch[k] \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(np\u001b[38;5;241m.\u001b[39mstack([f[k] \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m features]))\n\u001b[1;32m    135\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 136\u001b[0m             batch[k] \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[43mf\u001b[49m\u001b[43m[\u001b[49m\u001b[43mk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mf\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    138\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m batch\n",
      "\u001b[0;31mValueError\u001b[0m: expected sequence of length 426 at dim 1 (got 326)"
     ]
    }
   ],
   "source": [
    "# start to train\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if the training process is accidently paused\n",
    "# you can resume to train from the certain checkpoint, \n",
    "# which is generated automatically like `./checkpint-{save_steps*n}` \n",
    "# by setting `save_steps` in training arguments (default is 500)\n",
    "\n",
    "# trainer.train(resume_from_checkpoint='./checkpoints-500/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the final model ckpt\n",
    "save_dir = \"./model/llama2-7b-qlora-adapter-e27k\"\n",
    "trainer.save_model(save_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step4.1 report the training process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step4.2 evaluate on the validation set with the baselines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "if you deleted the bfp model and int8 model during fine-tuning the int4 model\n",
    "\n",
    "you should firstly load them back in step1.1 and step1.2 before runing step4.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### step4.3 evaluate on the public benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llms",
   "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.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}