Finetuning a LLM using LoRA Code Walkthrough .ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8c618e0",
   "metadata": {},
   "source": [
    "# Fine-tuning an LLM: the distributed edition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ca15607",
   "metadata": {},
   "source": [
    "This notebook was submitted in pursuance of the Deep Learning course, part of the MSDS Curriculum at University of San Francisco.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7ff42eb",
   "metadata": {},
   "source": [
    "**Group Members**\n",
    "- Maneel Reddy Karri\n",
    "- Gurusankar Gopalakrishnan\n",
    "- Devendra Govil"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f665bedf",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3a83976",
   "metadata": {},
   "source": [
    "## Introduction <a></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6e03446",
   "metadata": {},
   "source": [
    "**Large Language Models(LLMs)** refer to advanced artificial intelligence systems designed to process and generate human-like text. They are trained on vast amounts of text data and learn to understand and produce natural language in a variety of contexts. These belong to the family of auto-regressive transformers based models and have had a huge boon thanks to their remarkable performance on a range of Natural Language Processing tasks. There is no better way to exemplify the significance they hold in the Natural Language Processing domain than to look at ChatGPT, a large language model that has taken the world by storm. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0842454",
   "metadata": {},
   "source": [
    "ChatGPT is a large language model which is developed by OpenAI and based on the GPT 3.5 architecture and is coupled with a step of RLHF (reinforcement learning through human feedback)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b26f77c6",
   "metadata": {},
   "source": [
    "### Fine-tuning an LLM <a name='h3_1'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef643de0",
   "metadata": {},
   "source": [
    "A general purpose LLM like ChatGPT or Bard is trained on an enormous corpus of text data and has remarkable performance on general purpose NLP tasks. They provide very good base models that can help one get started with modelling NLP problems.  \n",
    "\n",
    "However, they still experience various issues. The data they are trained on may not be include stuff that is relevant to our use case. For example, if I am working in a finance company with proprietary data, our base model has never seen that (or possibly even similar) data, leading to significant performance degradation on many NLP tasks related to niche or specialized domains.Similarly, a general purpose LLM, which is built to support a variety of NLP tasks, is by design incentivized for general performance rather than for a specific task. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28ce06c4",
   "metadata": {},
   "source": [
    "Hence, finetuning is extremely important because it may help in getting significant performance boost on tasks of relevance. It may sometimes in fact be the only way to get acceptable performance on relevant tasks. Since training an LLM from scratch is prohibitively expensive, time consuming, requires a lot of expertise, and is resource cumbersome, fine-tuning a base LLM provides a very attractive option for organizations and individuals, to improve performance of LLMs on domain specific tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "880a5940",
   "metadata": {},
   "source": [
    "### Goals for this project"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b5cf81b",
   "metadata": {},
   "source": [
    "Our project entailed fine-tuning an LLM using the following bleeding edge techniques:\n",
    "- Different distributed training approaches like data, model and pipeline parallelism\n",
    "- Using parameter efficient fine-tuning techniques like LoRA (Low Rank Adapters)\n",
    "- Model inference using a distributed orchestration platform like rayserve"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c126b89d",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "438d0f89",
   "metadata": {},
   "source": [
    "## Data Preprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0acdd0ff",
   "metadata": {},
   "source": [
    "For finetuning our model we are required to take an instruction dataset that contains instruction, question, and answer pairs. These help us to transform our model to optimize for the purposes of question answering as well as for fine-tuning. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f91f150",
   "metadata": {},
   "source": [
    "We use the famous and reknowed [Alpaca dataset](https://huggingface.co/datasets/teknium/GPT4-LLM-Cleaned) that is available in [here](https://huggingface.co/datasets/teknium/GPT4-LLM-Cleaned) which has the benefit of being processed for the express purpose of finetuning an LLM. This is beneficial because it helps us in being able to devot more of our time and resources on tackling the challenging task of finetuning the LLM on the distributed architecture using the LoRA technique."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab7dc47",
   "metadata": {},
   "source": [
    "To start off, We first load the model from the Hugging Face library, initialize the tokenizer and add special tokens such as:\n",
    "- End of Sentence\n",
    "- Beginning of Sentence\n",
    "- Unknown Tokens\n",
    "   \n",
    "Adding these special tokens allows us to convert the text to generate embedding indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c8757e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import LlamaTokenizer, LlamaForCausalLM,LlamaConfig\n",
    "\n",
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "import torch\n",
    "model_name_or_path = \"openlm-research/open_llama_3b\"\n",
    "config = LlamaConfig.from_pretrained(model_name_or_path)\n",
    "model = LlamaForCausalLM.from_pretrained(model_name_or_path,config = config,load_in_8bit = True,torch_dtype = torch.float16,) \n",
    "tokenizer = LlamaTokenizer.from_pretrained(\n",
    "        model_name_or_path,trust_remote_code = True\n",
    "        )\n",
    "tokenizer.pad_token =\"[PAD]\"\n",
    "special_tokens = {'bos_token': \"<s>\",'eos_token': \"</s>\",'unk_token': \"<unk>\"}\n",
    "for k, val in special_tokens.items():\n",
    "    tokenizer.add_special_tokens({k: val})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b28af2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset, DatasetDict, load_dataset, load_from_disk\n",
    "tokenizer_name = tokenizer.__class__.__name__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf5a6f69",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset json (/home/mkarri/.cache/huggingface/datasets/teknium___json/teknium--GPT4-LLM-Cleaned-a71aa8ae1ac3982d/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n",
      "100%|████████████████████████████████████████████| 1/1 [00:00<00:00, 310.80it/s]\n"
     ]
    }
   ],
   "source": [
    "ds = load_dataset(\"teknium/GPT4-LLM-Cleaned\",streaming=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c37cce5f",
   "metadata": {},
   "source": [
    "**The below code has been adapted from the Open AI Access Collective which can be referenced [here](https://github.com/OpenAccess-AI-Collective/axolotl).**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d85d156f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The code below has been taken from the OpenAI Collective which is located \n",
    "# at the location: https://github.com/OpenAccess-AI-Collective/axolotl\n",
    "# Please note that this template code was modified for the purposes of this notebook\n",
    "\n",
    "import abc\n",
    "import functools\n",
    "from typing import List, Tuple, Union\n",
    "from datasets import IterableDataset\n",
    "from enum import Enum, auto\n",
    "from typing import Generator, List, Optional, Tuple, Union\n",
    "\n",
    "class InvalidDataException(Exception):\n",
    "    \"\"\"\n",
    "    Exception raised when the data is invalid\n",
    "    \"\"\"\n",
    "class PromptStyle(Enum):\n",
    "    \"\"\"\n",
    "    Enum for prompt styles\n",
    "    \"\"\"\n",
    "\n",
    "    INSTRUCT = \"instruct\"\n",
    "    CHAT = \"chat\"\n",
    "\n",
    "\n",
    "class AlpacaPrompter:\n",
    "    \"\"\"\n",
    "    Base class for alpaca prompters\n",
    "    \"\"\"\n",
    "\n",
    "    system_prompt = \"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\"\n",
    "    system_no_input_prompt = \"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n\"\n",
    "    prompt_style: Optional[PromptStyle] = None\n",
    "\n",
    "    def __init__(self, prompt_style=PromptStyle.INSTRUCT.value):\n",
    "        self.prompt_style = prompt_style if prompt_style else PromptStyle.INSTRUCT.value\n",
    "        self.match_prompt_style()\n",
    "\n",
    "    def match_prompt_style(self):\n",
    "        if self.prompt_style == PromptStyle.INSTRUCT.value:\n",
    "            self.prompt_input = (\n",
    "                self.system_prompt\n",
    "                + \"### Instruction:\\n{instruction}\\n\\n### Input:\\n{input}\\n\\n### Response:\\n\"\n",
    "            )\n",
    "            self.prompt_no_input = (\n",
    "                self.system_no_input_prompt\n",
    "                + \"### Instruction:\\n{instruction}\\n\\n### Response:\\n\"\n",
    "            )\n",
    "            self.response_split = \"### Response:\"\n",
    "        if self.prompt_style == PromptStyle.CHAT.value:\n",
    "            self.prompt_input = (\n",
    "                self.system_prompt + \"USER: {instruction}\\n{input}\\nASSISTANT:\"\n",
    "            )\n",
    "            self.prompt_no_input = (\n",
    "                self.system_no_input_prompt + \"USER: {instruction}\\nASSISTANT:\"\n",
    "            )\n",
    "            self.response_split = \"ASSISTANT:\"\n",
    "\n",
    "    def build_prompt(\n",
    "        self,\n",
    "        instruction: str,\n",
    "        input: Union[None, str] = None,  # pylint: disable=redefined-builtin\n",
    "        output: Union[None, str] = None,\n",
    "    ) -> Generator[str, None, None]:\n",
    "        # returns the full prompt from instruction and optional input\n",
    "        # if a label (=response, =output) is provided, it's also appended.\n",
    "        if input:\n",
    "            res = self.prompt_input.format(instruction=instruction, input=input)\n",
    "        else:\n",
    "            res = self.prompt_no_input.format(instruction=instruction)\n",
    "        if output:\n",
    "            res = f\"{res}{output}\"\n",
    "        yield res\n",
    "\n",
    "    def get_response(self, output: str) -> str:\n",
    "        return output.split(self.response_split)[1].strip()\n",
    "\n",
    "\n",
    "\n",
    "class PromptTokenizingStrategy(abc.ABC):\n",
    "    \"\"\"\n",
    "    Abstract class for tokenizing strategies\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        prompter,\n",
    "        tokenizer,\n",
    "        train_on_inputs: bool = False,\n",
    "        sequence_len: int = 2048,\n",
    "    ):\n",
    "        self.prompter = prompter\n",
    "        self.tokenizer: PreTrainedTokenizer = tokenizer\n",
    "        self.train_on_inputs = train_on_inputs\n",
    "        self.sequence_len = sequence_len\n",
    "\n",
    "    @abc.abstractmethod\n",
    "    def tokenize_prompt(self, prompt):\n",
    "        pass\n",
    "\n",
    "    @functools.lru_cache(maxsize=128)\n",
    "    def _get_user_token(self):\n",
    "        id_or_ids = self.tokenizer.convert_tokens_to_ids(\"<|USER|>\")\n",
    "        if isinstance(id_or_ids, (int,)):\n",
    "            return id_or_ids\n",
    "        return False\n",
    "\n",
    "    @functools.lru_cache(maxsize=128)\n",
    "    def _get_assistant_token(self):\n",
    "        id_or_ids = self.tokenizer.convert_tokens_to_ids(\"<|ASSISTANT|>\")\n",
    "        if isinstance(id_or_ids, (int,)):\n",
    "            return id_or_ids\n",
    "        return False\n",
    "\n",
    "    def _tokenize(self, prompt: str, add_eos_token=True, strip_bos_token=False):\n",
    "        result = self.tokenizer(\n",
    "            prompt,\n",
    "            truncation=True,\n",
    "            max_length=self.sequence_len,\n",
    "            padding=False,\n",
    "            return_tensors=None,\n",
    "        )\n",
    "        if (\n",
    "            result[\"input_ids\"][-1] != self.tokenizer.eos_token_id\n",
    "            and len(result[\"input_ids\"]) < self.sequence_len\n",
    "            and add_eos_token\n",
    "        ):\n",
    "            result[\"input_ids\"].append(self.tokenizer.eos_token_id)\n",
    "            result[\"attention_mask\"].append(1)\n",
    "\n",
    "        if result[\"input_ids\"][0] == self.tokenizer.bos_token_id and strip_bos_token:\n",
    "            result[\"input_ids\"] = result[\"input_ids\"][1:]\n",
    "            result[\"attention_mask\"] = result[\"attention_mask\"][1:]\n",
    "\n",
    "        result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "        return result\n",
    "\n",
    "\n",
    "class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy):\n",
    "    \"\"\"\n",
    "    Tokenizing strategy for instruction-based prompts.\n",
    "    \"\"\"\n",
    "\n",
    "    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def tokenize_prompt(self, prompt):\n",
    "        (\n",
    "            instruction,\n",
    "            input,  # pylint: disable=redefined-builtin\n",
    "            response,\n",
    "        ) = self.parse_instruction_fields(prompt)\n",
    "        \n",
    "        full_prompt = self._build_full_prompt(instruction, input, response)\n",
    "        tokenized_full_prompt = self._tokenize(full_prompt)\n",
    "        if not self.train_on_inputs:\n",
    "            user_prompt = next(\n",
    "                iter(\n",
    "                    self.prompter.build_prompt(\n",
    "                        instruction,\n",
    "                        input,\n",
    "                    )\n",
    "                )\n",
    "            )\n",
    "            tokenized_user_prompt = self._tokenize(user_prompt, add_eos_token=False)\n",
    "            user_prompt_len = len(tokenized_user_prompt[\"input_ids\"])\n",
    "            # TODO this could be sped up using numpy array slicing\n",
    "            tokenized_full_prompt[\"labels\"] = [\n",
    "                -100\n",
    "            ] * user_prompt_len + tokenized_full_prompt[\"labels\"][user_prompt_len:]\n",
    "\n",
    "        return tokenized_full_prompt\n",
    "\n",
    "    def _build_full_prompt(\n",
    "        self, instruction, input, response  # pylint: disable=redefined-builtin\n",
    "    ):\n",
    "        return next(\n",
    "            iter(\n",
    "                self.prompter.build_prompt(\n",
    "                    instruction,\n",
    "                    input,\n",
    "                    response,\n",
    "                )\n",
    "            )\n",
    "        )\n",
    "\n",
    "class TokenizedPromptDataset(IterableDataset):\n",
    "    \"\"\"\n",
    "    Iterable dataset that returns tokenized prompts from a stream of text files.\n",
    "        Args:\n",
    "            prompt_tokenizer (PromptTokenizingStrategy): The prompt tokenizing method for proccessing the data.\n",
    "            dataset (dataset.Dataset): Dataset with text files.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(  # pylint: disable=super-init-not-called\n",
    "        self,\n",
    "        prompt_tokenizer: PromptTokenizingStrategy,\n",
    "        dataset: IterableDataset,\n",
    "    ):\n",
    "        self.prompt_tokenizer = prompt_tokenizer\n",
    "        self.dataset = dataset\n",
    "\n",
    "    def __iter__(self):\n",
    "        iterator = iter(self.dataset)\n",
    "        count = 0\n",
    "        # Loop through the entire dataset\n",
    "        for example in iterator:\n",
    "            try:\n",
    "                yield self.prompt_tokenizer.tokenize_prompt(example)\n",
    "                count += 1\n",
    "            except InvalidDataException:\n",
    "                pass\n",
    "        if count == 0:\n",
    "            raise RuntimeError(\"Expected at least one datapoint in dataset.\")\n",
    "\n",
    "class AlpacaPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):\n",
    "    \"\"\"\n",
    "    Tokenizing strategy for Alpaca prompts.\n",
    "    \"\"\"\n",
    "\n",
    "    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:\n",
    "        return (\n",
    "            prompt[\"instruction\"],\n",
    "            prompt[\"input\"] if \"input\" in prompt else \"\",\n",
    "            prompt[\"output\"],\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14cb3aa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_strategy = AlpacaPromptTokenizingStrategy(AlpacaPrompter('instruct'),tokenizer,False,256)\n",
    "ds_wrapper = TokenizedPromptDataset(ds_strategy, ds['train'])\n",
    "samples = []\n",
    "for d in [ds_wrapper]:\n",
    "    samples = samples + list(d)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54fd519f",
   "metadata": {},
   "source": [
    "**So what is happening here?**  \n",
    "  \n",
    "**Step 1**: Create the appropriate prompt template for the dataset. Below is the prompt for the first example of the dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dab0fb70",
   "metadata": {},
   "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",
      "Give three tips for staying healthy.\n",
      "\n",
      "### Response:\n",
      "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",
      "\n",
      "2. 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",
      "\n",
      "3. 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.\n"
     ]
    }
   ],
   "source": [
    "print(next(AlpacaPrompter('instruct').build_prompt(ds['train'][0]['instruction'],ds['train'][0]['input'],ds['train'][0]['output'])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "141795a3",
   "metadata": {},
   "source": [
    "**Step 2**: Create the appropriate tokenizations. Here, we can see that inputs to the model consist of both the instruction \n",
    "    as well as the answer. The labels or the output based on which the model needs to be finetuned has some -100s to it.  However, the text form of the labels denoted as  -100 noted below:-<br>\n",
    "\n",
    "Below is an instruction that describes a task.  \n",
    "\n",
    "Write a response that appropriately completes the request.  \n",
    "***USER:*** *Give three tips for staying healthy.*  \n",
    "This is done so that we can parallely process the entire sample at one go.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80ee3995",
   "metadata": {},
   "source": [
    "**Step 3**: Split into train and test datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cef6008",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Dataset.from_list(samples).shuffle(seed=42)\n",
    "dataset = dataset.train_test_split(test_size=0.02, shuffle=False)\n",
    "train_dataset = dataset['train']\n",
    "eval_dataset = dataset['test']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c486bef",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d794d919",
   "metadata": {},
   "source": [
    "## Model Implementation and Methods"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "250f4fdc",
   "metadata": {},
   "source": [
    "### Loading the base Open Llama Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d5d921e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import LlamaTokenizer, LlamaForCausalLM,LlamaConfig\n",
    "\n",
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "import torch\n",
    "model_name_or_path = \"openlm-research/open_llama_3b\"\n",
    "config = LlamaConfig.from_pretrained(model_name_or_path)\n",
    "model = LlamaForCausalLM.from_pretrained(model_name_or_path,config = config,load_in_8bit = True,torch_dtype = torch.float16,) \n",
    "tokenizer = LlamaTokenizer.from_pretrained(\n",
    "        model_name_or_path,trust_remote_code = True\n",
    "        )\n",
    "tokenizer.pad_token =\"[PAD]\"\n",
    "special_tokens = {'bos_token': \"<s>\",'eos_token': \"</s>\",'unk_token': \"<unk>\"}\n",
    "for k, val in special_tokens.items():\n",
    "    tokenizer.add_special_tokens({k: val})\n",
    "inputs = tokenizer(\"Tell me about Alpacas.\", return_tensors=\"pt\",truncation=False).to('cuda')\n",
    "generation_output = model.generate(**inputs,max_new_tokens=32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afb18164",
   "metadata": {},
   "source": [
    "**Note: The below cell shows a way to call the model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b334a783",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s>Tell me about Alpacas.\n",
      "Alpacas are a type of camelid, a group of animals that includes llamas, guanacos, vicunas,\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.decode(generation_output[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99a34919",
   "metadata": {},
   "source": [
    "Hurray! We just loaded an Open Llama model with 3 billion parameters and are able to generate inferences. This is a huge step. To put this into perspective, the inference requires us to initialize the model and load **> 3 GB** of parameter weights alone."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2aefa40",
   "metadata": {},
   "source": [
    "### Ruminations on model architectures: Open Llama vs ChatGPT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02f7cf9f",
   "metadata": {},
   "source": [
    "ChatGPT is a SOTA LLM which is based on GPT-3.5 and follows the transformer architecture for GPT-3.5. This is based on a decoder only architecture whose primary aim is to predict the next token. More details can be found [here](https://towardsdatascience.com/gpt-3-explained-19e5f2bd3288)  \n",
    "  \n",
    "\n",
    "**To Note:** Decoder architecture which is the backbone for the GPT-3.5 architecture has the primary objective of predicting the next word.  \n",
    "  \n",
    "![Alt Text](./GPTarchitecture.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07550b4b",
   "metadata": {},
   "source": [
    "**Image Source**: The image above has been taken from the paper present [here](https://arxiv.org/pdf/2005.14165.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f9e0381",
   "metadata": {},
   "source": [
    "A traditional GPT model consists of stacked decoders. The $n_{\\texttt{layers}}$ = 96 layers listed above (for GPT-3 Large) consists of a number of stacked decoders. The decoders of the GPT 3.5 model alternate between sparse and self attention layers but more on that later. We can see this in the below image."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bbce637",
   "metadata": {},
   "source": [
    "In the below the most important parameters are:-\n",
    "- $d_{\\texttt{model}}$: The embedding dimensions of the input text\n",
    "- $n_{\\texttt{layers}}$: The number of decoder layers present."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b8a94ed",
   "metadata": {},
   "source": [
    "There are 96 of these encoders in the GPT architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03ab7a2a",
   "metadata": {},
   "source": [
    "<!-- ![Alt Text](decoder.png) -->\n",
    "<image src=\"./decoder.png\" alt=\"Decoder Architecture\" height=\"300\" width=\"300\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9caadc72",
   "metadata": {},
   "source": [
    "The model we have chosen is Open LLama-3B which is an Apache 2.0 version of Llama. It replicates the Llama architecture which is an open source model by Meta AI. We can check the model at the [location](https://arxiv.org/abs/2302.13971)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa4c0bc3",
   "metadata": {},
   "source": [
    "**We can have a detailed look at the architecture of Open Llama below.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4384285d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(32000, 3200, padding_idx=0)\n",
       "    (layers): ModuleList(\n",
       "      (0-25): 26 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaAttention(\n",
       "          (q_proj): Linear8bitLt(in_features=3200, out_features=3200, bias=False)\n",
       "          (k_proj): Linear8bitLt(in_features=3200, out_features=3200, bias=False)\n",
       "          (v_proj): Linear8bitLt(in_features=3200, out_features=3200, bias=False)\n",
       "          (o_proj): Linear8bitLt(in_features=3200, out_features=3200, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear8bitLt(in_features=3200, out_features=8640, bias=False)\n",
       "          (down_proj): Linear8bitLt(in_features=8640, out_features=3200, bias=False)\n",
       "          (up_proj): Linear8bitLt(in_features=3200, out_features=8640, bias=False)\n",
       "          (act_fn): SiLUActivation()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=3200, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0c900d",
   "metadata": {},
   "source": [
    "In the above the most important parameters are:-\n",
    "\n",
    "- $d_{\\texttt{model}}$: 3200\n",
    "- $n_{\\texttt{layers}}$: 26  \n",
    "  \n",
    "<br>The main difference between the traditional decoder and the LLama paper are as below(https://paperswithcode.com/method/llama) :-\n",
    "1. RMSNorm normalizing function is used to improve the training stability, by normalizing the input of each transformer sub-layer, instead of normalizing the output.\n",
    "2. The ReLU non-linearity is replaced by the SwiGLU activation function to improve performance.\n",
    "3. Absolute positional embeddings are removed and instead rotary positional embeddings (RoPE) are added at each layer of the network.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5241a3f4",
   "metadata": {},
   "source": [
    "Based on the code in the lLamma paper(https://github.com/facebookresearch/llama/blob/main/llama/model.py), the main changes have been done to the code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97162660",
   "metadata": {},
   "source": [
    "![Alt Text](llama.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b24672b",
   "metadata": {},
   "source": [
    "The feedforward layer has been changed as below:-"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abbae263",
   "metadata": {},
   "source": [
    "![Swish ReLU](./swishrelu.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f979c7bc",
   "metadata": {},
   "source": [
    "Both rotary embeddings and swish Relu are out of context for this paper. For more information refer to (https://blog.eleuther.ai/rotary-embeddings/#how-is-this-different-from-the-sinusoidal-embeddings-used-in-attention-is-all-you-need) for rotary embeddings and Swish Relu(https://arxiv.org/pdf/2002.05202.pdf) The main point we are trying to make is that there have been minor changes to the original transformer architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d6f5b6",
   "metadata": {},
   "source": [
    "### Model Inference and Finetuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "814d5800",
   "metadata": {},
   "source": [
    "Model Inference is done the following way as per the above example. There is a sequential way in which the output is determined.:-\n",
    "\n",
    "- **Input**: Tell me about Alpacas.   \n",
    "**Output**: [BOS] \n",
    "\n",
    "- **Input**: Tell me about Alpacas. <BOS>   \n",
    "**Output**: Alpacas  \n",
    "\n",
    "- **Input**: Tell me about Alpacas. <BOS> Alpacas   \n",
    "**Output**: are  \n",
    "\n",
    "- **Input**: Tell me about Alpacas. <BOS> Alpacas are   \n",
    "**Output**: a  \n",
    "\n",
    "- **Input**: Tell me about Alpacas. <BOS> Alpacas are a   \n",
    "**Output**: type  \n",
    "\n",
    "- **Input**: Tell me about Alpacas. <BOS> Alpacas are a type   \n",
    "**Output**: of  \n",
    "\n",
    "- **Input**: Tell me about Alpacas. <BOS> Alpacas are a type of    \n",
    "**Output**: camel  \n",
    "\n",
    "- **Input**: Tell me about Alpacas. <BOS> Alpacas are a type of camel   \n",
    "**Output**: [EOS]  \n",
    "\n",
    "- **Final Input**: Tell me about Alpacas. <BOS> Alpacas are a type of camel   \n",
    "**Output**: [EOS]  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d6e5ce4",
   "metadata": {},
   "source": [
    "Model finetuning is a supervised learning process in contrast to model training which is semi supervised. \n",
    "The cross entropy loss is the loss function. An example is shown below:-\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a0e4180",
   "metadata": {},
   "source": [
    "**Input**: Tell me about Alpacas.<br>\n",
    "**Output**: Alpacas are a type of camel found in South America.<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a335275",
   "metadata": {},
   "source": [
    "$$ CE_{loss} = - \\bigg(log(P(Alpacas)) + log(P(are|Alpacas)) + log(P(a|Alpacas \\space are))   \\\\\n",
    "+ log(P(type|Alpacas \\space are \\space a)) + log(P(camel|Alpacas \\space are \\space a \\space type)) + ...\\bigg) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec418c91",
   "metadata": {},
   "source": [
    "### Resource Requirement"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b6cd374",
   "metadata": {},
   "source": [
    "LLMs take an enormous amount of resources to train.  \n",
    "For e.g. from the Llama paper the 65B parameter model took 21 days on 2048 A100 GPUs. Hence the total cost of training a 65B model would take:  \n",
    "2048(# of GPUs) * 24(# hours per day) * 21(# days of training) * 8.8(# Cost of GPU per hour) = $449k"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0795b663",
   "metadata": {},
   "source": [
    "A100 GPUs are state of the art and are not consumer GPUs. While finetuning should likely cost a fraction of the cost, there is still a significant \n",
    "portion of compute time as well as power. Finetuning based on the Llama 7B based on the Stanford Alpaca(https://crfm.stanford.edu/2023/03/13/alpaca.html) took ~$600. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc71c1b",
   "metadata": {},
   "source": [
    "### LoRA for finetuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ac9b649",
   "metadata": {},
   "source": [
    "While we currently don't have access to a A100 GPU, a finetuning experiment will still take time. Based on the model configurations above, even ordinary finetuning can take time. Hence, in this case, we use LoRA (Low Rank Adapters) which ingest trainable parameters to each layer of the model.  \n",
    "\n",
    "**LoRA** is based on this [paper](https://arxiv.org/pdf/2106.09685.pdf) and notes that during finetuning, most of the delta weights are mainly zero and hence have a \"low instrinsic dimension\". Hence, they can still learn efficiently even if projected to a smaller subpace. Hence, when learning the delta weights, we follow the following equation:   \n",
    "   \n",
    "$$h = W_0x + \\Delta Wx = W_0x + BAx$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f332d212",
   "metadata": {},
   "source": [
    "The main benefit of this is simplicity and no additional latency inference. B & A are known as the Lora weight matrices. The main parameter of the LORA adapter is the rank($r$) and the idea goes as below: \n",
    "$$(\\Delta Wx).shape = (M, N)$$\n",
    "$$\\Delta Wx = A_{M \\times r} * B_{r \\times N}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea264018",
   "metadata": {},
   "source": [
    "### Preparing model for finetuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b17d1b6e",
   "metadata": {},
   "source": [
    "We use the Peft package from Huggingface to assist us in the LORA finetuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "babe5af5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, PeftModel, get_peft_model\n",
    "from peft import prepare_model_for_kbit_training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "69c4ef7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_target_modules = ['gate_proj','down_proj','up_proj','q_proj','v_proj','k_proj','o_proj']\n",
    "lora_alpha = 16\n",
    "lora_r = 8\n",
    "lora_config = LoraConfig(\n",
    "        r=lora_r,\n",
    "        lora_alpha=lora_alpha,\n",
    "        target_modules=lora_target_modules,\n",
    "        task_type=\"CAUSAL_LM\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8544fd6b",
   "metadata": {},
   "source": [
    "Here, $LoRA_r$ refers to the low rank taken in this case. $ LoRA_\\alpha $ refers to the alpha parameter which indicates the scaling of the weights \n",
    "when calculating gradients. The lora target modules defines the layers where LoRA Adapters are being added."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "25fb1ab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = prepare_model_for_kbit_training(\n",
    "            model, use_gradient_checkpointing=True\n",
    "        )\n",
    "model = get_peft_model(model, lora_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8592f51c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 12,712,960 || all params: 3,439,186,560 || trainable%: 0.36965020007521776\n"
     ]
    }
   ],
   "source": [
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c279688b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear8bitLt(\n",
       "  in_features=3200, out_features=3200, bias=False\n",
       "  (lora_dropout): ModuleDict(\n",
       "    (default): Identity()\n",
       "  )\n",
       "  (lora_A): ModuleDict(\n",
       "    (default): Linear(in_features=3200, out_features=8, bias=False)\n",
       "  )\n",
       "  (lora_B): ModuleDict(\n",
       "    (default): Linear(in_features=8, out_features=3200, bias=False)\n",
       "  )\n",
       "  (lora_embedding_A): ParameterDict()\n",
       "  (lora_embedding_B): ParameterDict()\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.base_model.model.model.layers[0].self_attn.q_proj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d741d72a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LORA A weights\n",
      "OrderedDict([('weight', tensor([[ 0.0119, -0.0117, -0.0060,  ...,  0.0118, -0.0083,  0.0060],\n",
      "        [-0.0040, -0.0096,  0.0049,  ..., -0.0099, -0.0024, -0.0158],\n",
      "        [-0.0110, -0.0039,  0.0003,  ...,  0.0128,  0.0003, -0.0120],\n",
      "        ...,\n",
      "        [-0.0010,  0.0118,  0.0046,  ..., -0.0135,  0.0029,  0.0147],\n",
      "        [-0.0059,  0.0155, -0.0033,  ..., -0.0075,  0.0021,  0.0076],\n",
      "        [ 0.0140,  0.0060, -0.0108,  ..., -0.0110,  0.0146,  0.0043]],\n",
      "       device='cuda:0'))])\n",
      "LORA B weights\n",
      "OrderedDict([('weight', tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        ...,\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.]], device='cuda:0'))])\n",
      "trainable params: 12,712,960 || all params: 3,439,186,560 || trainable%: 0.36965020007521776\n"
     ]
    }
   ],
   "source": [
    "print(\"LORA A weights\")\n",
    "print(model.base_model.model.model.layers[0].self_attn.q_proj.lora_A.default.state_dict())\n",
    "print(\"LORA B weights\")\n",
    "print(model.base_model.model.model.layers[0].self_attn.q_proj.lora_B.default.state_dict())\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72b4fba9",
   "metadata": {},
   "source": [
    "The LORA B parameters are set to 0 so that the randomly initialized weights do not affect the model performance at the start of the training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e035a5e",
   "metadata": {},
   "source": [
    "## Experiments and Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf3570a4",
   "metadata": {},
   "source": [
    "### Prepare the trainer for training\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b890902e",
   "metadata": {},
   "source": [
    "Here, we use the Huggingface Trainer which allows for a config driven management."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac84aaeb",
   "metadata": {},
   "source": [
    "Now we set some learning parameters for our trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a1506ab6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Learning Rate \n",
    "learning_rate = 0.0002\n",
    "# Weight decay associated with learning parameter\n",
    "weight_decay = 0.0\n",
    "batch_size = 16\n",
    "micro_batch_size = 4\n",
    "eval_steps = 50\n",
    "save_steps = 1000\n",
    "# Num of epochs\n",
    "num_epochs = 3\n",
    "# Optimizer(We use Adam 8 bit as the model has been loaded in 8 bit)\n",
    "optimizer = 'adamw_bnb_8bit'\n",
    "import math\n",
    "total_num_steps = int(\n",
    "    math.ceil(len(train_dataset) * num_epochs / batch_size)\n",
    ")\n",
    "warmup_steps = 10\n",
    "logging_steps = max(min(int(0.005 * total_num_steps), 10), 1)\n",
    "training_arguments_kwargs = {}\n",
    "## Train model in FP16 or float because it has been loaded in 8 bit\n",
    "training_arguments_kwargs[\"fp16\"] = True\n",
    "training_arguments_kwargs[\"tf32\"] = False\n",
    "## Gradient Checkpointing\n",
    "gradient_checkpointing = True\n",
    "training_arguments_kwargs[\"warmup_steps\"] = warmup_steps\n",
    "training_arguments_kwargs[\"logging_steps\"] = logging_steps\n",
    "training_arguments_kwargs[\"gradient_checkpointing\"] = gradient_checkpointing\n",
    "gradient_accumulation_steps = batch_size // micro_batch_size\n",
    "optimizer = \"adamw_bnb_8bit\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fb468df",
   "metadata": {},
   "source": [
    "We use a 8bit Adam which is quantized for 8bits. For more on quantization please refer this [blog](https://huggingface.co/blog/hf-bitsandbytes-integration)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7a6060d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "training_args = transformers.TrainingArguments(\n",
    "        per_device_train_batch_size=micro_batch_size,\n",
    "        per_device_eval_batch_size=micro_batch_size,\n",
    "        gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "        eval_accumulation_steps=gradient_accumulation_steps,\n",
    "        num_train_epochs=num_epochs,\n",
    "        learning_rate=learning_rate,\n",
    "        evaluation_strategy=\"steps\",\n",
    "        save_strategy=\"steps\",\n",
    "        eval_steps=eval_steps,\n",
    "        save_steps=save_steps,\n",
    "        output_dir='./lora-out',\n",
    "        save_total_limit=3,\n",
    "        group_by_length=False,\n",
    "        report_to=None,\n",
    "        run_name=None,\n",
    "        optim=optimizer,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        weight_decay=weight_decay,\n",
    "        **training_arguments_kwargs,\n",
    "    )\n",
    "trainer_kwargs = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6735a6a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.trainer_pt_utils import get_parameter_names\n",
    "from torch import nn\n",
    "import bitsandbytes as bnb\n",
    "decay_parameters = get_parameter_names(model, [nn.LayerNorm])\n",
    "decay_parameters = [name for name in decay_parameters if \"bias\" not in name]\n",
    "optimizer_grouped_parameters = [\n",
    "    {\n",
    "        \"params\": [\n",
    "            p\n",
    "            for n, p in model.named_parameters()\n",
    "            if (n in decay_parameters and p.requires_grad)\n",
    "        ],\n",
    "        \"weight_decay\": training_args.weight_decay,\n",
    "    },\n",
    "    {\n",
    "        \"params\": [\n",
    "            p\n",
    "            for n, p in model.named_parameters()\n",
    "            if (n not in decay_parameters and p.requires_grad)\n",
    "        ],\n",
    "        \"weight_decay\": 0.0,\n",
    "    },\n",
    "]\n",
    "\n",
    "optimizer = bnb.optim.Adam8bit(\n",
    "    optimizer_grouped_parameters,\n",
    "    betas=(training_args.adam_beta1, training_args.adam_beta2),\n",
    "    eps=training_args.adam_epsilon,\n",
    "    lr=training_args.learning_rate,\n",
    ")\n",
    "lr_scheduler = transformers.get_cosine_schedule_with_warmup(\n",
    "                optimizer,\n",
    "                training_args.warmup_steps,\n",
    "                total_num_steps,\n",
    "            )\n",
    "trainer_kwargs[\"optimizers\"] = (optimizer, lr_scheduler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "59066962",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator_kwargs = {\n",
    "        \"padding\": True,\n",
    "    }\n",
    "data_collator_kwargs[\"pad_to_multiple_of\"] = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70fae240",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model.config.use_cache = False\n",
    "import torch\n",
    "import transformers\n",
    "model = torch.compile(model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "47085e99",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = transformers.Trainer(\n",
    "        model=model,\n",
    "        train_dataset=train_dataset,\n",
    "        eval_dataset=eval_dataset,\n",
    "        args=training_args,\n",
    "        data_collator=transformers.DataCollatorForSeq2Seq(\n",
    "            tokenizer,\n",
    "            return_tensors=\"pt\",\n",
    "            **data_collator_kwargs,\n",
    "        ),\n",
    "        **trainer_kwargs,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f169275f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.4"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/wreckit/huggingface/runs/7pbswgnm' target=\"_blank\">comic-waterfall-20</a></strong> to <a href='https://wandb.ai/wreckit/huggingface' 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/wreckit/huggingface' target=\"_blank\">https://wandb.ai/wreckit/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/wreckit/huggingface/runs/7pbswgnm' target=\"_blank\">https://wandb.ai/wreckit/huggingface/runs/7pbswgnm</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='35' max='1431' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [  35/1431 16:44 < 11:47:52, 0.03 it/s, Epoch 0.07/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b884a945",
   "metadata": {},
   "source": [
    "Once the model training is completed (it takes 11 hours.), the model adapter weights can be found at (https://huggingface.co/maneel/OpenLLama3B_8bitQuantized_alpaca_finetuned)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03c4a49d",
   "metadata": {},
   "source": [
    "### Model Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d250e14e",
   "metadata": {},
   "source": [
    "Once, we have the model trained. we need to make an endpoint from it so that we can perform inference from it easily.  \n",
    "The framework we have used is Ray Serve and heavily inspired from this [notebook](https://docs.ray.io/en/latest/ray-air/examples/gptj_serving.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f693586b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from ray import serve\n",
    "from starlette.requests import Request\n",
    "\n",
    "\n",
    "@serve.deployment(ray_actor_options={\"num_gpus\": 1})\n",
    "class PredictDeployment:\n",
    "    def __init__(self, model_id: str, revision: str = None):\n",
    "        # from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "        from peft import PeftModel\n",
    "        from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig\n",
    "        import torch\n",
    "\n",
    "        model = LlamaForCausalLM.from_pretrained(model_id,\n",
    "                                                      torch_dtype=torch.float16, \n",
    "                                                      low_cpu_mem_usage=True,\n",
    "                                                      load_in_8bit=True, \n",
    "                                                      device_map=\"auto\")\n",
    "        #load the adapter delta weights on top of base model\n",
    "        self.model = PeftModel.from_pretrained(model, \"maneel/OpenLLama3B_8bitQuantized_alpaca_finetuned\",\\\n",
    "                                               adapter_name=\"maneel_openllama\")\n",
    "        \n",
    "        self.tokenizer = LlamaTokenizer.from_pretrained(model_id)\n",
    "        \n",
    "        self.generation_config = GenerationConfig(temperature=0.8,\n",
    "                                             top_p=0.75,\n",
    "                                             top_k=40,\n",
    "                                             num_beams=4,\n",
    "                                             no_repeat_ngram_size=3,\n",
    "                                             max_new_tokens=256)\n",
    "        \n",
    "    def generate(self, text: str) -> pd.DataFrame:\n",
    "        input_ids = self.tokenizer(text, return_tensors=\"pt\").input_ids.to(\n",
    "            self.model.device\n",
    "        )\n",
    "\n",
    "        gen_tokens = self.model.generate(\n",
    "            input_ids=input_ids,\n",
    "            generation_config=self.generation_config,\n",
    "        )\n",
    "        return pd.DataFrame(\n",
    "            self.tokenizer.batch_decode(gen_tokens), columns=[\"responses\"]\n",
    "        )\n",
    "\n",
    "    async def __call__(self, http_request: Request) -> str:\n",
    "        json_request: str = await http_request.json()\n",
    "        prompts = []\n",
    "        for prompt in json_request:\n",
    "            text = prompt[\"text\"]\n",
    "            if isinstance(text, list):\n",
    "                prompts.extend(text)\n",
    "            else:\n",
    "                prompts.append(text)\n",
    "        return self.generate(prompts)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15fc23a3",
   "metadata": {},
   "source": [
    "As we can see below in the outputs of the cell, a ray server has been launched."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d9ad2138",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-28 19:20:21,102\tINFO worker.py:1627 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32m127.0.0.1:8265 \u001b[39m\u001b[22m\n",
      "\u001b[2m\u001b[36m(ServeController pid=3936403)\u001b[0m INFO 2023-06-28 19:20:24,806 controller 3936403 deployment_state.py:1298 - Deploying new version of deployment default_PredictDeployment.\n",
      "\u001b[2m\u001b[36m(HTTPProxyActor pid=3936481)\u001b[0m INFO:     Started server process [3936481]\n",
      "\u001b[2m\u001b[36m(ServeController pid=3936403)\u001b[0m INFO 2023-06-28 19:20:24,876 controller 3936403 deployment_state.py:1537 - Adding 1 replica to deployment default_PredictDeployment.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m [2023-06-28 19:20:28,006] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m \n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m ===================================BUG REPORT===================================\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m Welcome to bitsandbytes. For bug reports, please run\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m \n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m python -m bitsandbytes\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m \n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m  and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m ================================================================================\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m bin /home/mkarri/anaconda3/envs/axototl/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda121.so\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m CUDA SETUP: CUDA runtime path found: /usr/local/cuda-12.1/lib64/libcudart.so\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m CUDA SETUP: Highest compute capability among GPUs detected: 7.5\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m CUDA SETUP: Detected CUDA version 121\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m CUDA SETUP: Loading binary /home/mkarri/anaconda3/envs/axototl/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda121.so...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m /home/mkarri/anaconda3/envs/axototl/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: /home/mkarri/anaconda3 did not contain ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] as expected! Searching further paths...\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m   warn(msg)\n",
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.\n",
      "\u001b[2m\u001b[36m(ServeController pid=3936403)\u001b[0m WARNING 2023-06-28 19:20:54,990 controller 3936403 deployment_state.py:1869 - Deployment default_PredictDeployment has 1 replicas that have taken more than 30s to initialize. This may be caused by a slow __init__ or reconfigure method.\n",
      "Downloading adapter_model.bin:   0%|          | 0.00/51.0M [00:00<?, ?B/s]\n",
      "Downloading adapter_model.bin:  21%|██        | 10.5M/51.0M [00:00<00:01, 20.6MB/s]\n",
      "Downloading adapter_model.bin:  41%|████      | 21.0M/51.0M [00:00<00:00, 33.2MB/s]\n",
      "Downloading adapter_model.bin:  62%|██████▏   | 31.5M/51.0M [00:00<00:00, 42.6MB/s]\n",
      "Downloading adapter_model.bin:  82%|████████▏ | 41.9M/51.0M [00:01<00:00, 43.2MB/s]\n",
      "Downloading adapter_model.bin: 100%|██████████| 51.0M/51.0M [00:01<00:00, 35.1MB/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "RayServeSyncHandle(deployment='default_PredictDeployment')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_id = \"openlm-research/open_llama_3b\"  #base openllama model\n",
    "revision = \"float16\"\n",
    "deployment = PredictDeployment.bind(model_id=model_id, revision=revision)\n",
    "serve.run(deployment)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efcf2d26",
   "metadata": {},
   "source": [
    "**Note:** The output of the above cell show that our ray serve is up and running and has detected GPUs for launching the inference endpoints."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb549f5e",
   "metadata": {},
   "source": [
    "#### Now, we test our inference endpoint by passing on a request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "dfa0668c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'responses': '<s> Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n### Instruction:\\nExplain in simple terms how the attention mechanism of a transformer model works\\n### Response:\\nA transformer network is a type of artificial neural network that is made up of multiple layers of interconnected nodes. At the topmost layer, called the input layer, the nodes are responsible for inputting the data into the network. The nodes in this layer will take in the input, process it, and pass it on to the next layer.\\nThe next layer, also known as the output layer, takes in the processed data from the previous layer, processes it further, and passes it on. This process continues until the end of the network, where the nodes at the final layer will have processed the input data and passed it on as output.\\nOnce the data has passed through all the layers, the final output will be the processed and pre-processed data. The attention mechanism is a mechanism used in the last layer to help to guide the processing of the data. Essentially, it allows the network to pay attention to the parts of the input that are relevant to the output, and ignore the rest. This helps to improve the accuracy and performance of the model.</s>'}]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(ServeReplica:default_PredictDeployment pid=3936595)\u001b[0m INFO 2023-06-28 19:29:31,585 default_PredictDeployment default_PredictDeployment#ajdfkL szbOguWTta / default replica.py:654 - __CALL__ OK 64223.7ms\n"
     ]
    }
   ],
   "source": [
    "import requests\n",
    "prompt = (\n",
    "    \"\"\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
    "### Instruction:\n",
    "Explain in simple terms how the attention mechanism of a transformer model works\n",
    "### Response:\"\"\"\n",
    ")\n",
    "\n",
    "sample_input = {\"text\": prompt}\n",
    "\n",
    "output = requests.post(\"http://localhost:8000/\", json=[sample_input]).json()\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8677cb71",
   "metadata": {},
   "source": [
    "***Voila! It works!***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1469c4fc",
   "metadata": {},
   "source": [
    "As we can see, the model is at an inference end-point that we can easily query."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e642951",
   "metadata": {},
   "source": [
    "We have used ray serve for inference and serving our model because:\n",
    "- It automatically does load balancing for us\n",
    "- It is able to effectively utilize multiple GPUs for inference\n",
    "- It automatically create RESTful API endpoints\n",
    "- It is easy and intuitive to use"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a14ac2d0",
   "metadata": {},
   "source": [
    "### Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc61e347",
   "metadata": {},
   "source": [
    "As we can see below that the loss continues to decrease as function of the step size and based on the requests above, we can show that the model has had a decent performance on our loss metric which is the Cross Entropy Loss as mentioned earlier here. This was done on our final evaluation dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f57f638",
   "metadata": {},
   "source": [
    "![Loss Results](./loss_results.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "622034d9",
   "metadata": {},
   "source": [
    "## Conclusion and Future Work"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b5f37c",
   "metadata": {},
   "source": [
    "### Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f7759d6",
   "metadata": {},
   "source": [
    "**Achievements**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c73f03b7",
   "metadata": {},
   "source": [
    "We were able to successfully do the following:\n",
    "\n",
    "1. Fine-tune an LLM model (Open Llama-3B) successfully\n",
    "2. Utilized multiple GPUs via distributed training using data parallelism approach\n",
    "3. Utilized parameter efficient finetuning using LoRA \n",
    "4. Successfully deployed the model using ray serve which created an inference endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "806cb5ec",
   "metadata": {},
   "source": [
    "**Implications**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf62c3d3",
   "metadata": {},
   "source": [
    "We were able to finetune an LLM on limited compute resources. This bodes well for a varitey of use cases. Training a large language model from scratch is not feasible for most organizations and individuals. However, using a large language model (particularly a closed source one) presents a number of challenges:\n",
    "1. It might not serve well for our specific use case\n",
    "2. We may not be comfortable in sharing proprietary or sensitive data with major corporations\n",
    "3. Training a Large Language Model from scratch might be extremely resource intensive and computationally burdensome. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db6a2a7a",
   "metadata": {},
   "source": [
    "Our approach is able to successfully resolve all of these issues:\n",
    "1. We were able to use extremely limited resources to get significant boost in performance.\n",
    "2. Finetuning an open source LLM bodes well for a lot of use cases where proprietary information can't be sent to closed source corporations.\n",
    "3. This also gives a big boost to open source development as open source projects with finetuning and proprietary information can beat large closed source models or come close to their performance.\n",
    "4. This democratizes the usage of large language models as using and training them is no longer the preserve of a few Large Corporations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d79c9cb8",
   "metadata": {},
   "source": [
    "**Limitations**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c05c37e4",
   "metadata": {},
   "source": [
    "Our approach however is not without the following limitations:\n",
    "1. Our approach works only for models that can fit on a single GPU which necessarily excludes models above a certain threshold size.\n",
    "2. Our approach also necessitates the usage of parameter efficient finetuning or training approaches which have their own limitations in performance and stability.\n",
    "3. Our approach also requires a large swath of data which might not be readily available in the format that we desire\n",
    "4. Our approach requires us to use low dimensional quantization which is again not without its potential issues as it might lead to the model loosing performance or inference capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd0e2a91",
   "metadata": {},
   "source": [
    "### Future Work"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c9af97",
   "metadata": {},
   "source": [
    "To address some of the limitations as well as level up our work, the next steps that we need to follow are rather clear:\n",
    "1. We need to find ways to train/finetune even larger models that can't fit in one GPU. This can be done by using the following:\n",
    "   1. Model Parallelism using multiple libraries like DeepSpeed (Library by Microsoft), FSDP (Fully Sharded Data Parallelism by Meta AI)\n",
    "   2. Pipeline Parallelism using the same libraries\n",
    "   3. Using custome data types like Bfloat16, FP16 and other quantization techniques like QLoRA\n",
    "2. Use a real world dataset that pertains to a specific domain like Finance or Medicine \n",
    "3. Increase context length using techniques like Landmark-Attention/Superhot/Alibi\n",
    "4. Use other libraries for model serving and inference endpoints\n",
    "5. Try deploying an inference endpoint using multiple GPUs\n",
    "6. Deploying our solution and making it publicly available via a REST API. This may require us to create a platform for the same"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23dd85c1",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e2642f3",
   "metadata": {},
   "source": [
    "------------------------------ **END OF FILE** ----------------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41563548",
   "metadata": {},
   "source": [
    "***"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "e0c937fdf9ccae87d0c83294b3046bb90558c3dea393d9c359be1166f2ad2fa3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}