OpenAI_GPT_v2.5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Welcome to OpenAI GPT v2.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE: This notebook is the continuation of <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/GPT%20from%20Scratch.ipynb\">GPT from Scratch</a> and <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/GPT%20Tokenizer.ipynb\">GPT Tokenizer</a> notebooks.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we are going to reproduce the <a href=\"https://github.com/openai/gpt-2\">OpenAI's GPT 2</a> model, the (`124M`) version of it.\n",
    "\n",
    "Now, when OpenAI released GPT 2, they released it with this <a href=\"https://openai.com/index/better-language-models/\">blog post</a> and this <a href=\"https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf\">paper</a>. And on top of that, they released this <a href=\"https://github.com/openai/gpt-2\">code</a> on GitHub..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But, when reproducing GPT 2, we have to be careful, because we are going to be reproducing the `124M` parameter model. And the thing to be careful with it is there's always a sub-series of models of different sizes when these model releases are made and usually the biggest model is called the **\"GPT\"**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's consider the charts that we have in the paper for a second: \\\n",
    "![OpenAIGPT2 Graphs](ExplanationMedia/Images/OpenAIGPT2Graphs.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the reason we have multiple models is because, according to the above graphs we see that we consider the `Number of parameters in the Language Model` in the `x-axis` and the `y-axis` we put a lot of *downstream metrics* that we are interested in like (\"Translation\", \"Summerization\", \"Question Answering\") and so on and we can chart out the *downstream metrics* as the model size increases.\n",
    "\n",
    "And in the paper we see a table like this:\n",
    "\n",
    "| Parameters | Layers | $d_{model}$ |\n",
    "|------------|--------|-----------|\n",
    "| 117M       | 12     | 768       |\n",
    "| 345M       | 24     | 1024      |\n",
    "| 762M       | 36     | 1280      |\n",
    "| 1542M      | 48     | 1600      |\n",
    "\n",
    "And we see `4` models in the `GPT-2` sub series, starting at `124M` all the way up to `1558M`...\n",
    "\n",
    "But you might be thinking that I might have made a mistake because, in the table the numbers are different and the numbers I spoke of are different. And the reason my numbers disagree with this table is because this entire table is wrong and if we go to their <a href=\"https://github.com/openai/gpt-2\">GitHub repository</a> we see a note that says:\n",
    "> * *Note that our original parameter counts were wrong due to an error (in our previous blog posts and paper). Thus you may have seen small referred to as 117M and medium referred to as 345M.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And in the `124M` parameter model, we see that they used `12 Layers` in the Transformer and `768 Channel Dimensions` in the Transformer.\n",
    "\n",
    "And by the end of this notebook we will try to beat the original `GPT-2 124M` model and will be looking at loss graphs to see our model perform better.\n",
    "\n",
    "The thing to note here is, this paper is more than 5 years old now and it was probably a very complicated optimization at the time and the computation was very low at the time, but today we can reproduce the same model's performance in roughly an hour or so and it will cost us around $10 (if we want to do this on a cloud compute, or in other words, a computer that we can all rent). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And one more thing to mention is, OpenAI did release it's model's weights and it is available on it's GitHub repository, but it's paper is not good with all of it's details with the training.\n",
    "\n",
    "So, in addition to the GPT-2 paper, we will also be referring to the <a href=\"https://arxiv.org/abs/2005.14165\">GPT-3 paper</a>, which is a lot more concrete and a lot of the hyper-parameters and optimization settings and so on, which is not a huge departure from the architecture of GPT-2 version of the model.\n",
    "\n",
    "\n",
    "So, let's do this..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Understanding Hugging Face Pre-Trained Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, the first thing we'd like to do is start at the very end. Or in other words, we'll load the `GPT-2 124M` model as it was released by OpenAI and take it for a spin and sample some `tokens` from it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, the issue is...\n",
    "\n",
    "When we look at the code base and look for the <a href=\"\">`model.py`</a> we see these imports:\n",
    "```python\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.contrib.training import HParams\n",
    "```\n",
    "\n",
    "And we realise that the code is written in <a href=\"https://www.tensorflow.org/\">TensorFlow</a> (another alternative for creating and training deep learning models offered by Google). Meaning that the original `GPT-2` code was written in TensorFlow and is not used anymore...\n",
    "\n",
    "And as per our previous notebooks, we'd like to use PyTorch. And it will be a lot easier if we'd be able to work with the old explanations.\n",
    "\n",
    "But the problem with that is that the initial code is in TensorFlow and we'd like to use PyTorch. So, in order to get the targets we'd like to use the <a href=\"https://huggingface.co/docs/transformers/en/index\">`Hugging Face Transformers Library`</a> released at PyPi. We can use this <a href=\"https://huggingface.co/docs/transformers/en/installation\">installation documentaiton</a> to walk through the steps to install the library in our system..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also check out Hugging Face's implementation of that transformer in their <a href=\"https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py\">`modeling_gpt2.py`</a>. Which did a lot of work to convert all those TensorFlow code to PyTorch such that it becomes easier to load and work with.\n",
    "\n",
    "So in particular we can look at the <a href=\"https://huggingface.co/openai-community/gpt2\">Hugging Face GPT-2</a> model and load it using the Hugging Face transformers...\n",
    "\n",
    "\n",
    "So this is what the code looks like:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from transformers import GPT2LMHeadModel\n",
    "\n",
    "huggingface_model = GPT2LMHeadModel.from_pretrained(\"gpt2\")\n",
    "huggingfaceStateDictionary = huggingface_model.state_dict()\n",
    "\n",
    "for key, value in huggingfaceStateDictionary.items():\n",
    "    print(key, value.shape)\n",
    "```\n",
    "Which gives us the result:\n",
    "```python\n",
    "transformer.wte.weight torch.Size([50257, 768])\n",
    "transformer.wpe.weight torch.Size([1024, 768])\n",
    "transformer.h.0.ln_1.weight torch.Size([768])\n",
    "transformer.h.0.ln_1.bias torch.Size([768])\n",
    "transformer.h.0.attn.c_attn.weight torch.Size([768, 2304])\n",
    "transformer.h.0.attn.c_attn.bias torch.Size([2304])\n",
    "transformer.h.0.attn.c_proj.weight torch.Size([768, 768])\n",
    "transformer.h.0.attn.c_proj.bias torch.Size([768])\n",
    "transformer.h.0.ln_2.weight torch.Size([768])\n",
    "transformer.h.0.ln_2.bias torch.Size([768])\n",
    "transformer.h.0.mlp.c_fc.weight torch.Size([768, 3072])\n",
    "transformer.h.0.mlp.c_fc.bias torch.Size([3072])\n",
    "transformer.h.0.mlp.c_proj.weight torch.Size([3072, 768])\n",
    "transformer.h.0.mlp.c_proj.bias torch.Size([768])\n",
    "transformer.h.1.ln_1.weight torch.Size([768])\n",
    "transformer.h.1.ln_1.bias torch.Size([768])\n",
    "transformer.h.1.attn.c_attn.weight torch.Size([768, 2304])\n",
    "transformer.h.1.attn.c_attn.bias torch.Size([2304])\n",
    "transformer.h.1.attn.c_proj.weight torch.Size([768, 768])\n",
    "transformer.h.1.attn.c_proj.bias torch.Size([768])\n",
    "transformer.h.1.ln_2.weight torch.Size([768])\n",
    "transformer.h.1.ln_2.bias torch.Size([768])\n",
    "transformer.h.1.mlp.c_fc.weight torch.Size([768, 3072])\n",
    "transformer.h.1.mlp.c_fc.bias torch.Size([3072])\n",
    "transformer.h.1.mlp.c_proj.weight torch.Size([3072, 768])\n",
    "transformer.h.1.mlp.c_proj.bias torch.Size([768])\n",
    "...\n",
    "transformer.h.11.mlp.c_proj.bias torch.Size([768])\n",
    "transformer.ln_f.weight torch.Size([768])\n",
    "transformer.ln_f.bias torch.Size([768])\n",
    "lm_head.weight torch.Size([50257, 768])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One awkward thing about this is, when we say `gpt2` it actually loads the `124M` parameter model and if we want the actual `GPT-2` model we'd specify it as `gpt2-xl`...\n",
    "\n",
    "Now when we actually get this `GPT-2` initialized, we want to get the **state dictionary** which is the **raw tensors loaded with values** and we can get those using the `.state_dict()` method. and we can print the `key` (which are the tensors) and the `value` (which are the tensor values) and we can look at the shapes of the `value` tensors to get an idea of the shapes of the states in the model...\n",
    "\n",
    "So, we can now look at the different parameters inside the `GPT-2` model and their shapes...\n",
    "\n",
    "And we can see that there are a lot of short forms of the terms that we already know of, so let's recall that:\n",
    "1. **wte**: Word Token Embeddings\n",
    "2. **wpe**: Word Position Embeddings\n",
    "3. **ln**: Layer Normalization\n",
    "4. **attn**: Attention\n",
    "5. **c_attn**: Cross Attention (awkward because `GPT-2` is a decoder only architecture and should be named **self attention**)\n",
    "6. **c_proj**: Projection layer within attention or MLP\n",
    "7. **mlp**: Multi-Layer Perceptron\n",
    "8. **lm_head**: Language Model Head (output layer)\n",
    "9. **c_fc**: Current/Common Fully Connected Layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We initially can recall the very first key-value pair `transformer.wte.weight torch.Size([50257, 768])` as the `Word Token Embeddings` having a shape of `[50257, 768]` and it comes from the `50257` vocabulary of tokens (which is exactly the number of tokens we spoke about in our <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/GPT%20Tokenizer.ipynb\">Tokenizer Notebook</a>) having `768` feature space (or embedding vector space, or `768 dimensional embedding`)...\n",
    "\n",
    "We can also look at the second key-value pair `transformer.wpe.weight torch.Size([1024, 768])`, we can recall them as `Word Positional Embeddings` having a shape of `[1024, 768]`. So, because `GPT-2` has a maximum sequence length of `1024` we have upto `1024` positions that each token can attend to in the past. And every one of those positions in `GPT-2` has a fixed vector of `768` that is learnt by optimization.\n",
    "\n",
    "And everything else is just the other weights and biases of this transformer..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, now for example, if we take just the `Positional Embeddings` and we flatten it out (we get a `[1, 768]` vector) and take just the first `20` elements of the `768` embeddings we can see that we get the proper weights as an output for this code:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "huggingfaceStateDictionary['transformer.wpe.weight'].view(-1)[:20]\n",
    "```\n",
    "We get:\n",
    "```python\n",
    "tensor([-0.0188, -0.1974,  0.0040,  0.0113,  0.0638, -0.1050,  0.0369, -0.1680,\n",
    "        -0.0491, -0.0565, -0.0025,  0.0135, -0.0042,  0.0151,  0.0166, -0.1381,\n",
    "        -0.0063, -0.0461,  0.0267, -0.2042])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can plot these weights and try to see what they represent like this:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "plt.imshow(huggingfaceStateDictionary['transformer.wpe.weight'], cmap='gray')\n",
    "```\n",
    "\n",
    "![GPT-2.transformer.wpe.weight](ExplanationMedia/Images/GPT-2.transformer.wpe.weight.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can see that this has structure, because these positional embeddings end up learning these **sinusoids** and **cosines** to represent each of these positions and each row here stands in for that position and is processed by the transformer to recover all the relative positions and realise which token is where and attend to them depending on their position not just their content...\n",
    "\n",
    "So now if we look at the individual columns of these we see:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "plt.plot(huggingfaceStateDictionary['transformer.wpe.weight'][:, 150])\n",
    "plt.plot(huggingfaceStateDictionary['transformer.wpe.weight'][:, 200])\n",
    "plt.plot(huggingfaceStateDictionary['transformer.wpe.weight'][:, 250])\n",
    "```\n",
    "![GPT-2Graphs.transformer.wpe.weight](ExplanationMedia/Images/GPT-2Graphs.transformer.wpe.weight.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we still don't know what these embeddings are doing and why they are the way they are.\n",
    "\n",
    "But we can still see that the lines are a little noisy and jittery and that is because this model was not fully trained, and the more trained this model becomes the more we'd expect these graphs to smooth out, which also tells us that the original `GPT-2` is an **under-trained** model.\n",
    "\n",
    "If I remember correctly, in the original \"Attention-Is-All-You-Need\" paper, the `positional embeddings` are actually initialized and fixed to sinusoids and cosines of different frequencies, but in `GPT-2` these are trained from scratch and they seem to recover these features during the optimization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, using the `Hugging Face Transformers` we can not just get all the raw weights but also get something called `pipeline` and sample from it...\n",
    "\n",
    "Here is the sample code snippet for `5` different generations of the same context window of tokens `\"Hello, I'm a language model,\"`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from transformers import pipeline, set_seed\n",
    "\n",
    "set_seed(42)\n",
    "\n",
    "generator = pipeline(\"text-generation\", model=\"gpt2\")\n",
    "generator(\"Hello, I'm a language model,\", max_length=50, num_return_sequences=5)\n",
    "```\n",
    "For which we get:\n",
    "```python\n",
    "[{'generated_text': \"Hello, I'm a language model, but what I'm really doing is making a human-readable document. There are other languages, but those are the ones I like the most. To do your research, please contact me, this isn't your\"},\n",
    " {'generated_text': \"Hello, I'm a language model, not a syntax model. That's why I like it. I've done a lot of programming projects.\\n\\nBut my job as a C programmer is to sort through every single line of the script so I\"},\n",
    " {'generated_text': \"Hello, I'm a language model, and I'll do it in no time!\\n\\nOne of the things we learned from talking to my friend from college a bit earlier, and in the context of the current language model I think it's important\"},\n",
    " {'generated_text': 'Hello, I\\'m a language model, not a command line tool.\\n\\nIf my code is simple enough:\\n\\nif (use (string-replace \"\\\\r\" ))) {\\n\\nconsole. log\\n\\n}\\n\\nthat\\'s'},\n",
    " {'generated_text': \"Hello, I'm a language model, I've been using Language in all my work. Just a small example, let's see a simplified example. I'm making an API for a game where I want a character to play a little bit of a\"}]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sadly, even though we are setting a seed we get different generations from both the code and the official <a href=\"https://huggingface.co/openai-community/gpt2\">Hugging Face GPT-2 Hosted Inference API</a>.\n",
    "\n",
    "But at this stage what is important is, we are getting coherent text and we were successfully able to load the model and look at all of it's parameters and the keys tell us, where in the model these come from...\n",
    "\n",
    "But we want to actually write our own `GPT-2` class so that we have a full understanding of what's happening there and we also don't want to work with something like the <a href=\"https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py\">`modeling_gpt2.py`</a> because it's too complicated and we want to write this from scratch ourselves.\n",
    "\n",
    "So we are going to be implementing our `GPT-2` model in `GPT_v2.5.py` script inside our `GPT Scripts` directory in parallel...\n",
    "\n",
    "But first let's load the `GPT-2 124M` into our `GPT_v2.5.py` for the class that we are going to develop from scratch, which is going to give us confidence that we can load the OpenAI model and there's a setting of weights that exactly is the `124M` model and we will try to surpass our own created `GPT` class...\n",
    "\n",
    "So, we're going to get different weights and everything is going to look different and hopefully even better and we will have the confidence that we are in the same model family and same model class and we just have to re-discover a good setting of the weights from scratch... \n",
    "\n",
    "So let's now write the `GPT-2` model and let's load the weights and make sure that we can also generate text that looks coherent..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hugging Face Pre-Trained Weight Transfer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer - Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now swing over to the <a href=\"https://arxiv.org/abs/1706.03762\">\"Attention Is All You Need\" paper</a> that started everything and look at the Transformer architecture:\n",
    "\n",
    "![Transformer_Model_Architecture](ExplanationMedia/Images/Transformer_Model_Architecture.png)\n",
    "\n",
    "Now, once again, like the last notebook, we mentioned that this architecture has changed over the years and `GPT-2` is slightly modified than the original `Transformer`... \n",
    "\n",
    "In particular, we do **NOT** have the **Encoder**, and `GPT-2` is a **Decoder** only `Transformer` as we call it. In addition to that the **Cross-Attention** that is used by that **Encoder** is also **missing**. Everything else stays almost the same, but there are some differences that we are going to see next..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, there are two main differences: \\\n",
    "When we go to the `GPT-2` <a href=\"https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf\">paper</a>, under section `2.3 Model` we see that there's a re-shuffling of the layer-normalizations (they change place) and an additional layer normalization was added after the final self-attention block..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPT-2 Skeleton"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyper-Parameters (`GPTConfiguration`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now implement the skeleton of the `nn.Module`(s) in our GPT Script and in particular we want to match up the schema that we got from `Hugging Face GPT-2`...\n",
    "\n",
    "And we will use a decorator called `@dataclass` which provides a decorator and functions for automatically adding generated special methods such as `__init__()` and `__repr__()` to user-defined classes...\n",
    "\n",
    "And we will use it to define all the hyper-parameters as a Class called `GPTConfiguration`...\n",
    "\n",
    "Now because we are going to be implementing the `124M GPT-2 Model`, when we go to the paper we see these hyper-parameters:\n",
    "1. block-size (context window) → 1024\n",
    "2. vocabulary-size (token vocabulary) → 50257\n",
    "3. n-layer (number of layers) → 12\n",
    "4. n-head (number of self-attention heads) → 12\n",
    "5. embedding-dimensions ($d_{model}$) → 768\n",
    "\n",
    "So let's implement this now..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For now our code looks like this:\n",
    "```python\n",
    "from dataclasses import dataclass\n",
    "\n",
    "# GPT configuration hyper-parameters\n",
    "@dataclass\n",
    "class GPTConfiguration:\n",
    "    blockSize: int = 1024\n",
    "    vocabularySize: int = 50257\n",
    "    numberOfLayers: int = 12\n",
    "    numberOfHeads: int = 12\n",
    "    numberOfEmbeddingDimensions: int = 768\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `GPTModel`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will be able to use this configuration under the `GPTModel` class that we are going to write...\n",
    "\n",
    "For now our empty `GPTModel` class looks like this:\n",
    "```python\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.configuration = configuration\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to **copy** the schema from the **Hugging Face `GPT-2` model** by utilizing the `huggingfaceStateDictionary`...\n",
    "\n",
    "And here's what I came up with...\n",
    "\n",
    "We see that the container in the schema is called `transformer` which contains all the modules and we can create something like that using `torch.nn.ModuleDict` which is just a dictionary of torch `Module`(s) which let's us index into `Module`(s) using **keys**, just like a normal python dictionary...\n",
    "\n",
    "Within that we can create something called `wordTokenEmbeddings` which corresponds with `wte` and create something called `wordPositionalEmbeddings` which corresponds with `wpe`, and we can match the shapes and create our initial layers...\n",
    "\n",
    "Then in the **Hugging Face `GPT-2` model** we see that we have a long list of **hidden** layers represented by a `.h` and followed by a range of number `.0` to `.11` hinting us about the number of layers as `12`, so we can now utilize our `numberOfLayers` hyper-parameter to construct these long list of layers. And instead of a `torch.nn.ModuleDict` we can use a `torch.nn.ModuleList` instead, which is just a list of `Module`(s).\n",
    "\n",
    "The important thing to note is, in those hidden layers we see different kinds of layer **weights** and **biases** of different **layers** all having their own shapes and sizes, but we do see a pattern that they repeat themselves in terms of **layer number**, so for now we can just consider these **layers** as `Block`(s) and iterate them through a list and return itself to the list. Keep in mind that the `Block`'s defination has not been defined yet, and we will define it later, but we want all the `Block`(s) to take in the same `configuration` object and construct the layer objects through it, because we already have all the hyper-parameters set inside it...\n",
    "\n",
    "Now that we have our long list of **hidden layers** it is time to construct the final **layer normalization** layer according to the `GPT-2` paper, so we can create something like `finalLayerNorm` and match the shapes which corresponds to the `ln_f`...\n",
    "\n",
    "And lastly we can construct our **final classifier** (or the **language model head**) which is just a **Linear Layer** that projects all the **embeddings** to their respective **tokens**, having **no bias**. So, we can easily construct this **languageModelingHead** which corresponds to the `lm_head` and finish with our skeleton of the `GPT-2` model...\n",
    "\n",
    "Now we end up with a code like this:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from dataclasses import dataclass\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "# GPT configuration hyper-parameters\n",
    "@dataclass\n",
    "class GPTConfiguration:\n",
    "    blockSize: int = 1024\n",
    "    vocabularySize: int = 50257\n",
    "    numberOfLayers: int = 12\n",
    "    numberOfHeads: int = 12\n",
    "    numberOfEmbeddingDimensions: int = 768\n",
    "\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.configuration = configuration\n",
    "\n",
    "        self.transformer = torch.nn.ModuleDict(dict(\n",
    "            wordTokenEmbeddings = torch.nn.Embedding(configuration.vocabularySize, configuration.numberOfEmbeddingDimensions),\n",
    "            wordPositionalEmbeddings = torch.nn.Embedding(configuration.blockSize, configuration.numberOfEmbeddingDimensions),\n",
    "            hidden = torch.nn.ModuleList(Block(configuration) for _ in range(configuration.numberOfLayers)),\n",
    "            finalLayerNorm = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        ))\n",
    "\n",
    "        self.languageModelingHead = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, configuration.vocabularySize, bias=False)\n",
    "\n",
    "model = GPTModel(GPTConfiguration())\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `Block`(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now create the `Block` class which is currently undefined...\n",
    "\n",
    "Now, here the `Block` refers to the `Transformer Block` that gets repeated again and again as hidden layers...\n",
    "\n",
    "Now, according to our <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/GPT%20from%20Scratch.ipynb\">GPT from Scratch</a> notebook, we already defined this `Block`, and as we mentioned `GPT-2` also has a slightly modified `Transformer Block`...\n",
    "\n",
    "For now we are going to use this template:\n",
    "```python\n",
    "# Transformer Block\n",
    "class Block(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        pass\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For, now we understand that we are left with the following modules & properties:\n",
    "1. Add & Norm\n",
    "2. Attention\n",
    "3. Feed Forward Network\n",
    "4. Residual Pathways\n",
    "\n",
    "Let's start with **Addition and Normalization (Add & Norm)**.\n",
    "\n",
    "According to the diagram, the **Add & Norm** is there **AFTER** the **Attention & Feed Forward Network**, but in our case we will use them **BEFORE** the **Attention & Feed Forward Network**, making it a **Pre-Normalization (Pre-Norm)**...\n",
    "\n",
    "Then we have our **Attention** module, and for now we can relate it to the attention we built in our last <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/GPT%20from%20Scratch.ipynb\">GPT from Scratch</a> notebook. Specifically, we built two modules (`MultiHeadAttention` and `Head`), but here we will implement both modules in a combined and mathematically optimized class called `CausalSelfAttention`...\n",
    "\n",
    "Then we have our **Feed Forward Network**, and we will call it our **MultiLayerPerceptron**.\n",
    "\n",
    "The thing to note is, once again two of the modules (`CausalSelfAttention` and `MultiLayerPerceptron`) remain undefined, and we are going to define them later...\n",
    "\n",
    "For now, we end up with a code like this:\n",
    "```python\n",
    "# Transformer Block\n",
    "class Block(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.layerNormalization1 = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        self.attention = CausalSelfAttention(configuration)\n",
    "        self.layerNormalization2 = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        self.multiLayerPerceptron = MultiLayerPerceptron(configuration)\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        pass\n",
    "```\n",
    "\n",
    "And lastly, we arrive at the **Residual Pathways**, and we see that the normalizations are **inside** the residual stream, or in other words, the residual pathway has normalizations **inside** them (which is not very good or desirable from an optimization perspective) and we actually prefer to have a single and clean residual stream all the way from **supervision** to all the way to the **inputs (or `tokens`)**, which is desirable because the gradients that flow from the top distributes the gradients equally because of additions, indicating that the gradients from the top flow straight to the inputs through the residual pathway (unchanged) but then addition to that, the gradient also flows through the blocks and the blocks contribute their own contribution over time when the optimization kicks in.\n",
    "\n",
    "Which means that we want to apply a **clean residual pathway**. And to do that we need the normalization to be applied to the residual pathway (result of layer normalization) **before** adding it back to the original **inputs**. This ensures that the residual connections are additive and do not interfere with the normalization process, facilitating better gradient flow and optimization stability.\n",
    "\n",
    "Therefore, we end up with a code like this:\n",
    "```python\n",
    "# Transformer Block\n",
    "class Block(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.layerNormalization1 = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        self.attention = CausalSelfAttention(configuration)\n",
    "        self.layerNormalization2 = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        self.multiLayerPerceptron = MultiLayerPerceptron(configuration)\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        inputs = inputs + self.attention(self.layerNormalization1(inputs))\n",
    "        inputs = inputs + self.multiLayerPerceptron(self.layerNormalization2(inputs))\n",
    "        return inputs\n",
    "```\n",
    "\n",
    "Here's how the residual pathways are structured:\n",
    "- **Attention Layer**: After applying `self.layerNormalization1`, the residual (inputs) are added to `attention_output`. This adheres to the clean residual pathway because the normalization (`layerNormalization1`) is applied **before** adding to inputs.\n",
    "- **MLP Layer**: Similarly, after applying `self.layerNormalization2`, the residual (inputs) is added to `mlp_output`. This also adheres to the clean residual pathway for the same reason as above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And one more thing that is interesting to note is that, **Attention** is a **communication operation**, and it is where all the tokens line-up in a sequence and this is where the tokens communicate and exchange information. And **Attention** is an **aggregation function**, it's a **pooling function**, it's a **weighted sum function**, it's a **reduce operation**.\n",
    "\n",
    "Whereas, **Multi Layer Perceptron** happens at every single token individually (**mapped**), and there is no information being exchanged or collected between the tokens.\n",
    "\n",
    "So, the **Attention** is the **reduce** and **Multi Layer Perceptron** is the **map**. And what we end up with is a repeated application of **Map-Reduce**. And this is where the `tokens` communicate and this is where they *think* individually about the information that they gathered. And every one of these blocks, iteratively refines the representation inside the residual stream...\n",
    "\n",
    "And now we can move on the to implementation of `CausalSelfAttention` and `MultiLayerPerceptron`..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `MultiLayerPerceptron`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now move on to the `MultiLayerPerceptron (MLP)`, and I implemented the class as follows...\n",
    "\n",
    "```python\n",
    "# Multi Layer Perceptron (MLP)\n",
    "class MultiLayerPerceptron(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, inputs):\n",
    "        pass\n",
    "```\n",
    "\n",
    "It is relatively straight forward. For now, we just have two `Linear` layers which wrap around a `GELU` non-linearity layer. So our block now becomes something like this:\n",
    "```python\n",
    "# Multi Layer Perceptron (MLP)\n",
    "class MultiLayerPerceptron(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.currentFullyConnected = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, 4 * configuration.numberOfEmbeddingDimensions)\n",
    "        self.gelu = torch.nn.GELU(approximate=\"tanh\")\n",
    "        self.currentProjection = torch.nn.Linear(4 * configuration.numberOfEmbeddingDimensions, configuration.numberOfEmbeddingDimensions)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        inputs = self.currentFullyConnected(inputs)\n",
    "        inputs = self.gelu(inputs)\n",
    "        inputs = self.currentProjection(inputs)\n",
    "        return inputs\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now when we swing over to the <a href=\"https://pytorch.org/docs/stable/generated/torch.nn.GELU.html\">GELU PyTorch Documentation</a>, we see **two** different `GELU`(s) being hinted there:\n",
    "1. Original GELU formulation (We will discuss this in a bit):\n",
    "   $$ \\text{GELU}(x) = x \\cdot \\Phi(x) $$\n",
    "2. Approximate GELU formulation:\n",
    "   $$ \\text{GELU}(x) = 0.5 \\cdot x \\cdot \\left(1 + \\tanh \\left( \\frac{2}{\\pi} \\cdot \\left(x + 0.044715 \\cdot x^3 \\right) \\right) \\right) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![GELU](ExplanationMedia/Images/GELU.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just as a preview, we can see that `GELU` is basically like a `ReLU`, except there's **no exactly flat tail at exactly `0`**. Otherwise, it just looks more like a slightly *smoother* `ReLU`. And it comes from this paper <a href=\"https://arxiv.org/abs/1606.08415\">\"Gaussian Error Linear Units (GELUs)\"</a> and there's a little bit of history here and I also invite you to step through the paper if you'd like. But for now, we will use the **approximate** version of the `GELU`, because that's what `GPT-2` in their model used..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, one other reason of why we prefer to use `GELU` is that, in previous notebooks we have spoken about the **Dead-ReLU-Neuron-Problem** where, in the tail of a `ReLU`, where it's exactly flat at `0`, any activations that fall there will get exactly `0` gradient (meaning that there's no change, there's no adaptation, there's no development of the network), but `GELU` always contributes to a **local-gradient** and so there's always going to be a change and there's always going to be an adaptation in a *smoothed-out* way which empirically working better, as demonstrated in the paper.\n",
    "\n",
    "And we also followed the rule of *\"Position-wise Feed-Forward Networks\"* section of the original \"Attention-is-all-you-need\" paper, which is why we have the `4 * numberOfEmbeddingDimensions` in the shapes..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally we can now move on to implement the `CausalSelfAttention` part of the code..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `CausalSelfAttention`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now start implementing our `CausalSelfAttention` block which is none other than the combination of **Scaled Dot-Product Attention** and **Multi-Head Attention**...\n",
    "\n",
    "For now we have a skeleton like this:\n",
    "```python\n",
    "# Causal Self Attention (Scaled Dot-Product Attention + Multi-Head Attention)\n",
    "class CausalSelfAttention(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        pass\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember from our <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/GPT%20from%20Scratch.ipynb\">GPT from Scratch notebook</a>, that **Multi-Head Attention** is just multiple **Scaled Dot-Product Attention**(s) running in parellel and their outputs are just being concatenated and that becomes the output.\n",
    "\n",
    "Instead, we do a bunch of tensor *gymnastics* of mathematical operations of the same logic used behind both these **Multi-Head Attention** & **Scaled Dot-Product Attention** modules in a single block. But fundamentally, and algorithmically, nothing is different from what we implemented previously...\n",
    "\n",
    "And this is what we end up with:\n",
    "```python\n",
    "# Causal Self Attention (Scaled Dot-Product Attention + Multi-Head Attention)\n",
    "class CausalSelfAttention(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        assert configuration.numberOfEmbeddingDimensions % configuration.numberOfHeads == 0\n",
    "        self.causalAttention = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, 3 * configuration.numberOfEmbeddingDimensions)\n",
    "        self.causalProjection = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, configuration.numberOfEmbeddingDimensions)\n",
    "        self.numberOfHeads = configuration.numberOfHeads\n",
    "        self.numberOfEmbeddingDimensions = configuration.numberOfEmbeddingDimensions\n",
    "        self.register_buffer(\"bias\", torch.tril(torch.ones(configuration.blockSize, configuration.blockSize)).view(1, 1, configuration.blockSize, configuration.blockSize))\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        B, T, C = inputs.size()\n",
    "        query_key_value = self.causalAttention(inputs)\n",
    "        query, key, value = query_key_value.split(self.numberOfEmbeddingDimensions, dim=2)\n",
    "        query = query.view(B, T, self.numberOfHeads, C // self.numberOfHeads).transpose(1, 2)\n",
    "        key = key.view(B, T, self.numberOfHeads, C // self.numberOfHeads).transpose(1, 2)\n",
    "        value = value.view(B, T, self.numberOfHeads, C // self.numberOfHeads).transpose(1, 2)\n",
    "        attention = (query @ key.transpose(-2, -1)) * (1.0 / math.sqrt(key.size(-1)))\n",
    "        attention = attention.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))\n",
    "        attention = F.softmax(attention, dim=-1)\n",
    "        outputs = attention @ value\n",
    "        outputs = outputs.transpose(1, 2).contiguous().view(B, T, C)\n",
    "        outputs = self.causalProjection(outputs)\n",
    "        return outputs\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPT-2 Weights Transfer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the skeleton ready, we can now move on to transfer the weights of the `Hugging Face GPT-2` to our `Custom GPT-2`..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `from_pretrained()` - Skeleton"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by understanding what we want to do first...\n",
    "\n",
    "We want to have a `from_pretrained()` method in our `GPTModel` class, that will transfer the weights for any kind of model we pass it (among the `4` models that are there in `GPT-2`), and copy the weights of each of those parameters and ensure their sizes and shapes match perfectly...\n",
    "\n",
    "And we also want our `from_pretrained()` method to be decorated by a `@classmethod` such that it could be accessed directly using the class reference and it is also able to modify the state of the class and return the appropriate model along with their appropriate parameters...\n",
    "\n",
    "So for now we can have a skeleton like this:\n",
    "```python\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        ... # Rest of the code\n",
    "\n",
    "    # Method to transfer weights from Hugging Face GPT-2\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, modelType):\n",
    "        assert modelType in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}\n",
    "        from transformers import GPT2LMHeadModel\n",
    "        print(\"Loading weights from pretrained gpt: %s\" % modelType)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `from_pretrained()` - Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create the separate configuartion arguements for all `4` `GPT-2` configurations...\n",
    "\n",
    "And our code will look like this:\n",
    "```python\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        ... # Rest of the code\n",
    "\n",
    "    # Method to transfer weights from Hugging Face GPT-2\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, modelType):\n",
    "        assert modelType in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}\n",
    "        from transformers import GPT2LMHeadModel\n",
    "        print(\"Loading weights from pretrained gpt: %s\" % modelType)\n",
    "\n",
    "        # Creating separate configurations for separate GPT-2 models\n",
    "        configurationArguements = {\n",
    "            'gpt2':         dict(numberOfLayers=12, numberOfHeads=12, numberOfEmbeddingDimensions=768),  # 124M parameters\n",
    "            'gpt2-medium':  dict(numberOfLayers=24, numberOfHeads=16, numberOfEmbeddingDimensions=1024), # 350M parameters\n",
    "            'gpt2-large':   dict(numberOfLayers=36, numberOfHeads=20, numberOfEmbeddingDimensions=1280), # 774M parameters\n",
    "            'gpt2-xl':      dict(numberOfLayers=48, numberOfHeads=25, numberOfEmbeddingDimensions=1600), # 1558M parameters\n",
    "        }[modelType]\n",
    "        configurationArguements['vocabularySize'] = 50257\n",
    "        configurationArguements['blockSize'] = 1024\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can unpack our configurations into a variable called `configuration` based on the `modelType` arguement we pass as a parameter. And then we can initialize our `GPTModel` based on the `configuration` that we initialize. And then we can also copy the state-dictionary containing all the layers in our model in a variable called `stateDictionary` with the method `state_dict()` and unpack it's keys using the `keys()` method into a variable called `stateDictionaryKeys` (We have to keep in mind that we discard all the buffers that are not a part of the parameters like **Attention Mask** and **Attention Bias**)...\n",
    "\n",
    "So, now we have a code like this:\n",
    "```python\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        ... # Rest of the code\n",
    "\n",
    "    # Method to transfer weights from Hugging Face GPT-2\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, modelType):\n",
    "        assert modelType in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}\n",
    "        from transformers import GPT2LMHeadModel\n",
    "        print(\"Loading weights from pretrained gpt: %s\" % model_type)\n",
    "\n",
    "        # Creating separate configurations for separate GPT-2 models\n",
    "        configurationArguements = {\n",
    "            'gpt2':         dict(numberOfLayers=12, numberOfHeads=12, numberOfEmbeddingDimensions=768),  # 124M parameters\n",
    "            'gpt2-medium':  dict(numberOfLayers=24, numberOfHeads=16, numberOfEmbeddingDimensions=1024), # 350M parameters\n",
    "            'gpt2-large':   dict(numberOfLayers=36, numberOfHeads=20, numberOfEmbeddingDimensions=1280), # 774M parameters\n",
    "            'gpt2-xl':      dict(numberOfLayers=48, numberOfHeads=25, numberOfEmbeddingDimensions=1600), # 1558M parameters\n",
    "        }[modelType]\n",
    "        configurationArguements['vocabularySize'] = 50257\n",
    "        configurationArguements['blockSize'] = 1024\n",
    "\n",
    "        configuration = GPTConfiguration(**configurationArguements)\n",
    "        model = GPTModel(configuration)\n",
    "\n",
    "        stateDictionary = model.state_dict()\n",
    "        stateDictionaryKeys = stateDictionary.keys()\n",
    "        stateDictionaryKeys = [key for key in stateDictionaryKeys if not key.endswith('.attention.bias')]\n",
    "\n",
    "        return model\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `from_pretrained()` - Hugging Face Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our own custom model completely initialized, we can go ahead and initialize the `Hugging Face GPT-2 Model` into a single variable and it's state dictionary and keys into other variables called `huggingfaceStateDictionary` and `huggingfaceStateDictionaryKeys`...\n",
    "\n",
    "And then we can start copying the weights after **ignoring the buffers**. But now, before copying we have to keep in mind that the original code for the `GPT-2` model was trained using the `TensorFlow` library and some of the weights are **transposed** in that architecture, so we will manually hard-code those weights and **copy them after transposing them to their original PyTorch form**...\n",
    "\n",
    "One last thing to be careful about is, in our model we are using **custom names for our variables** but the `Hugging Face GPT-2 Model` has an architecture of **short forms**, so it is better to have a `parameterKeyMapping` that **maps our custom keys with the Hugging Face GPT-2 keys of the state-dictionary** such that it becomes much easier to iterate through...\n",
    "\n",
    "So, now we have a final code like this:\n",
    "```python\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        ... # Rest of the code\n",
    "\n",
    "    # Method to transfer weights from Hugging Face GPT-2\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, modelType):\n",
    "        assert modelType in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}\n",
    "        from transformers import GPT2LMHeadModel\n",
    "        print(\"Loading weights from pretrained gpt: %s\" % modelType)\n",
    "        \n",
    "        # Creating separate configurations for separate GPT-2 models\n",
    "        blockSize = 1024\n",
    "        vocabularySize = 50257\n",
    "        configurationArguements = {\n",
    "            'gpt2':         dict(numberOfLayers=12, numberOfHeads=12, numberOfEmbeddingDimensions=768),  # 124M parameters\n",
    "            'gpt2-medium':  dict(numberOfLayers=24, numberOfHeads=16, numberOfEmbeddingDimensions=1024), # 350M parameters\n",
    "            'gpt2-large':   dict(numberOfLayers=36, numberOfHeads=20, numberOfEmbeddingDimensions=1280), # 774M parameters\n",
    "            'gpt2-xl':      dict(numberOfLayers=48, numberOfHeads=25, numberOfEmbeddingDimensions=1600), # 1558M parameters\n",
    "        }[modelType]\n",
    "        configurationArguements['vocabularySize'] = 50257\n",
    "        configurationArguements['blockSize'] = 1024\n",
    "\n",
    "        configuration = GPTConfiguration(**configurationArguements)\n",
    "        model = GPTModel(configuration)\n",
    "        stateDictionary = model.state_dict()\n",
    "        stateDictionaryKeys = stateDictionary.keys()\n",
    "        stateDictionaryKeys = [key for key in stateDictionaryKeys if not key.endswith('.attention.bias')]\n",
    "\n",
    "        huggingfaceModel = GPT2LMHeadModel.from_pretrained(modelType)\n",
    "        huggingfaceStateDictionary = huggingfaceModel.state_dict()\n",
    "        huggingfaceStateDictionaryKeys = huggingfaceStateDictionary.keys()\n",
    "        huggingfaceStateDictionaryKeys = [key for key in huggingfaceStateDictionaryKeys if not key.endswith('.attn.masked_bias')]\n",
    "        huggingfaceStateDictionaryKeys = [key for key in huggingfaceStateDictionaryKeys if not key.endswith('.attn.bias')]\n",
    "        transposedParameters = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']\n",
    "\n",
    "        assert len(huggingfaceStateDictionaryKeys) == len(stateDictionaryKeys), f\"Mismatched Keys: {len(huggingfaceStateDictionaryKeys)} != {len(stateDictionaryKeys)}\"\n",
    "\n",
    "        parameterKeyMapping = {\n",
    "            customKey: huggingfaceKey\n",
    "            for customKey, huggingfaceKey in zip(stateDictionaryKeys, huggingfaceStateDictionaryKeys)\n",
    "            }\n",
    "\n",
    "        for customKey, huggingfaceKey in parameterKeyMapping.items():\n",
    "            if (huggingfaceStateDictionary[huggingfaceKey].shape != stateDictionary[customKey].shape):\n",
    "                # Special treatment for the Conv1D weights (Transposed Weights)\n",
    "                if (huggingfaceKey.endswith(word) for word in transposedParameters):\n",
    "                    assert huggingfaceStateDictionary[huggingfaceKey].shape[::-1] == stateDictionary[customKey].shape\n",
    "                    with torch.no_grad():\n",
    "                        stateDictionary[customKey].copy_(huggingfaceStateDictionary[huggingfaceKey].t())\n",
    "            # Vanilla copy for other parameters\n",
    "            else:\n",
    "                assert huggingfaceStateDictionary[huggingfaceKey].shape == stateDictionary[customKey].shape\n",
    "                with torch.no_grad():\n",
    "                    stateDictionary[customKey].copy_(huggingfaceStateDictionary[huggingfaceKey])\n",
    "        \n",
    "        return model\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As soon as we do this, you will see that it will start downloading the model from Hugging Face like this:\n",
    "```bash\n",
    "Loading weights from pretrained gpt: gpt2\n",
    "config.json: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 665/665 [00:00<?, ?B/s]\n",
    "model.safetensors:  34%|████████████████████████████████████████████████████▋                                                                                                    | 189M/548M [00:18<00:35, 10.1MB/s]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once it's done we can now move on to the generation phase of the model..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPT-2 Forward Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, before we can generate from this model that we have implemented, we need a forward function that forwards the token sequence through the model to get the logits...\n",
    "\n",
    "Now the inputs are none other than the `token indeces` that is represented in a shape of `(Batch, Time)` tensor, where `Batch` dimension is the **independent sequences** and `Time` dimension is the **maximum sequence length**. Meaning that we have inputs in the shape of a matrix with each row having independent sequences of a maximum length of a sequence.\n",
    "\n",
    "So, we first unpack the `Batch` and `Time` dimensions in a variable. Then we forward the `positional embeddings` by creating a different tensor as `tokenPositions` which is none other than the `(0, Time)` dimension of the `token indeces`...\n",
    "\n",
    "Then we forward the `positional embeddings` and the `token embeddings` and when we get their respective outputs, we will concatenate them in a variable called `inputs`.\n",
    "\n",
    "And lastly, we will loop through every block of the `transformer` and forward the `inputs` through them, and finally forward them through the final `layer normalization` and `language modeling head` to get the logits...\n",
    "\n",
    "So now we have a forward loop inside the `GPTModel` class like this:\n",
    "```python\n",
    "def forward(self, indeces):\n",
    "    Batch, Time = indeces.size()\n",
    "    assert T <= self.configuration.blockSize, f\"Cannot forward sequence of length {Time}, Block Size is only {self.configuration.blockSize}\"\n",
    "\n",
    "    tokenPositions = torch.arange(0, Time, dtype=torch.long)\n",
    "    positionalEmbeddings = self.transformer.wordPositionalEmbeddings(tokenPositions)\n",
    "    tokenEmbeddings = self.transformer.wordTokenEmbeddings(indeces)\n",
    "    inputs = tokenEmbeddings + positionalEmbeddings\n",
    "\n",
    "    for block in self.transformer.hidden:\n",
    "        inputs = block(x)\n",
    "    inputs = self.transformer.finalLayerNorm(inputs)\n",
    "    logits = self.languageModelingHead(inputs)\n",
    "    return logits\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also want to use this model on the `GPU`, so for now we can change the `tokenPositions` that we create to and specify the correct device like this:\n",
    "`tokenPositions = torch.arange(0, Time, dtype=torch.long, device=indeces.device)`.\n",
    "\n",
    "And then we can move on to the generation phase of the model..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPT-2 Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we generate from the model, we need to think through what we will generate from the model and what will go through the inputs...\n",
    "\n",
    "We will firstly put the model into `eval()` mode because the even though we don't have layers that uses different mechanisms during training time and inference time, it is a good practice to keep the model change it's state if we do any further changes to the model.\n",
    "\n",
    "Well, we will forward `encoded tokens` into the model and get `encoded tokens` as output which we need to decode again to see the generation. Now, these tokens will have two hyper-parameters that we will define, one being the `maximum generation length` (denotes how much tokens will the model generate for each independant token sequence) and one being the `number of sequences to generate` (denotes the number of sequences we are trying to generate in a single run in parallel).\n",
    "\n",
    "You might be wondering \"where suddenly these `encoded tokens` are coming from?\".\n",
    "\n",
    "To answer this, we can go back to the <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/GPT%20Tokenizer.ipynb\">GPT Tokenizer</a> notebook and refer to it to understand that we will use a library to encode and decode tokens called `tiktoken`...\n",
    "\n",
    "And we will encode the same sequence:\n",
    "> \"Hello, I'm a language model,\"\n",
    "\n",
    "And if we go to <a href=\"https://tiktokenizer.vercel.app/?model=gpt2\">tiktokenizer</a> and see the token sequence size, we will see that there are `8` tokens in this sequence...\n",
    "\n",
    "So, we can now create a tensor called `tokens` which will be of shape `(numberOfSequences, numberOfTokensInSequence)` (Let's say we want to generate `5` sequences, then the shape of this `tokens` tensor will be of shape `(5, 8)`).\n",
    "\n",
    "Then we will create loop with `maximumGenerationLength` and forward the `encoded tokens` to get the `logits`. Now, because the `logits` will be of shape `(Batch, Time, Channel)` where `Batch` represents the **independent sequences**, `Time` represents the **tokens in a sequence** and `Channel` represents the **vocabulary that `logits` will classify into**, we will select only the **last** generated `Time` dimension from the `logits` tensor and pass them through a `softmax()` to get the `probabilites`...\n",
    "\n",
    "Now, by default the Hugging Face pipeline uses **top-k probabilites of 50 by default**, so, we will also implement this in our loop by using PyTorch's `topk()` method and pass in our probabilites and specify the correct dimension. This `topk()` method will then return the final probabilites (`topKProbabilites`) and the indeces of these probabilites(`tokKIndeces`). **This helps us to never sample very rare tokens.**\n",
    "\n",
    "Then we can sample a token from the `multinomial distribution` of these `topKProbabilites` and get the correct `tokenIndeces` (which is a single sampled `token` in a batch). Then we can create a column of `tokenIndeces` using <a href=\"https://pytorch.org/docs/stable/generated/torch.gather.html\">`torch.gather()`</a> and append them on the original `tokens` to have our generated tokens in a sequence...\n",
    "\n",
    "We will also keep this entire section inside the loop into `torch.no_grad()` to let PyTorch know that we won't be needing any backward processing (gradient calculation and intermediate operation caching) to save us some memory and hopefully some time...\n",
    "\n",
    "And finally, when we have our `encoded tokens in a sequence` which will be in the shape of `(numberOfSequences, maximumGenerationLength)` and we can decode them using a loop based on `numberOfSequences`..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we end up with a final code like this:\n",
    "```python\n",
    "from dataclasses import dataclass\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "import tiktoken\n",
    "\n",
    "# Causal Self Attention (Scaled Dot-Product Attention + Multi-Head Attention)\n",
    "class CausalSelfAttention(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        assert configuration.numberOfEmbeddingDimensions % configuration.numberOfHeads == 0\n",
    "        self.causalAttention = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, 3 * configuration.numberOfEmbeddingDimensions)\n",
    "        self.causalProjection = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, configuration.numberOfEmbeddingDimensions)\n",
    "        self.numberOfHeads = configuration.numberOfHeads\n",
    "        self.numberOfEmbeddingDimensions = configuration.numberOfEmbeddingDimensions\n",
    "        self.register_buffer(\"bias\", torch.tril(torch.ones(configuration.blockSize, configuration.blockSize)).view(1, 1, configuration.blockSize, configuration.blockSize))\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        B, T, C = inputs.size()\n",
    "        query_key_value = self.causalAttention(inputs)\n",
    "        query, key, value = query_key_value.split(self.numberOfEmbeddingDimensions, dim=2)\n",
    "        query = query.view(B, T, self.numberOfHeads, C // self.numberOfHeads).transpose(1, 2)\n",
    "        key = key.view(B, T, self.numberOfHeads, C // self.numberOfHeads).transpose(1, 2)\n",
    "        value = value.view(B, T, self.numberOfHeads, C // self.numberOfHeads).transpose(1, 2)\n",
    "        attention = (query @ key.transpose(-2, -1)) * (1.0 / math.sqrt(key.size(-1)))\n",
    "        attention = attention.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))\n",
    "        attention = F.softmax(attention, dim=-1)\n",
    "        outputs = attention @ value\n",
    "        outputs = outputs.transpose(1, 2).contiguous().view(B, T, C)\n",
    "        outputs = self.causalProjection(outputs)\n",
    "        return outputs\n",
    "\n",
    "# Multi Layer Perceptron (MLP)\n",
    "class MultiLayerPerceptron(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.currentFullyConnected = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, 4 * configuration.numberOfEmbeddingDimensions)\n",
    "        self.gelu = torch.nn.GELU(approximate=\"tanh\")\n",
    "        self.currentProjection = torch.nn.Linear(4 * configuration.numberOfEmbeddingDimensions, configuration.numberOfEmbeddingDimensions)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        inputs = self.currentFullyConnected(inputs)\n",
    "        inputs = self.gelu(inputs)\n",
    "        inputs = self.currentProjection(inputs)\n",
    "        return inputs\n",
    "\n",
    "\n",
    "# Transformer Block\n",
    "class Block(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.layerNormalization1 = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        self.attention = CausalSelfAttention(configuration)\n",
    "        self.layerNormalization2 = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        self.multiLayerPerceptron = MultiLayerPerceptron(configuration)\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        inputs = inputs + self.attention(self.layerNormalization1(inputs))\n",
    "        inputs = inputs + self.multiLayerPerceptron(self.layerNormalization2(inputs))\n",
    "        return inputs\n",
    "\n",
    "# GPT configuration hyper-parameters\n",
    "@dataclass\n",
    "class GPTConfiguration:\n",
    "    blockSize: int = 1024\n",
    "    vocabularySize: int = 50257\n",
    "    numberOfLayers: int = 12\n",
    "    numberOfHeads: int = 12\n",
    "    numberOfEmbeddingDimensions: int = 768\n",
    "\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.configuration = configuration\n",
    "\n",
    "        self.transformer = torch.nn.ModuleDict(dict(\n",
    "            wordTokenEmbeddings = torch.nn.Embedding(configuration.vocabularySize, configuration.numberOfEmbeddingDimensions),\n",
    "            wordPositionalEmbeddings = torch.nn.Embedding(configuration.blockSize, configuration.numberOfEmbeddingDimensions),\n",
    "            hidden = torch.nn.ModuleList(Block(configuration) for _ in range(configuration.numberOfLayers)),\n",
    "            finalLayerNorm = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        ))\n",
    "\n",
    "        self.languageModelingHead = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, configuration.vocabularySize, bias=False)\n",
    "    \n",
    "    def forward(self, indeces):\n",
    "        Batch, Time = indeces.size()\n",
    "        assert Time <= self.configuration.blockSize, f\"Cannot forward sequence of length {Time}, Block Size is only {self.configuration.blockSize}\"\n",
    "\n",
    "        tokenPositions = torch.arange(0, Time, dtype=torch.long, device=indeces.device)\n",
    "        positionalEmbeddings = self.transformer.wordPositionalEmbeddings(tokenPositions)\n",
    "        tokenEmbeddings = self.transformer.wordTokenEmbeddings(indeces)\n",
    "        inputs = tokenEmbeddings + positionalEmbeddings\n",
    "\n",
    "        for block in self.transformer.hidden:\n",
    "            inputs = block(inputs)\n",
    "        inputs = self.transformer.finalLayerNorm(inputs)\n",
    "        logits = self.languageModelingHead(inputs)\n",
    "        return logits\n",
    "\n",
    "    # Method to transfer weights from Hugging Face GPT-2\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, modelType):\n",
    "        assert modelType in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}\n",
    "        from transformers import GPT2LMHeadModel\n",
    "        print(\"Loading weights from pretrained GPT: %s\" % modelType)\n",
    "        \n",
    "        # Creating separate configurations for separate GPT-2 models\n",
    "        blockSize = 1024\n",
    "        vocabularySize = 50257\n",
    "        configurationArguements = {\n",
    "            'gpt2':         dict(numberOfLayers=12, numberOfHeads=12, numberOfEmbeddingDimensions=768),  # 124M parameters\n",
    "            'gpt2-medium':  dict(numberOfLayers=24, numberOfHeads=16, numberOfEmbeddingDimensions=1024), # 350M parameters\n",
    "            'gpt2-large':   dict(numberOfLayers=36, numberOfHeads=20, numberOfEmbeddingDimensions=1280), # 774M parameters\n",
    "            'gpt2-xl':      dict(numberOfLayers=48, numberOfHeads=25, numberOfEmbeddingDimensions=1600), # 1558M parameters\n",
    "        }[modelType]\n",
    "        configurationArguements['vocabularySize'] = 50257\n",
    "        configurationArguements['blockSize'] = 1024\n",
    "\n",
    "        configuration = GPTConfiguration(**configurationArguements)\n",
    "        model = GPTModel(configuration)\n",
    "        stateDictionary = model.state_dict()\n",
    "        stateDictionaryKeys = stateDictionary.keys()\n",
    "        stateDictionaryKeys = [key for key in stateDictionaryKeys if not key.endswith('.attention.bias')]\n",
    "\n",
    "        huggingfaceModel = GPT2LMHeadModel.from_pretrained(modelType)\n",
    "        huggingfaceStateDictionary = huggingfaceModel.state_dict()\n",
    "        huggingfaceStateDictionaryKeys = huggingfaceStateDictionary.keys()\n",
    "        huggingfaceStateDictionaryKeys = [key for key in huggingfaceStateDictionaryKeys if not key.endswith('.attn.masked_bias')]\n",
    "        huggingfaceStateDictionaryKeys = [key for key in huggingfaceStateDictionaryKeys if not key.endswith('.attn.bias')]\n",
    "        transposedParameters = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']\n",
    "\n",
    "        assert len(huggingfaceStateDictionaryKeys) == len(stateDictionaryKeys), f\"Mismatched Keys: {len(huggingfaceStateDictionaryKeys)} != {len(stateDictionaryKeys)}\"\n",
    "\n",
    "        parameterKeyMapping = {\n",
    "            customKey: huggingfaceKey\n",
    "            for customKey, huggingfaceKey in zip(stateDictionaryKeys, huggingfaceStateDictionaryKeys)\n",
    "            }\n",
    "\n",
    "        for customKey, huggingfaceKey in parameterKeyMapping.items():\n",
    "            if (huggingfaceStateDictionary[huggingfaceKey].shape != stateDictionary[customKey].shape):\n",
    "                # Special treatment for the Conv1D weights (Transposed Weights)\n",
    "                if (huggingfaceKey.endswith(word) for word in transposedParameters):\n",
    "                    assert huggingfaceStateDictionary[huggingfaceKey].shape[::-1] == stateDictionary[customKey].shape\n",
    "                    with torch.no_grad():\n",
    "                        stateDictionary[customKey].copy_(huggingfaceStateDictionary[huggingfaceKey].t())\n",
    "            # Vanilla copy for other parameters\n",
    "            else:\n",
    "                assert huggingfaceStateDictionary[huggingfaceKey].shape == stateDictionary[customKey].shape\n",
    "                with torch.no_grad():\n",
    "                    stateDictionary[customKey].copy_(huggingfaceStateDictionary[huggingfaceKey])\n",
    "        return model\n",
    "\n",
    "model = GPTModel.from_pretrained('gpt2')\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "model.eval()\n",
    "model.to(device=device)\n",
    "\n",
    "# Generation\n",
    "maximumGenerationLength = 30\n",
    "numberOfSequences = 5\n",
    "\n",
    "encoder = tiktoken.get_encoding('gpt2')\n",
    "encodedTokens = encoder.encode(\"Hello, I'm a language model,\")\n",
    "encodedTokens = torch.tensor(encodedTokens, dtype=torch.long)\n",
    "encodedTokens = encodedTokens.unsqueeze(0).repeat(numberOfSequences, 1)\n",
    "inputs = encodedTokens.to(device=device)\n",
    "\n",
    "torch.manual_seed(42)\n",
    "torch.cuda.manual_seed(42)\n",
    "\n",
    "while inputs.size(1) < maximumGenerationLength:\n",
    "    with torch.no_grad():\n",
    "        logits = model(inputs)\n",
    "        logits = logits[:, -1, :]\n",
    "        probabilites = F.softmax(logits, dim=-1)\n",
    "\n",
    "        topKProbabilites, tokKIndeces = torch.topk(input=probabilites, k=50, dim=-1)\n",
    "\n",
    "        tokenIndeces = torch.multinomial(input=topKProbabilites, num_samples=1)\n",
    "        columnOfTokenIndeces = torch.gather(input=tokKIndeces, dim=-1, index=tokenIndeces)\n",
    "\n",
    "        inputs = torch.cat((inputs, columnOfTokenIndeces), dim=1)\n",
    "\n",
    "for i in range(numberOfSequences):\n",
    "    tokensToDecode = inputs[i, :maximumGenerationLength].tolist()\n",
    "    decodedTokens = encoder.decode(tokensToDecode)\n",
    "    print(\">\", decodedTokens)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we get the following generations:\n",
    "```plaintext\n",
    "Loading weights from pretrained GPT: gpt2\n",
    "> Hello, I'm a language model, as's the - a. and for they.. were also. is of -\n",
    " to/ ' can\n",
    "> Hello, I'm a language model, I (, the\n",
    " ( have \" to ( \" of are\n",
    "., ' the'sa. that\n",
    "> Hello, I'm a language model, - on the<|endoftext|> was will also's't: of or<|endoftext|>.. are is to he, is the\n",
    "> Hello, I'm a language model, or \"\n",
    " a will. will the the.. The and in.,- ofThe's- for\n",
    "> Hello, I'm a language model,.. willThe's. and was and, was I would of a's his's's's. or\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, comes the interesting part... We want to initialize everything from scratch... We don't want to use any of these weights, and we want to use random numbers and initialize them and train them and generate from them...\n",
    "\n",
    "So let's now move on to the next part of this notebook..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Moving Model to GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, in case you do not have a `GPU` available, you can still follow along the notebook to some extent, but probably not to the very end because we will be actually using multiple `GPU`(s) and an actually perform a serious training run, but for now you can actually follow along with the notebook...\n",
    "\n",
    "And the one thing that I'd like to do is to *auto-detect* the **device** that is available to you and run the code on the highest compute capability...\n",
    "\n",
    "And you can do that with a code like this:\n",
    "```python\n",
    "device = \"cpu\"\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "elif hasattr(torch.backends, \"mps\") and torch.backends.mps.is_available():\n",
    "    device = \"mps\"\n",
    "print(f\"Using Device: {device}\")\n",
    "```\n",
    "\n",
    "We see that by default, the device is the `CPU` which is available everywhere, but then, we can detect the `GPU` using `CUDA`, and then if we don't have a `CUDA` we can detect if it atleast has `MPS` which is the backend for `Apple Silicon` (Newer Macbook Models)...\n",
    "\n",
    "And once we have this `device` we can potentially use this in the model:\n",
    "```python\n",
    "model.to(device=device)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you'll remember that we used `tokenPositions` by carefully setting the device to `device=indeces.device` in the `GPTModel`, and we did this to carefully set the location of initialization of this tensor to the correct device to **prevent the device mis-match**...\n",
    "\n",
    "Now, I do want to loop back around to this section to 'what it means to have different devices in PyTorch, and what it is exactly that PyTorch does in the background when we do something like `model.to(device=device)` and how it works', but for now we'd like to get to training of this model and we'd like to start training the model and for now let's just say the `device` makes the code go fast..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing Dataset and Encoding Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It turns out that we are initializing our weights at random already because PyTorch already initializes our layers randomly and by default...\n",
    "\n",
    "Right now for the Hugging Face model initialization we are using the code:\n",
    "```python\n",
    "model = GPTModel.from_pretrained('gpt2')\n",
    "```\n",
    "But if we want to use our default initialization we can use our older code:\n",
    "```python\n",
    "model = GPTModel(GPTConfiguration())\n",
    "```\n",
    "\n",
    "And for now if we try to run our code it *blabbers* garbage like this:\n",
    "```plaintext\n",
    "Using Device: cpu\n",
    "> Hello, I'm a language model,FCConnect Sandwich 64 Seed SHARparam Bloodyivil Sketch arrang Deaths backdoor Steeledoorccording bathingParentiven revers cafeteria trustees\n",
    "> Hello, I'm a language model, Brenblescler Awakens55 collar foe COUNboarding70710erry Evidence promotionMeet Icononticient Copobyl survivor Advoc Gro\n",
    "> Hello, I'm a language model, asylumacaninventoryQuantity physician OHBillyhirt controlling doctrines Summers wallet disdain Test repercussions Nighthew Goblinbreeding flight amuse Most paradox\n",
    "> Hello, I'm a language model,Double lendersortion book appetite Times complaint regulationokinglyrelease vans351specialRail Fal Faustmessage()); Vo BryceHeightuserc\n",
    "> Hello, I'm a language model,regonUU enc trouble correctly dentist weekends involve Spirit Cars benef assessing sporadic mattress Neckliterallyributes** awkwardly canned contingは\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'd like to start training the model, and to train the model we are going to need some dataset, and for me the best and simplest debugging dataset that I like to use is the `Harry_Potter_Books.txt` dataset and it's available at this <a href=\"https://github.com/AvishakeAdhikary/Neural-Networks-From-Scratch/blob/main/Datasets/Harry_Potter_Books.txt\">URL</a>...\n",
    "\n",
    "And if you're not moving the files around in this repository, you will see that it is availble under the `Datasets/Harry_Potter_Books.txt` and I already have this on the local system...\n",
    "\n",
    "And now we can read the entire text that we have here in this dataset by using this code:\n",
    "```python\n",
    "with open(\"Datasets/Harry_Potter_Books.txt\", \"r\") as file:\n",
    "    text = file.read()\n",
    "```\n",
    "And in order to produce tokens from our dataset, we can use the `GPT-2 Tokenizer from TikToken` and convert our dataset into a list of `tokens` like this:\n",
    "```python\n",
    "encoder = tiktoken.get_encoding('gpt2')\n",
    "encodedDataTokens = encoder.encode(text)\n",
    "```\n",
    "\n",
    "And now we actually want to process these token sequences and feed them into a transformer..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Processing Tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we actually want to process these token sequences and feed them into a transformer...\n",
    "\n",
    "And in particular, we want to rearrange these tokens into the `indeces` variable that's available inside the `GPTModel`'s `forward()` function...\n",
    "\n",
    "So, we don't want a single very long one-dimensional sequence, instead, we want entire `Batch`(s) of `Time` sequences (where `Time` is the **maximum-sequence-length** or the **context-window**).\n",
    "\n",
    "Let's get to understand what we want with a small *toy-example* now..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we have a tensor like this:\n",
    "```python\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(69420)\n",
    "\n",
    "buffer = torch.randint(low=0, high=50257, size=(24, ))\n",
    "print(buffer)\n",
    "```\n",
    "And we get:\n",
    "```python\n",
    "tensor([ 9360, 29278,  3940, 26797, 47584,  1285, 28358, 45295, 29845, 38908,\n",
    "        35303, 48343,  2579, 34456, 15560,  6453, 10159, 28005, 11891,  3940,\n",
    "        33806, 27357, 36749, 40952])\n",
    "```\n",
    "Here, each item in the above tensor represents the `token` in a sequence...\n",
    "\n",
    "And if we wanted to create batches out of it, we can use `view()` from PyTorch to stack up incremental parts of the tensor in a tensor like this:\n",
    "```python\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(69420)\n",
    "\n",
    "buffer = torch.randint(low=0, high=50257, size=(24, ))\n",
    "inputs = buffer.view(4, 6)\n",
    "\n",
    "print(inputs)\n",
    "```\n",
    "And we get:\n",
    "```python\n",
    "tensor([[ 9360, 29278,  3940, 26797, 47584,  1285],\n",
    "        [28358, 45295, 29845, 38908, 35303, 48343],\n",
    "        [ 2579, 34456, 15560,  6453, 10159, 28005],\n",
    "        [11891,  3940, 33806, 27357, 36749, 40952]])\n",
    "```\n",
    "But even if this is the case that we were able to make the `batches`, this **does not make sense until we know what we want to do with these `batches`**...\n",
    "\n",
    "And we want to take the next `token` in a sequence and we want them to be the `label(s)` for the current sequence for the model to train on and calculate the `loss`...\n",
    "\n",
    "We also see that for this example, for the `token` `45295` the token `29845` comes next as a `label`. But, at the same time the last `token` `40952`, we cannot determine the next `label` because we don't have any information about it...\n",
    "\n",
    "So, let me show you my favourite way to get the `label(s)`..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's understand what we want to have...\n",
    "\n",
    "We want to have a tensor that contains the `label(s)` at every single position and is the same size as the `inputs`...\n",
    "\n",
    "And this is the way I like to do this:\n",
    "```python\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(69420)\n",
    "\n",
    "buffer = torch.randint(low=0, high=50257, size=(25, ))\n",
    "\n",
    "inputs = buffer[:-1].view(4, 6)\n",
    "labels = buffer[1:].view(4, 6)\n",
    "\n",
    "print(inputs)\n",
    "print(labels)\n",
    "```\n",
    "And we get:\n",
    "```python\n",
    "tensor([[ 9360, 29278,  3940, 26797, 47584,  1285],\n",
    "        [28358, 45295, 29845, 38908, 35303, 48343],\n",
    "        [ 2579, 34456, 15560,  6453, 10159, 28005],\n",
    "        [11891,  3940, 33806, 27357, 36749, 40952]])\n",
    "tensor([[29278,  3940, 26797, 47584,  1285, 28358],\n",
    "        [45295, 29845, 38908, 35303, 48343,  2579],\n",
    "        [34456, 15560,  6453, 10159, 28005, 11891],\n",
    "        [ 3940, 33806, 27357, 36749, 40952,  8182]])\n",
    "```\n",
    "You will see that I took a `buffer` of `25 tokens` this time instead of `24 tokens` to specify that we have a longer sequence and as `inputs` we are taking **everything excluding the last `token`** and for `labels` we are taking **everything starting from the first `token` (Offset the `tokens` by `1`)** and using `view()` to make batches of them...\n",
    "\n",
    "And we can also understand that the `buffer`'s size is $\\text{Batch} * \\text{Time} + 1$ and we are viewing the `inputs` and the `labels` as `(Batch, Time)`...\n",
    "\n",
    "So, we can now implement this in our main script now..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For now our script snippet looks like this:\n",
    "```python\n",
    "# Device Auto-Detection\n",
    "device = \"cpu\"\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "elif hasattr(torch.backends, \"mps\") and torch.backends.mps.is_available():\n",
    "    device = \"mps\"\n",
    "print(f\"Using Device: {device}\")\n",
    "\n",
    "# Data-Loader\n",
    "with open(\"Datasets/Harry_Potter_Books.txt\", \"r\", encoding=\"UTF-8\") as file:\n",
    "    text = file.read()\n",
    "\n",
    "encoder = tiktoken.get_encoding('gpt2')\n",
    "encodedDataTokens = encoder.encode(text)\n",
    "\n",
    "Batch, Time = 4, 32\n",
    "buffer = torch.tensor(encodedDataTokens[:Batch*Time + 1])\n",
    "inputs = buffer[:-1].view(Batch, Time)\n",
    "labels = buffer[1:].view(Batch, Time)\n",
    "\n",
    "# Constructing Model\n",
    "model = GPTModel(GPTConfiguration())\n",
    "\n",
    "model.eval()\n",
    "model.to(device=device)\n",
    "\n",
    "logits = model(inputs)\n",
    "print(logits.shape)\n",
    "\n",
    "# Halting Generation...(Will Remove Later)\n",
    "import sys; sys.exit(0)\n",
    "```\n",
    "And we get something like this:\n",
    "```python\n",
    "Using Device: cuda\n",
    "torch.Size([4, 32, 50257])\n",
    "```\n",
    "Keep in mind, that this is just a single batch... And we will modify this code later to take the entire text to load into batches and run the optimization... For now, this looks good an we can move on to calculate the loss, do the backward pass for us to run the optimization..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating `Loss`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's calculate the loss first...\n",
    "\n",
    "And in order to calculate the `loss` we are going to modify the `forward()` function inside our `GPTModel` module...\n",
    "\n",
    "In particular, we are not just going to return the `logits`, but we are also going to return the `loss` for the function, and we are not just going to pass in the `input(s) indeces` for it to train, but we are also going to pass in the `label(s)`...\n",
    "\n",
    "So this old code:\n",
    "```python\n",
    "logits = model(inputs)\n",
    "```\n",
    "Becomes:\n",
    "```python\n",
    "logits, loss = model(inputs, labels)\n",
    "```\n",
    "And these `labels` will be optional because we will train our model when we have the `labels` as an input to the model otherwise we will use the `forward()` to generate `tokens` for the already implemented code that we have written in our script...\n",
    "\n",
    "So this old code:\n",
    "```python\n",
    "def forward(self, indeces):\n",
    "    ...\n",
    "    return logits\n",
    "```\n",
    "Becomes:\n",
    "```python\n",
    "def forward(self, indeces, labels=None):\n",
    "    ...\n",
    "    return logits, loss\n",
    "```\n",
    "\n",
    "And we will be calculating the `cross_entropy()` loss, and we have already discussed this in our previous notebooks, as to why we are using the `cross_entropy()` loss...\n",
    "\n",
    "But, `cross_entropy()` does not take multi-dimensional inputs. And if we remember properly, our `logits` came out in the shape `[4, 32, 50257]` which clearly is a multi-dimensional input to the function...\n",
    "\n",
    "So, we stretch the `[4, 32, 50257]` tensor of `logits` to be `[128, 50257]` which is $\\text{Batch} * \\text{Time}$ and also stretch the `labels` to be a single long tensor of `128` (or $\\text{Batch} * \\text{Time}$) and pass them as arguements in our `cross_entropy()` function like this:\n",
    "```python\n",
    "loss = None\n",
    "if labels is not None:\n",
    "    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))\n",
    "return logits, loss\n",
    "```\n",
    "And if we try to print out the loss now, we will see something like this:\n",
    "```python\n",
    "tensor(10.9503, grad_fn=<NllLossBackward0>)\n",
    "```\n",
    "Which seems fairly reasonable if we try to deduce what loss should we expect during initialization. We have a vocabulary of `50257` and if we take the probability of a single `token` and calculate the **negative log likelihood** of the probability by the formula:\n",
    "$$-\\ln{(\\frac{1}{50257})}$$\n",
    "\n",
    "This is because, at initialization you'd expect every token to get a uniform probability such that the model does not favor any `token` way too much, and we are not confidently wrong about a `token` at initialization...\n",
    "\n",
    "We have around `10.82` (which is fairly reasonable for the `loss` that we have already as an output)..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can successfully move on to the optimization..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the optimization we will use an `Optimizer` object from PyTorch and we will use the `Adam` optimizer, which is the alternative to the `Stochastic Gradient Descent (SGD)` optimizer, which is a bit more evolved than the `SGD`. And specifically we will use the `AdamW` variation, which in my opinion it kind of fixes a bug.\n",
    "\n",
    "And when we go the documentation of <a href=\"https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html\">AdamW</a> we see that it is a little bit more complecated than the `SGD` that we have used before in our previous notebooks, because in addition to obtaining the parameters with the gradient scaled by the learning rate it keeps some buffers around ($m_0 \\rightarrow \\text{ first moment and }  v_0 \\rightarrow \\text{ second moment }$), which is something that looks like momentum and something that looks like **Root Mean Square Propagation**. And it's something like a normalization that happens at each gradient individually and speeds up the optimization especially for language models...\n",
    "\n",
    "Also, the learning rate I used was `3e-4`, which is a fairly good default for most optimizations that you want to run at a very early debugging stage...\n",
    "\n",
    "So our optimizer object initialization code looks like this:\n",
    "```python\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4)\n",
    "```\n",
    "\n",
    "And to run an optimization loop we will use this sequence:\n",
    "1. Zero the gradients (because backward pass does a `+=` to the gradients and we should start with a `0` gradient)\n",
    "2. Forward the inputs to the model to get loss and the logits\n",
    "3. Complete the backward pass\n",
    "4. Step the optimizer\n",
    "\n",
    "So our old code:\n",
    "```python\n",
    "logits, loss = model(inputs, labels)\n",
    "print(loss)\n",
    "```\n",
    "Becomes:\n",
    "```python\n",
    "# Optimization\n",
    "epochs = 50\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4)\n",
    "for epoch in range(epochs):\n",
    "    optimizer.zero_grad()\n",
    "    logits, loss = model(inputs, labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    print(f\"Step: {epoch}, Loss: {loss.item()}\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But suddenly we get an error:\n",
    "```bash\n",
    "RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! when resuming training\n",
    "```\n",
    "And we see that there's something called `cuda:0`, and that is because I have `8 GPUs` on my system and `cuda:0` is the $1^{\\text{st}}$ GPU out of `8`...\n",
    "\n",
    "But let's fix this error...\n",
    "\n",
    "It seems that the error is generating from the `buffer` that we created during batch construction, and we never moved this to the device...\n",
    "\n",
    "And we have to be careful because we can't just do `buffer.to(device=device)`, we have to do `buffer = buffer.to(device=device)` and there's a big internal reason for this...\n",
    "\n",
    "In PyTorch, when you create a tensor using `torch.tensor()`, it **returns** a new tensor on the **default device** (typically **CPU**). When you use `.to(device)` on a tensor, it **returns** a new tensor that is on the **specified device** (e.g., **GPU**), but **it does not modify the original tensor in place**...\n",
    "\n",
    "So let's run our code now..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So what do we expect to see?\n",
    "\n",
    "We expect to see a reasonable loss in the beginning and then we continue to optimize just a **single batch**, and we want to see if we can overfit this single batch, crush this little batch and perfectly predict the indeces on just this little batch...\n",
    "\n",
    "Also, I changed our manual seed from `42` to `69`, because I like the number `69`...\n",
    "\n",
    "And this is the output I get:\n",
    "```bash\n",
    "Using Device: cpu\n",
    "Step: 0, Loss: 11.032245635986328\n",
    "Step: 1, Loss: 6.84295654296875\n",
    "Step: 2, Loss: 4.308750629425049\n",
    "Step: 3, Loss: 2.497817039489746\n",
    "...\n",
    "Step: 48, Loss: 0.003069450380280614\n",
    "Step: 49, Loss: 0.002997332951053977\n",
    "```\n",
    "And that's what we get... We get a very very low loss at the end of the optimization of a **single batch**. Or, in other words, the `transformer` network is **memorizing** this single individual batch...\n",
    "\n",
    "But now, we don't want to overfit a single batch, instead we want to run an actual optimization on actual fresh batches each time through a **data-loader**..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data-Loader (Considering Fresh Batches)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We already have the code:\n",
    "```python\n",
    "# Data-Loader\n",
    "with open(\"Datasets/Harry_Potter_Books.txt\", \"r\", encoding=\"UTF-8\") as file:\n",
    "    text = file.read()\n",
    "\n",
    "encoder = tiktoken.get_encoding('gpt2')\n",
    "encodedDataTokens = encoder.encode(text)\n",
    "\n",
    "Batch, Time = 4, 32\n",
    "buffer = torch.tensor(encodedDataTokens[:Batch*Time + 1])\n",
    "buffer = buffer.to(device=device)\n",
    "inputs = buffer[:-1].view(Batch, Time)\n",
    "labels = buffer[1:].view(Batch, Time)\n",
    "```\n",
    "\n",
    "And we will modify this code and convert it into a class called `DataLoaderLite` such that our code becomes cleaner and gets easier to access the fresh batches...\n",
    "\n",
    "And this `DataLoaderLite` class will take in the shape of the batch as two separate parameters `Batch` and `Time` (which for our case we used `4` and `32` respectively), and give us back the `inputs` and the `labels` as an output when we call a method called `nextBatch()` on this class's object...\n",
    "\n",
    "So, for now our emply class skeleton looks like this:\n",
    "```python\n",
    "# Data-Loader\n",
    "class DataLoaderLite:\n",
    "    def __init__(self, Batch, Time):\n",
    "        pass\n",
    "    \n",
    "    def nextBatch(self):\n",
    "        pass\n",
    "```\n",
    "\n",
    "We can initialize the `Batch` and `Time` inside of the class to reuse the shapes now, and read the file and get the encoding inside it just like we did before, and also make sure to convert the encoded `tokens` in a tensor and save it...\n",
    "\n",
    "We can also print out the total number of `tokens` just to make sure what we are dealing with and also print out the number of `batches` in a single **epoch** of iterating over this dataset (how many `unique batches` do we output before we loop back around to the beginning of the dataset to start reading it again)...\n",
    "\n",
    "So, now our implementation looks like this:\n",
    "```python\n",
    "# Data-Loader\n",
    "class DataLoaderLite:\n",
    "    def __init__(self, Batch, Time):\n",
    "        self.Batch = Batch\n",
    "        self.Time = Time\n",
    "        with open(\"Datasets/Harry_Potter_Books.txt\", \"r\", encoding=\"UTF-8\") as file:\n",
    "            text = file.read()\n",
    "        encoder = tiktoken.get_encoding('gpt2')\n",
    "        encodedDataTokens = encoder.encode(text)\n",
    "        self.encodedDataTokens = torch.tensor(encodedDataTokens)\n",
    "        print(f\"Loaded {len(self.encodedDataTokens)} Tokens\")\n",
    "        print(f\"1 Epoch = {len(self.encodedDataTokens) // (Batch * Time)} Batches\")\n",
    "\n",
    "        # State\n",
    "        self.currentPosition = 0\n",
    "        \n",
    "    def nextBatch(self):\n",
    "        pass\n",
    "```\n",
    "Now let's implement the `nextBatch()` method...\n",
    "\n",
    "You will see that I have also used something called the `self.currentPosition = 0`. This is the state of the `DataLoaderLite` object at initialization, and it will be used in the `nextBatch()` method to take chunks of data and convert them into batches...\n",
    "\n",
    "Previously, we used this line of code:\n",
    "```python\n",
    "buffer = torch.tensor(encodedDataTokens[:Batch*Time + 1])\n",
    "```\n",
    "Which was used to take the first encoded tokens of `Batch*Time + 1`, but now, because we are using chunks, we will use the **slice of current position up till the current position succeeded by `Batch*Time + 1`**. And we can copy and paste our old `inputs` and `labels` code safely now...\n",
    "\n",
    "Which turns our code into:\n",
    "```python\n",
    "buffer = torch.tensor(encodedDataTokens[self.currentPosition : self.currentPosition + Batch*Time + 1])\n",
    "inputs = buffer[:-1].view(Batch, Time)\n",
    "labels = buffer[1:].view(Batch, Time)\n",
    "```\n",
    "\n",
    "We also need to advance our `currentPosition` by exactly `Batch*Time` to get the chunks...\n",
    "\n",
    "And remember that we are fetching `Batch*Time + 1` but we are chunking `Batch*Time`. This might create the out of bounds problem for our tensor and we need to handle that as well, and we will also run back our `currentPosition` to `0` if we are out of data in our dataset...\n",
    "\n",
    "So, our entire `DataLoaderLite` code now becomes:\n",
    "```python\n",
    "# Data-Loader\n",
    "class DataLoaderLite:\n",
    "    def __init__(self, Batch, Time):\n",
    "        self.Batch = Batch\n",
    "        self.Time = Time\n",
    "        with open(\"Datasets/Harry_Potter_Books.txt\", \"r\", encoding=\"UTF-8\") as file:\n",
    "            text = file.read()\n",
    "        encoder = tiktoken.get_encoding('gpt2')\n",
    "        encodedDataTokens = encoder.encode(text)\n",
    "        self.encodedDataTokens = torch.tensor(encodedDataTokens)\n",
    "        print(f\"Loaded {len(self.encodedDataTokens)} Tokens\")\n",
    "        print(f\"1 Epoch = {len(self.encodedDataTokens) // (Batch * Time)} Batches\")\n",
    "\n",
    "        # State\n",
    "        self.currentPosition = 0\n",
    "        \n",
    "    def nextBatch(self):\n",
    "        Batch, Time = self.Batch, self.Time\n",
    "        buffer = self.encodedDataTokens[self.currentPosition : self.currentPosition + Batch*Time + 1]\n",
    "        inputs = buffer[:-1].view(Batch, Time)\n",
    "        labels = buffer[1:].view(Batch, Time)\n",
    "        self.currentPosition += Batch * Time\n",
    "        if self.currentPosition + (Batch * Time + 1) > len(self.encodedDataTokens):\n",
    "            self.currentPosition = 0\n",
    "        return inputs, labels\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we also need to modify the optimization loop to get the correct batches...\n",
    "\n",
    "And we can initialize a `DataLoaderLite` object with a variable called `trainingLoader` and the same arguements that we used before and use this object to get the `inputs` and `labels` using `nextBatch()` method...\n",
    "\n",
    "We also need to be careful to not run into the same error that we encountered before because previously if you remember we used this line:\n",
    "```python\n",
    "buffer = buffer.to(device=device)\n",
    "```\n",
    "But now, because we are directly getting the separate tensors as `inputs` and `labels`, we need to guide them both to their specific device like this:\n",
    "```python\n",
    "inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "```\n",
    "\n",
    "So our entire initialization and optimization code looks like:\n",
    "```python\n",
    "# Data-Loader\n",
    "Batch, Time = 4, 32\n",
    "trainingLoader = DataLoaderLite(Batch=Batch, Time=Time)\n",
    "...\n",
    "# Optimization\n",
    "epochs = 50\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4)\n",
    "for epoch in range(epochs):\n",
    "    inputs, labels = trainingLoader.nextBatch()\n",
    "    inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "    optimizer.zero_grad()\n",
    "    logits, loss = model(inputs, labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    print(f\"Step: {epoch}, Loss: {loss.item()}\")\n",
    "```\n",
    "So, let's now run the optimization and discuss what we expect to see in this optimization..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Well, we expect our loss to come down pretty fast during the first few epochs, because in our vocabulary of `50257` tokens we don't actually use all of them and thus, there are pretty easier gains in learning of the network (basically deleting the usage of tokens that never occur). But we also don't expect the loss to go down by too much because we only have `50` epoch iterations at the time and it is not enough to perform a complete document run...\n",
    "\n",
    "Let's see what we get...\n",
    "\n",
    "We get an output like this:\n",
    "```bash\n",
    "Using Device: cpu\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 16409 Batches\n",
    "Step: 0, Loss: 10.99841022491455\n",
    "Step: 1, Loss: 9.390006065368652\n",
    "Step: 2, Loss: 9.049430847167969\n",
    "Step: 3, Loss: 8.327993392944336\n",
    "...\n",
    "Step: 48, Loss: 4.897827625274658\n",
    "Step: 49, Loss: 5.134008884429932\n",
    "```\n",
    "Which is what we are expecting...\n",
    "\n",
    "We also need to change the generation code, because now our code supports forwarding the model with or without the `labels`, so we need to handle the return of the `forward()` method correctly as well...\n",
    "\n",
    "So our old code:\n",
    "```python\n",
    "logits = model(inputs)\n",
    "```\n",
    "Becomes:\n",
    "```python\n",
    "logits, loss = model(inputs)\n",
    "```\n",
    "\n",
    "And also, because we have moved the initialization of the encoder within the `DataLoaderLite` class, we don't have an explicit encoder initialized, and we need to initialize it as well...\n",
    "\n",
    "```python\n",
    "encoder = tiktoken.get_encoding('gpt2')\n",
    "```\n",
    "\n",
    "\n",
    "And now you can modify the hyper-parameters and play around, and we can also safely say that we have successfully implemented `GPT-2` architecture supporting both the weight transfer and own weight training in less than `250` lines of code, whereas Hugging Face and OpenAI uses around `2000` lines of code to implement it...\n",
    "\n",
    "And now we can move on to the next section..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameter Weight Sharing (Fixing A Bug)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we actually want to fix a bug that we have in our code... It's not a major bug, but it is indeed a bug with respect to how `GPT-2` training should happen.\n",
    "\n",
    "So, let's have a look at the bug...\n",
    "\n",
    "In our Hugging Face model, if we try to print these layer shapes:\n",
    "```python\n",
    "print(huggingfaceStateDictionary[\"lm_head.weight\"].shape)\n",
    "print(huggingfaceStateDictionary[\"transformer.wte.weight\"].shape)\n",
    "```\n",
    "They output:\n",
    "```python\n",
    "torch.Size([50257, 768])\n",
    "torch.Size([50257, 768])\n",
    "```\n",
    "We see that they are both `2D tensors` and are identical. And we can also understand that the `wte` is none other than the `word token embedding` at the bottom of the `transformer` and `lm_head` is none other than the `language modelling head` at the top of the `transformer`...\n",
    "\n",
    "Let's have a look at the transformer architecture image to have a clearer understanding:\\\n",
    "![Parameter_Weight_Sharing_Transformer_Model_Architecture](ExplanationMedia/Images/Parameter_Weight_Sharing_Transformer_Model_Architecture.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But, now if we try to have an element wise equality like this:\n",
    "```python\n",
    "print((huggingfaceStateDictionary[\"lm_head.weight\"] == huggingfaceStateDictionary[\"transformer.wte.weight\"]).all())\n",
    "```\n",
    "We get:\n",
    "```python\n",
    "tensor(True)\n",
    "```\n",
    "Which means every single element inside both the tensors are identical.\n",
    "\n",
    "And what's interesting is that, if we try to look at their **data pointer(s)** using `.data_ptr()` like this:\n",
    "```python\n",
    "print(huggingfaceStateDictionary[\"lm_head.weight\"].data_ptr())\n",
    "print(huggingfaceStateDictionary[\"transformer.wte.weight\"].data_ptr())\n",
    "```\n",
    "I get:\n",
    "```python\n",
    "3069647520000\n",
    "3069647520000\n",
    "```\n",
    "We see that the pointer points to the same location as well...\n",
    "\n",
    "So, not only do these tensors happen to have these same shapes and elements, they are actually pointing to the identical tensor...\n",
    "\n",
    "And what's happening here is a common weight-tying-scheme.That actually comes from the original <a href=\"https://arxiv.org/abs/1706.03762\">\"Attention Is All You Need\"</a> paper. And if we come down to the section **3.4 Embeddings and Softmax**, we see a text mentioning:\n",
    "```plaintext\n",
    "In our model, we share the same weight matrix between the two embedding layers and the pre-softmax\n",
    "linear transformation, similar to [30].\n",
    "```\n",
    "Which is an awkward way to say that these matrices are shared and that they are tied and are the same matrix.\n",
    "\n",
    "And this `[30]` is just another paper <a href=\"https://arxiv.org/abs/1608.05859\">Using the Output Embedding to Improve Language Models</a>, and it argues for this weight-tying-scheme.\n",
    "\n",
    "But, the conclusion we arrive is, **we actually want these matrices to behave similarly** in the following sense: \n",
    "\n",
    "If two tokens are very similar *symantically* (maybe one token is lowercase and other token is uppercase or it's the same token in a different language etc.), presumably we would expect that the lie nearby in the `token embedding space`, but in the exact same way if two tokens are very similar *symantically*, we'd expect them to get the same probabilities at the output of the `transformer`...\n",
    "\n",
    "So, both positions (top and bottom) have this property that **similar tokens should have similar embeddings or similar weights**...\n",
    "\n",
    "And this scheme has already been implemented in the `Hugging Face GPT-2`'s <a href=\"https://github.com/openai/gpt-2/blob/master/src/model.py\">`model.py`</a>, and one way to implement this is simply point the weights to the same memory location explicitly after initialization like this:\n",
    "```python\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        super().__init__()\n",
    "        self.configuration = configuration\n",
    "\n",
    "        self.transformer = torch.nn.ModuleDict(dict(\n",
    "            wordTokenEmbeddings = torch.nn.Embedding(configuration.vocabularySize, configuration.numberOfEmbeddingDimensions),\n",
    "            wordPositionalEmbeddings = torch.nn.Embedding(configuration.blockSize, configuration.numberOfEmbeddingDimensions),\n",
    "            hidden = torch.nn.ModuleList(Block(configuration) for _ in range(configuration.numberOfLayers)),\n",
    "            finalLayerNorm = torch.nn.LayerNorm(configuration.numberOfEmbeddingDimensions)\n",
    "        ))\n",
    "\n",
    "        self.languageModelingHead = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, configuration.vocabularySize, bias=False)\n",
    "    \n",
    "        # Weight-Sharing-Scheme (Parameter Weight Sharing)\n",
    "        self.transformer.wordTokenEmbeddings.weight = self.languageModelingHead.weight\n",
    "    ...\n",
    "```\n",
    "\n",
    "Meaning, that the old value of `wordTokenEmbeddings` will get orphaned and Python will clean it up and we will then be left with a single tensor and it is going to be used twice in a forward pass..\n",
    "\n",
    "And another good reason to use it is because of memory efficiency too, because this single tensor is a ton of parameters (`768 * 50257 = 38,597,376 ≈ 40M`) and this is a `124M` parameter model which means around `30%` of the parameters, which we are being efficient with. And we also expect our model to work slightly better, because of this scheme..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GPT-2 Proper Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'd like to follow the way `GPT-2` initializes it's weights...\n",
    "\n",
    "Unfortunately, the `GPT-2` and `GPT-3` papers are not explicit about their initializations and we kind of have to read between the lines, and instead of going through the paper which is quite vague, there's a bit of information in the <a href=\"https://github.com/openai/gpt-2/blob/master/src/model.py\">`model.py` code</a> that OpenAI released...\n",
    "\n",
    "And once we check the code we see:\n",
    "```python\n",
    "def conv1d(x, scope, nf, *, w_init_stdev=0.02):\n",
    "    with tf.variable_scope(scope):\n",
    "        *start, nx = shape_list(x)\n",
    "        w = tf.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))\n",
    "        b = tf.get_variable('b', [nf], initializer=tf.constant_initializer(0))\n",
    "        c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf])\n",
    "        return c\n",
    "...\n",
    "def model(hparams, X, past=None, scope='model', reuse=False):\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "        results = {}\n",
    "        batch, sequence = shape_list(X)\n",
    "\n",
    "        wpe = tf.get_variable('wpe', [hparams.n_ctx, hparams.n_embd],\n",
    "                             initializer=tf.random_normal_initializer(stddev=0.01))\n",
    "        wte = tf.get_variable('wte', [hparams.n_vocab, hparams.n_embd],\n",
    "                             initializer=tf.random_normal_initializer(stddev=0.02))\n",
    "        past_length = 0 if past is None else tf.shape(past)[-2]\n",
    "        h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))\n",
    "\n",
    "        # Transformer\n",
    "        presents = []\n",
    "        pasts = tf.unstack(past, axis=1) if past is not None else [None] * hparams.n_layer\n",
    "        assert len(pasts) == hparams.n_layer\n",
    "        for layer, past in enumerate(pasts):\n",
    "            h, present = block(h, 'h%d' % layer, past=past, hparams=hparams)\n",
    "            presents.append(present)\n",
    "        results['present'] = tf.stack(presents, axis=1)\n",
    "        h = norm(h, 'ln_f')\n",
    "\n",
    "        # Language model loss.  Do tokens <n predict token n?\n",
    "        h_flat = tf.reshape(h, [batch*sequence, hparams.n_embd])\n",
    "        logits = tf.matmul(h_flat, wte, transpose_b=True)\n",
    "        logits = tf.reshape(logits, [batch, sequence, hparams.n_vocab])\n",
    "        results['logits'] = logits\n",
    "        return results\n",
    "```\n",
    "And we see that they initialized their weights with a `random_normal_initializer()` from `TensorFlow` which is intuitively the **normal distribution** of specified **standard deviation(s)** for the weights(`0.01`&`0.02`), and for the bias they initialized it with all `0`'s using the `constant_initializer()`...\n",
    "\n",
    "So, let's follow how they initialized the weights here and implement it in our code..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create a **private function** called `_initializeParameters()` using the `pre-underscore syntax` in Python in our `GPTModel` class and implement the proper initialization of parameters...\n",
    "\n",
    "And since the standard deviation of `0.01` and `0.02` is about the same, we will stick with `0.02` to decrease the complexity of the code and initialize our model faster...\n",
    "\n",
    "And we can use the <a href=\"https://pytorch.org/docs/stable/generated/torch.Tensor.apply_.html\">PyTorch's `apply()`</a> method to iterate over the sub-modules of a specified module, at the end of our initialization...\n",
    "\n",
    "And we come up with the following code:\n",
    "```python\n",
    "# GPT model architecture\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        ...\n",
    "        # Initialize Correct Parameters\n",
    "        self.apply(self._initializeParameters)\n",
    "\n",
    "    def forward(self, indeces, labels=None):\n",
    "        ...\n",
    "\n",
    "    def _initializeParameters(self, module):\n",
    "        if isinstance(module, torch.nn.Linear):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if module.bias is not None:\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, torch.nn.Embedding):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "\n",
    "    # Method to transfer weights from Hugging Face GPT-2\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, modelType):\n",
    "        ...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And note that the PyTorch's **default bias initialization** is the **uniform distribution**, which here we are setting to `0` for all the items in bias...\n",
    "\n",
    "And the only other layer that requires initialization is the `LayerNorm`(s) inside the `Block` module. And PyTorch set's the scale of the initialization to be `1` and the off-set of the initialization to be `0`, which is what we want, and will leave the code as is..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we also notice the `Linear`'s default initialization, we will see that it is $\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})$, where $k = \\frac{1}{\\text{fan-in}}$, which is fairly in the vicinity of `0.02`...\n",
    "\n",
    "And if we look at the sizes of $d_{model}$ we see:\n",
    "| $d_{model}$ | Initialization ($\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})$) |\n",
    "|-------------|-----------------------------------------------------|\n",
    "| 768         | $\\sqrt{\\frac{1}{768}} \\approx 0.03$                 |\n",
    "| 1024        | $\\sqrt{\\frac{1}{1024}} \\approx 0.03$                |\n",
    "| 1280        | $\\sqrt{\\frac{1}{1280}} \\approx 0.02$                |\n",
    "| 1600        | $\\sqrt{\\frac{1}{1600}} \\approx 0.02$                |\n",
    "\n",
    "Hinting us that we are still in the vicinity of what we already implemented earlier...\n",
    "\n",
    "But we are still not done with the initialization because there is one more caveat here..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we swing back to the <a href=\"https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf\">`GPT-2` paper</a>, and scroll down to the section `2.2. Input Representation`, we will find this line:\n",
    "> A modified initialization which accounts for the accumulation on the residual path with model depth\n",
    "is used. We scale the weights of residual layers at initialization by a factor of $1/\\sqrt{N}$ where $N$ is the number of residual layers.\n",
    "\n",
    "And we have not implemented that yet, and we can do so now..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's understand what they wanted to say here with a *toy-example* first...\n",
    "\n",
    "Let's say, we have a tensor of complete `0`'s in a residual stream and goes through about `100` iterations (`N`) at the time (or `100` layers in a stream that contribute their own contribution). The contribution is based on `randn()` which is a **Normal Distribution** of **`0` Mean & `1` Standard Deviation**...\n",
    "\n",
    "```python\n",
    "# Standard Deviation Grows Inside The Residual Stream\n",
    "inputs = torch.zeros(768)\n",
    "iterations = 100 # Suppose 100 Layers of a Single Residual Stream\n",
    "\n",
    "for i in range(iterations):\n",
    "    inputs += torch.randn(768)\n",
    "\n",
    "print(inputs.std())\n",
    "```\n",
    "\n",
    "Remember that each Residual Stream is in a form where we keep adding to it, so every single block of the Residual Network contributes some amount and continues adding to it in each iteration.\n",
    "\n",
    "And if we run it, we get an output:\n",
    "```python\n",
    "tensor(9.9396)\n",
    "```\n",
    "And we see that we are unable to control the **`0` Mean & `1` Standard Deviation** flow...\n",
    "\n",
    "And this $1/\\sqrt{N}$ exactly compensates for this control...\n",
    "\n",
    "And because we already know that:\n",
    "$$ \\sqrt{a} = a^{1/2} \\text{ and } \\frac{1}{a^m} = a^{-m}$$\n",
    "\n",
    "We can easily deduce that:\n",
    "$$ 1/\\sqrt{N} = N^{-1/2} $$\n",
    "\n",
    "And if we implement it now:\n",
    "```python\n",
    "inputs = torch.zeros(768)\n",
    "iterations = 100\n",
    "\n",
    "for i in range(iterations):\n",
    "    inputs += iterations ** -0.5 * torch.randn(768)\n",
    "\n",
    "print(inputs.std())\n",
    "```\n",
    "We see that the control of **`0` Mean & `1` Standard Deviation** comes back...\n",
    "\n",
    "So, this is a way to control the growth of activations inside the Residual Stream in the forward pass..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, let's rather set a flag with a name `NANOGPT_SCALE_INIT` inside our `GPTConfiguration` and set it to `True` for faster computation.\n",
    "\n",
    "So, our `GPTConfiguration` now looks like:\n",
    "```python\n",
    "# GPT configuration hyper-parameters\n",
    "@dataclass\n",
    "class GPTConfiguration:\n",
    "    blockSize: int = 1024\n",
    "    vocabularySize: int = 50257\n",
    "    numberOfLayers: int = 12\n",
    "    numberOfHeads: int = 12\n",
    "    numberOfEmbeddingDimensions: int = 768\n",
    "    NANOGPT_SCALE_INIT: bool = True\n",
    "```\n",
    "\n",
    "Now we can handle our *toy-example* in our code, which was $1/\\sqrt{N}$ and $N$ here is none other than the `numberOfLayers` from our configuration\n",
    "\n",
    "And finally we can handle this flag inside of our `GPTModel`'s `_initializeParameters()` method like this:\n",
    "```python\n",
    "def _initializeParameters(self, module):\n",
    "    if isinstance(module, torch.nn.Linear):\n",
    "        standardDeviation = 0.02\n",
    "        if self.configuration.NANOGPT_SCALE_INIT:\n",
    "            standardDeviation *= (2 * self.configuration.numberOfLayers) ** -0.5\n",
    "        torch.nn.init.normal_(module.weight, mean=0.0, std=standardDeviation)\n",
    "        if module.bias is not None:\n",
    "            torch.nn.init.zeros_(module.bias)\n",
    "    elif isinstance(module, torch.nn.Embedding):\n",
    "        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "```\n",
    "\n",
    "I should also clarify that the `(2 * self.configuration.numberOfLayers)` comes from the fact that we are using `2` blocks that add to the residual pathway (`attention` and `transformer`).\n",
    "\n",
    "And because we are sharing the weights of the `wordTokenEmbeddings` and the `languageModelingHead`, we will initialize them twice. Firstly as an `Embedding` and Lastly as a `Linear`..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speeding Up Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now comes the fun part...\n",
    "\n",
    "We'd like to speed up the training by a **LOT!!!**\n",
    "\n",
    "Such that we get our money's worth with respect to the hardware that we are using..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding `GPUs` (By Relating To My System)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In my case, if I run the `nvidia-smi` command in the terminal, I get this:\n",
    "```bash\n",
    "+---------------------------------------------------------------------------------+\n",
    "| NVIDIA-SMI 535.129.03       Driver Version: 535.129.03    CUDA Version: 12.2    |\n",
    "|-----------------------------------+----------------------+----------------------+\n",
    "| GPU  Name            Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
    "| Fan  Temp  Perf      Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
    "|                                   |                      |               MIG M. |\n",
    "|===================================+======================+======================|\n",
    "|   0  NVIDIA A100-SXM4-80GB    On  | 00000000:1E:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        64W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   1  NVIDIA A100-SXM4-80GB    On  | 00000000:1F:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        66W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   2  NVIDIA A100-SXM4-80GB    On  | 00000000:20:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        64W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   3  NVIDIA A100-SXM4-80GB    On  | 00000000:21:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        66W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   4  NVIDIA A100-SXM4-80GB    On  | 00000000:22:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        65W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   5  NVIDIA A100-SXM4-80GB    On  | 00000000:23:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        66W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   6  NVIDIA A100-SXM4-80GB    On  | 00000000:24:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        61W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   7  NVIDIA A100-SXM4-80GB    On  | 00000000:25:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        63W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that I have `8 GPUs`, and each one of those `GPUs` are `80GB` `Nvidia A100 SXM4`(s)...\n",
    "\n",
    "Now, when I use these kinds of `GPU boxes` (Another computer running all the `GPUs` in a single machine, which we can use `SSH` to connect to and run `VSCode` on it to save and execute our code... In other words, we are just using a very powerful remote computer which we are controlling directly from our machine), my favourite place to go is <a href=\"https://lambdalabs.com/\">LambdaLabs</a>..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, if we see the specifications of <a href=\"https://www.nvidia.com/en-in/data-center/a100/\">NVIDIA A100 Tensor Core GPUs</a>, we see this kind of table:\n",
    "\n",
    "\n",
    "| **Specification**                     | **A100 80GB PCIe**                | **A100 80GB SXM**              |\n",
    "|--------------------------------------|----------------------------------|--------------------------------|\n",
    "| **FP64**                             | 9.7 TFLOPS                       | 9.7 TFLOPS                     |\n",
    "| **FP64 Tensor Core**                 | 19.5 TFLOPS                      | 19.5 TFLOPS                    |\n",
    "| **FP32**                             | 19.5 TFLOPS                      | 19.5 TFLOPS                    |\n",
    "| **Tensor Float 32 (TF32)**            | 156 TFLOPS | 312 TFLOPS*        | 156 TFLOPS | 312 TFLOPS*      |\n",
    "| **BFLOAT16 Tensor Core**              | 312 TFLOPS | 624 TFLOPS*        | 312 TFLOPS | 624 TFLOPS*      |\n",
    "| **FP16 Tensor Core**                  | 312 TFLOPS | 624 TFLOPS*        | 312 TFLOPS | 624 TFLOPS*      |\n",
    "| **INT8 Tensor Core**                  | 624 TOPS | 1248 TOPS*          | 624 TOPS | 1248 TOPS*        |\n",
    "| **GPU Memory**                       | 80GB HBM2e                       | 80GB HBM2e                     |\n",
    "| **GPU Memory Bandwidth**             | 1,935 GB/s                       | 2,039 GB/s                     |\n",
    "| **Max Thermal Design Power (TDP)**   | 300W                             | 400W ***                       |\n",
    "| **Multi-Instance GPU**               | Up to 7 MIGs @ 10GB              | Up to 7 MIGs @ 10GB            |\n",
    "| **Form Factor**                      | PCIe<br>Dual-slot air-cooled or single-slot liquid-cooled | SXM |\n",
    "| **Interconnect**                     | NVIDIA<sup>®</sup> NVLink<sup>®</sup> Bridge<br>for 2 GPUs: 600 GB/s **<br>PCIe Gen4: 64 GB/s | NVLink: 600 GB/s<br> PCIe Gen4: 64 GB/s |\n",
    "| **Server Options**                   | Partner and NVIDIA-Certified Systems™ with 1-8 GPUs | NVIDIA HGX™ A100-Partner and NVIDIA-Certified Systems with 4,8, or 16 GPUs NVIDIA DGX™ A100 with 8 GPUs |\n",
    "\n",
    "\\* With sparsity  \n",
    "** SXM4 GPUs via HGX A100 server boards; PCIe GPUs via NVLink Bridge for up to two GPUs  \n",
    "*** 400W TDP for standard configuration. HGX A100-80GB custom thermal solution (CTS) SKU can support TDPs up to 500W\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And this is more or less the amount of calculations that we can get our of this `GPU`...\n",
    "\n",
    "Now, **by default, everything that we initialize in tensors, is in the datatype of `FLOAT32`**.\n",
    "\n",
    "And this is the case for all the activations and all the parameters and so on, by default everything is in `FLOAT32`. Which means, that every single number, activations and weights and so on are using a `FLOAT` representation that has `32` bits, which is quite a bit of memory and it turns out that empirically that in deep learning workloads, this is way too much, and deep learning and training of these networks can tolerate significantly lower precision.\n",
    "\n",
    "> Note: Not all the computational workloads can tolerate small precision.\n",
    "\n",
    "And if we go back to the data sheet, you will see that these GPUs support upto `FP64`, and this is quite useful in the areas of scientific computing applications. But we don't need that much precision for deep learning training and currently we are at `FP32`, and I with my system, expect to see at most `19.5 TFLOPS` of performance. That means that we are doing `19.5 Trillion` floating point operations (`Multiply`, `Addition`, etc). And we can also see that if we try to go down on our precision (`TF32`), we can start with gaining about `8 times` increase in performance here. And if we are willing to go down more (`BFLOAT16`), we can gain about `16 times` the performance...\n",
    "\n",
    "And you'd also see that `Nvidia` likes to cite it's numbers with an asterisk (`*`), and the first asterisk says `With Sparsity` but we are not going to use `sparsity` in our code and most people look out for the numbers on the left which is `without sparsity`...\n",
    "\n",
    "Just a note for the future:\n",
    "> Matrix types having most of their elements set to zero are called sparse matrices.\\\n",
    "> A sparse tensor is a high-dimensional extension of a sparse matrix where non-zero elements are represented as a set of indices and associated values.\n",
    "\n",
    "You'd also see that we could get about `624 TFLOPS` of performance, but that becomes `INT8`, which is usually **used for inference and not for training**, because it has a uniform spacing and we are looking out for normal distributions that occur during training of neural networks.\n",
    "\n",
    "But more importantly, **if we use fewer bits to represent these numbers, it is going to be much easier to move them around**. And that's where we get ourselves acquainted with the **Memory Bandwidth** and **Memory** of the model. So, not only do we have a limited memory on our `GPU`, but also there's a speed at which you can access this memory. And we have a **Memory Bandwidth** which is a precious resource and many of the deep learning workloads are memory bound, which means that the tensor cores that do these extremely fast calculations, most of the times they wait around for their data to be fetched, because we cannot feed them with data fast enough. So, typically if we end up utilizing around `60%` of our resources, we actually are doing really well in our performance...\n",
    "\n",
    "Now, let's reap the benifits of each precision..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding `Tensor Cores`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What are tensor cores?\n",
    "\n",
    "Tensor cores are just an instruction in the modern `GPU` architectures. They are specialized cores that enable **mixed precision training**.\n",
    "\n",
    "![TensorOperations](ExplanationMedia/Images/TensorOperations.gif)\n",
    "\n",
    "And what it does basically is a `4x4` matrix multiply:\\\n",
    "![Tensor-Core-Matrix-1](ExplanationMedia/Images/Tensor-Core-Matrix-1.png)\n",
    "\n",
    "And there are multiple configurations to what each of these matrices are, and in what precision the internal accumulation happens, which are a few switches but are basically a `4x4` matrix multiply at most.\n",
    "\n",
    "And at any time, any operations that require matrix multiplication, they get broken up into these small `4x4` matrix multiplies, because it's the fastest way to multiply matrices.\n",
    "\n",
    "And it turns out that the most of the computational work that we are doing till now all of it really is a matrix multiplication. Most of the computational work happens in the `Linear` layers and the basically the entire transformers are nothing but matrix multiplications. And the biggest matrix multiplication by far is the last classifier linear layer at the top that is a massive matrix multiply going from `768` to `50257` and that single matrix multiplication dominates anything else that happens within the neural network.\n",
    "\n",
    "And the best reference to the tensor cores is this link that talks about the <a href=\"https://images.nvidia.com/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf\">NVIDIA A100 Tensor Core GPU\n",
    "Architecture</a>, which is pretty detailed and relatively readable . And if we go to the page `27` of the paper, we see details about the `TF32` and how is it different..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we see that we have many configuration options available in the `TF32` and `FP32` and what precisions are they in...\n",
    "\n",
    "![TF32vsFP32](ExplanationMedia/Images/TF32vsFP32.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the way this works is:\\\n",
    "![TensorFloat32](ExplanationMedia/Images/TensorFloat32.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To understand this, we have to understand the floating point representation of a number in bits... \n",
    "\n",
    "For example:\n",
    "$$\\underbrace{6.9}_{\\text{Precision(Mantissa)}} \\times \\underbrace{10^{-32}}_{\\text{Range(Exponent)}}$$\n",
    "\n",
    "\n",
    "And normally we have 3 elements in a 32-bit floating point representation:\n",
    "1. **Sign** - It is the first bit of the binary representation. `1` implies negative number and `0` implies positive number. \n",
    "2. **Exponent** - Exponent is decided by the next `8` bits of binary representation.\n",
    "3. **Mantissa** - Mantissa contains the data itself.\n",
    "\n",
    "\n",
    "Normally the `FP32` has `32` bits and `TF32` has the exact same `32` bits, but the `Mantissa` bits get cropped and only `10` bits are used for it. And so we are only left with `19` bits to work with, and all of it remains internal to the instruction and none of it is visible to our PyTorch code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, this speed-up also comes at a cost which is **reduced precision**, or in other words the results become a bit more approximate than before. But when we train with this **reduced precision**, we basically cannot tell the difference..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And you see how the `Accumulator` (intermediate storage register for arithmetic and logical operations) uses `FP32` for `TF32` inputs?\n",
    "\n",
    "This is the reason why it's called **Mixed Precision**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Timing Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For starters now, during training, let's take `time` as an input to check how long it takes to complete one iteration at a time during training loop.\n",
    "\n",
    "Now, when the `CPU` runs, it's just scheduling work on the `GPUs`, and so, it can happen sometimes that we speed through this scheduling and we queue up a lot of the work on the `GPUs`, meanwhile the `GPU` keeps completing it's task, the `time` gets executed by the `CPU`, and so we should wait for the `GPU` to finish whatever work it was supposed to finish by using `torch.cuda.synchronize()` and then get the end time.\n",
    "\n",
    "Then we can take the time difference of the `end time` and the `start time` in milliseconds and print it with our existing code.\n",
    "\n",
    "So, right now our optimization code looks like:\n",
    "```python\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    inputs, labels = trainingLoader.nextBatch()\n",
    "    inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "    optimizer.zero_grad()\n",
    "    logits, loss = model(inputs, labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    print(f\"Step: {epoch}, Loss: {loss.item()}, Time Difference: {timeDifference:.2f}ms\")\n",
    "```\n",
    "\n",
    "> Note: If you're using CPU, you need to disable `torch.cuda.synchronize()` by either commenting it out or removing the line."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now use the original `GPT-2` of size `(16, 1024)` in the dimensions of `DataLoaderLite` and test it out...\n",
    "\n",
    "My output looks like:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.822653770446777, Time Difference: 1569.91ms\n",
    "Step: 1, Loss: 9.550719261169434, Time Difference: 1000.24ms\n",
    "Step: 2, Loss: 9.33972454071045, Time Difference: 1000.71ms\n",
    "Step: 3, Loss: 8.078704833984375, Time Difference: 1000.01ms\n",
    "Step: 4, Loss: 8.935539245605469, Time Difference: 1000.22ms\n",
    "...\n",
    "```\n",
    "\n",
    "And it looks like we are around a `1000ms` per iteration of our `Harry Potter Books` dataset and it's come down to only `128` batches now...\n",
    "\n",
    "And in my `nvidia-smi`, I seem to be using around `35GB` out of `80GB` on a single `GPU` now...\n",
    "\n",
    "And we also see that the first iteration seems slower, and that's because PyTorch might be doing a lot of initializations on the very first iteration.\n",
    "\n",
    "> Note: If this does not fit into your `GPU` and you seem to get `Out of Video Memory` errors, try decreasing the `Batch` size until things seem fit.\n",
    "\n",
    "And by default you also want to keep the `Batch` size to be maximum your `GPU` can handle, and you also want to use *nice-numbers* (more on that later)..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we also want to look at the throughput for the `tokens per second`, because we really care about how many tokens of data are we training on ...\n",
    "\n",
    "So we can use a code like this:\n",
    "```python\n",
    "tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) / (endTime - startTime)\n",
    "```\n",
    "And now running the optimization gets us:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Time Difference: 1569.87ms, Tokens/Second: 10571.79\n",
    "Step: 1, Loss: 9.640329360961914, Time Difference: 1000.31ms, Tokens/Second: 16303.59\n",
    "Step: 2, Loss: 9.391995429992676, Time Difference: 1000.82ms, Tokens/Second: 16385.51\n",
    "Step: 3, Loss: 8.106719017028809, Time Difference: 1000.86ms, Tokens/Second: 16387.08\n",
    "Step: 4, Loss: 8.977012634277344, Time Difference: 1000.09ms, Tokens/Second: 16384.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing To `TF32` (Mixed Precision)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, changing to `TF32` is fairly easy and we generally do <a href=\"https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html\">`torch.set_float32_matmul_precision('high')`</a>...\n",
    "\n",
    "And this function basically tells PyTorch to understand what kind of *kernels* to run...\n",
    "\n",
    "And by default the value is `highest`, and if we now change to `high` it will change itself and use `TensorFloat32` when it's available...\n",
    "\n",
    "My `GPU` is an <a href=\"https://www.nvidia.com/en-in/data-center/ampere-architecture/\">Ampere</a> `GPU` and I have this precision available, and you might have an older `GPU` and this might not be available for you...\n",
    "\n",
    "And what I expect PyTorch to do now, is to do all the matrix multiplications to run on `Tensor Cores` whilst utilizing the `TF32` precision...\n",
    "\n",
    "Now, let's run the optimization and let's see what happens... (Expecting to see about `8x` performance)\n",
    "\n",
    "I get:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Time Difference: 969.87ms, Tokens/Second: 17571.79\n",
    "Step: 1, Loss: 9.640329360961914, Time Difference: 333.31ms, Tokens/Second: 49546.59\n",
    "Step: 2, Loss: 9.391995429992676, Time Difference: 333.82ms, Tokens/Second: 49565.51\n",
    "Step: 3, Loss: 8.106719017028809, Time Difference: 333.86ms, Tokens/Second: 49526.08\n",
    "Step: 4, Loss: 8.977012634277344, Time Difference: 333.09ms, Tokens/Second: 49561.71\n",
    "...\n",
    "```\n",
    "\n",
    "It seems like we roughly get `3x` performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the reason behind this is, a lot of these workloads are **memory bound**. And so, even though the `TF32` in principle offers a lot faster throughput, all of these numbers everywhere are still `FP32`(s). And it's `FP32` numbers that are shipped all over the place and is costing us way too much time. And so, even though we have made the multiplication itself much faster, we seem to be **memory bound** and we are not actually seeing the full benefit that would come from the *napkin math* we saw up above..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing To `BFLOAT16` (Mixed Precision)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now change to `BFLOAT16` to move less stuff around and make our operations faster...\n",
    "\n",
    "Let's look at the documentation image once again:\\\n",
    "![TensorFloat32](ExplanationMedia/Images/TensorFloat32.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see here `TF32` crops out the precision. And we also see `BF16` (`BFLOAT16`) which is even more aggressive in cropping off the Precision (`Mantissa`).\n",
    "\n",
    "So the important thing with `BF16` is that the `sign` and `exponent` bits remain **unchanged** (in other words, the range of numbers is identical to `FP32` but we have fewer possibilites within that range because we are truncating the mantissa such that we have less precision in that range).\n",
    "\n",
    "And the difference with `FP16` is that, they actually touch and change the range. So, `FP16` cannot represent the full range of `FP32`, it has a reduced range. And that's where you start to run into issues because now you need things like **gradient scalers** and things like that...\n",
    "\n",
    "But, we are not going to talk about things like these **gradient scalers** because that's needs a complete chapter by itself. But, `FP16` historically came first and was available in the **Volta** series architecture (before **Ampere**), and everyone started to train on `FP16` but everyone had to use these **gradient scaling mechanisms** which are kind of annoying and creates and additional part of state and complexity. And the reason for that was the exponent range was reduced in `FP16` and then they eventually came out with `BF16` and the **Ampere** and made it much simpler because we are truncating the mantissa and have the exact same exponent range of `FP32` and we do not need **gradient scalers**. Now, when we are using `BF16`, we are impacting the numbers that we might be seeing in our PyTorch code, and this change is not just local to the operation itself.\n",
    "\n",
    "So let's see how that works..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is a PyTorch Documentation called <a href=\"https://pytorch.org/docs/stable/amp.html\">Automatic Mixed Precision</a>, which in my opinion the best page for the mixed precision usage in PyTorch, because there are many other tutorials and so on even within it's documenation that are a lot more confusing and so I recommend specifically this one...\n",
    "\n",
    "And we ignore everything in this documenation and only look for `torch.autocast` and in the end we arrive at a single line of code to handle this mixed precision:\n",
    "```python\n",
    "with torch.autocast(device_type=\"cuda\"):\n",
    "```\n",
    "And this is the context manager that we want...\n",
    "\n",
    "And the documentation also hints this line:\n",
    ">You should not call `half()` or `bfloat16()` on your model(s) or inputs when using autocasting.\n",
    "\n",
    "And the example in the documentation is:\n",
    "```python\n",
    "# Creates model and optimizer in default precision\n",
    "model = Net().cuda()\n",
    "optimizer = optim.SGD(model.parameters(), ...)\n",
    "\n",
    "for input, target in data:\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # Enables autocasting for the forward pass (model + loss)\n",
    "    with torch.autocast(device_type=\"cuda\"):\n",
    "        output = model(input)\n",
    "        loss = loss_fn(output, target)\n",
    "\n",
    "    # Exits the context manager before backward()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "```\n",
    "Which is essentially telling us that we should not use `bfloat16()` on any of our tensors and just use the context manager and **only surround the forward pass of the model and the loss calculation of the model** within it and **leave the backward pass and the optimizer step alone**...\n",
    "\n",
    "And we will be following that guidance now..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our old training loop looks like this for now:\n",
    "```python\n",
    "# Optimization\n",
    "epochs = 50\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4)\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    inputs, labels = trainingLoader.nextBatch()\n",
    "    inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "    optimizer.zero_grad()\n",
    "    logits, loss = model(inputs, labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) / (endTime - startTime)\n",
    "    print(f\"Step: {epoch}, Loss: {loss.item()}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```\n",
    "\n",
    "And in our case we already have both the forward pass of the model and the loss calculation of the model in a single call (`logits, loss = model(inputs, labels)`) we will only surround that line with the context manager like this:\n",
    "```python\n",
    "with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "    logits, loss = model(inputs, labels)\n",
    "```\n",
    "\n",
    "So, our new code training loop like:\n",
    "```python\n",
    "# Optimization\n",
    "epochs = 50\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4)\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    inputs, labels = trainingLoader.nextBatch()\n",
    "    inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "    optimizer.zero_grad()\n",
    "    with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "        logits, loss = model(inputs, labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) / (endTime - startTime)\n",
    "    print(f\"Step: {epoch}, Loss: {loss.item()}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And I would also like to point out that not everything get's converted to `BF16` here, and for now the documentation tells us that only these things get converted into `BF16`:\n",
    "`__matmul__`, `addbmm`, `addmm`, `addmv`, `addr`, `baddbmm`, `bmm`, `chain_matmul`, `multi_dot`, `conv1d`, `conv2d`, `conv3d`, `conv_transpose1d`, `conv_transpose2d`, `conv_transpose3d`, `GRUCell`, `linear`, `LSTMCell`, `matmul`, `mm`, `mv`, `prelu`, `RNNCell`\n",
    "\n",
    "And a lot of operations remain in `float32`:\n",
    "`__pow__`, `__rdiv__`, `__rpow__`, `__rtruediv__`, `acos`, `asin`, `binary_cross_entropy_with_logits`, `cosh`, `cosine_embedding_loss`, `cdist`, `cosine_similarity`, `cross_entropy`, `cumprod`, `cumsum`, `dist`, `erfinv`, `exp`, `expm1`, `group_norm`, `hinge_embedding_loss`, `kl_div`, `l1_loss`, `layer_norm`, `log`, `log_softmax`, `log10`, `log1p`, `log2`, `margin_ranking_loss`, `mse_loss`, `multilabel_margin_loss`, `multi_margin_loss`, `nll_loss`, `norm`, `normalize`, `pdist`, `poisson_nll_loss`, `pow`, `prod`, `reciprocal`, `rsqrt`, `sinh`, `smooth_l1_loss`, `soft_margin_loss`, `softmax`, `softmin`, `softplus`, `sum`, `renorm`, `tan`, `triplet_margin_loss`\n",
    "\n",
    "Now, that we have everything, we an take our model for a spin...\n",
    "\n",
    "And for me, I get:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Time Difference: 956.87ms, Tokens/Second: 16729.79\n",
    "Step: 1, Loss: 9.640329360961914, Time Difference: 299.31ms, Tokens/Second: 54981.59\n",
    "Step: 2, Loss: 9.391995429992676, Time Difference: 299.82ms, Tokens/Second: 54991.51\n",
    "Step: 3, Loss: 8.106719017028809, Time Difference: 299.86ms, Tokens/Second: 54916.08\n",
    "Step: 4, Loss: 8.977012634277344, Time Difference: 299.09ms, Tokens/Second: 54915.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we are definately running faster, but not a lot faster, and that's because there are many many bottlenecks in our `GPT-2` and we are just getting started. But we have dropped down the precision as far as we can. And we can just run the training longer to make up for that precision decrease as well...\n",
    "\n",
    "And we will start to use some of the heavy weaponary now in PyTorch..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using <a href=\"https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html\">`torch.compile`</a> & Understanding **Kernel Fusion**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html\">`torch.compile`</a> is quite infrastructure from the PyTorch team and it's basically a compiler for neural networks...\n",
    "\n",
    "And all we need to do is write this single line of code:\n",
    "```python\n",
    "model = torch.compile(model)\n",
    "```\n",
    "\n",
    "Now, this line of code is going to cost us **compilation time**, but as you can guess, this line of code can execute the code a lot faster..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we run our code now, we get:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Time Difference: 31098.87ms, Tokens/Second: 565.79\n",
    "Step: 1, Loss: 9.640329360961914, Time Difference: 145.31ms, Tokens/Second: 112949.59\n",
    "Step: 2, Loss: 9.391995429992676, Time Difference: 129.82ms, Tokens/Second: 126679.51\n",
    "Step: 3, Loss: 8.106719017028809, Time Difference: 129.86ms, Tokens/Second: 126984.08\n",
    "Step: 4, Loss: 8.977012634277344, Time Difference: 129.09ms, Tokens/Second: 126984.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I'd also like to explain what `torch.compile()` does under the hood...\n",
    "\n",
    "And I think we all should just use `torch.compile()` by default, unless we are debugging...\n",
    "\n",
    "And there's one line in this documentation that tells us why this `torch.compile()` is faster:\n",
    "> Speedup mainly comes from reducing Python overhead and GPU read/writes, and so the observed speedup may vary on factors such as model architecture and batch size.\n",
    "\n",
    "So, what is happening under the hood?\n",
    "\n",
    "Well, when we pass the model to `torch.compile()`, all of our code is just algorithmic description of what we'd like to happen in our network and `torch.compile()` will analyze the entire thing and will see what operations we'd like to use and with the benefit of knowing what's going to happen, it doesn't have to run in *\"eager-mode\"*, and it would optimize that process. The first thing it will do is, it will take out the Python interpreter and it will kind of compile this entire neural network as a single object with no Python interpreter involved (such that it knows exactly what it's going to run, and will just run that). And it's all going to be running in effecient code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second thing that happens is this Read/Write that they mention in the documentation above...\n",
    "\n",
    "Let's take a small *toy-example* to explain this:\n",
    "```python\n",
    "def powerOfThree(inputs):\n",
    "    return torch.pow(inputs, 3.0)\n",
    "```\n",
    "\n",
    "In the above example, the Python interpreter will make it's way to the `torch.pow(inputs, 3.0)`, and it's going to dispatch a kernel that takes the `input` and raise it to the power of `3.0`, and that kernel will run, and when this kernel runs, what ends up happening is, this `input` is stored in the memory of the `GPU`...\n",
    "\n",
    "So here's a helpful example of the layout of what's happening...\n",
    "\n",
    "![GPU-memory-oversubscription](ExplanationMedia/Images/GPU-memory-oversubscription.png)\n",
    "\n",
    "We have our `CPU` connected to `RAM` which is already well known to us. And now, we have added the `GPU` to the system, and the `GPU` has a slightly different architecture than `CPU`, but the `GPU` and `CPU` communicate with each other. And the `GPU` is different from `CPU` in the sense that it has a lot more cores than the `CPU`, and all of those cores are individually a lot simpler but it also has a memory called `High Bandwidth Memory (HBM)` which is very equivalent to the `RAM` in the computer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And what's happening is that `input` is living in the `HBM` and this has to travel through the `GPU` throught the cores and all the **caches & registers** and it has to calculate all the elements to the power of `3.0` and it saves the result back to the `HBM`. And it's this travel time that actually causes a lot of issues and is very slow. And we keep on doing multiple round trips for a single operation.\n",
    "\n",
    "And PyTorch, without using `torch.compile()`, doesn't know how to optimize this, because it doesn't know what kind of operations are we running later. And what `torch.compile()` will try to do a single round trip to calculate and save a large operation. And it's one example of what's called a **Kernel Fusion** and is a major way in which everything is sped up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we move on to the next part, I wanted to suppliment the discussion of what's happening here...\n",
    "\n",
    "So, the `GPU` that we see in the above diagram, also does have a little bit of memory inside it. But most of the memory by far is outside the chip, which is inside the `HBM` and they are connected to each other (they are two separate chips).\n",
    "\n",
    "And this is a zoom in of a cartoon diagram of a `GPU`:\\\n",
    "![SM_GA100](ExplanationMedia/Images/SM_GA100.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You see those black bars on the sides, those are `HBMs`, which is still outside the actual `GPU`, but inside the `GPU` there are a large number of these `Streaming-Multiprocessors (SM)`, which is where a lot of the computations happen.\n",
    "\n",
    "And on the left, we have a zoom-in of a single `SM`. Now, we have all this logic here to do the calculations but in addition to that we still have memory sprinkled throughout the chip.\n",
    "\n",
    "So, here `L2 Cache` is some amount of memory that lives inside the `GPU`, and then, on the `SMs` themselves they have the `L1 Cache` and there's also registers. But the memory stored in these stores are very different from the way memory is stored in the `HBM`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And this is a memory diagram for a slightly different `GPU`:\\\n",
    "![GPU_Memory_Bandwidth_Hierarchy](ExplanationMedia/Images/GPU_Memory_Bandwidth_Hierarchy.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And it shows the typical numbers for a `CPU - DRAM` memory which is very expensive to access. And everything is extremely fast within the chip but we only have small amounts of memory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Flash Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, `torch.compile()` is amazing, but there are operations that `torch.compile()` will not find, and an amazing example of that is **`Flash Attention`**.\n",
    "\n",
    "And `Flash Attention` comes from <a href=\"https://arxiv.org/abs/2205.14135#\">this paper</a> from Stanforn in 2022.\n",
    "\n",
    "And flash attention will take these lines of code:\n",
    "```python\n",
    "attention = (query @ key.transpose(-2, -1)) * (1.0 / math.sqrt(key.size(-1)))\n",
    "attention = attention.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))\n",
    "attention = F.softmax(attention, dim=-1)\n",
    "outputs = attention @ value\n",
    "```\n",
    "And flash attention will implement these lines really really quickly...\n",
    "\n",
    "And how does it do that?\n",
    "\n",
    "Well `Flash Attention` is a **kernel fusion** operation:\n",
    "![Flash_Attention](ExplanationMedia/Images/Flash_Attention.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the reason that `torch.compile()` cannot find it is that it requires an **algorithmic re-write** of how attention is implemented in this case. And what's remarkable about it is, `Flash Attention` does more `FLOPS` the attention we have implemented, but `Flash Attention` is significanly faster, and that's because they are very mindful of the memory hierarchy we discussed earlier.\n",
    "\n",
    "And the $N \\times N$ attention matrix is never *materialized* at any point and never gets read or written to the `HBM`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the fundamental algorithmic rewrite relies on this ***online-softmax*** trick, which shows how you can incrementally evaluate a softmax without having to realize all of the inputs to the softmax to do the normalization, and you do that some intermediate variables $m$ & $l$ and there's an update to them that allows you to evaluate the softmax in an incremental manner.\n",
    "\n",
    "And recently <a href=\"https://arxiv.org/pdf/2307.08691\">Flash Attention - 2</a> came out as well. And the original paper that this is based on is - <a href=\"https://arxiv.org/abs/1805.02867\">Online normalizer calculation for softmax</a>, and remarkably it came out of NVIDIA and really early(2018) which is `4` years before `Flash Attention`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the way to implement `Flash Attention` is to take out these lines:\n",
    "```python\n",
    "attention = (query @ key.transpose(-2, -1)) * (1.0 / math.sqrt(key.size(-1)))\n",
    "attention = attention.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))\n",
    "attention = F.softmax(attention, dim=-1)\n",
    "outputs = attention @ value\n",
    "```\n",
    "\n",
    "And replace it with a single line:\n",
    "```python\n",
    "outputs = F.scaled_dot_product_attention(query=query, key=key, value=value, is_causal=True)\n",
    "```\n",
    "\n",
    "And you see how we are setting `is_causal=True`, well that is the traingular mask, and we have a small cleanup to do, which is this `bias` buffer we created during the construction of the class.\n",
    "```python\n",
    "self.register_buffer(\"bias\", torch.tril(torch.ones(configuration.blockSize, configuration.blockSize)).view(1, 1, configuration.blockSize, configuration.blockSize))\n",
    "```\n",
    "\n",
    "So now if we try to run our code, we get:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Time Difference: 24098.87ms, Tokens/Second: 650.79\n",
    "Step: 1, Loss: 9.640329360961914, Time Difference: 106.31ms, Tokens/Second: 152949.59\n",
    "Step: 2, Loss: 9.391995429992676, Time Difference: 96.82ms, Tokens/Second: 166679.51\n",
    "Step: 3, Loss: 8.106719017028809, Time Difference: 95.86ms, Tokens/Second: 176984.08\n",
    "Step: 4, Loss: 8.977012634277344, Time Difference: 96.09ms, Tokens/Second: 176984.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nice & Ugly Numbers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now getting to one of my favourite optimizations, and it's simultaniously the dumbest and the most brilliant optimizations ever...\n",
    "\n",
    "And I previously mentioned that some numbers are nice and some numbers that are ugly...\n",
    "\n",
    "So `64` is a nice number, `128` is a nice number, `256` is a beautiful number... And what makes these numbers nice is that, they are the powers of `2` and you can divide these numbers by `2` many times...\n",
    "\n",
    "And examples of ugly numbers are `13` and `17` and so on (Prime Numbers, Numbers that are not even and so on)...\n",
    "\n",
    "And we should always want to use nice numbers when we are dealing with neural networks or `CUDA` because, everything in `CUDA` works in the powers of `2` and lots of `Kernels` are written in terms of powers of `2` and there's lots of extra logic that's being written in the handling of ugly numbers...\n",
    "\n",
    "So, let's see what that looks like..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's scan our code and look for ugly numbers first...\n",
    "\n",
    "Here `3` is an ugly number:\n",
    "```python\n",
    "self.causalAttention = torch.nn.Linear(configuration.numberOfEmbeddingDimensions, 3 * configuration.numberOfEmbeddingDimensions)\n",
    "```\n",
    "\n",
    "Here `1024` is a nice number, `50257` is a very ugly number, `12` is a suspisiously ugly number and `768` is a nice number:\n",
    "```python\n",
    "# GPT configuration hyper-parameters\n",
    "@dataclass\n",
    "class GPTConfiguration:\n",
    "    blockSize: int = 1024\n",
    "    vocabularySize: int = 50257\n",
    "    numberOfLayers: int = 12\n",
    "    numberOfHeads: int = 12\n",
    "    numberOfEmbeddingDimensions: int = 768\n",
    "    NANOGPT_SCALE_INIT: bool = True\n",
    "```\n",
    "\n",
    "And here we have mostly nice numbers, except for `25`:\n",
    "```python\n",
    "configurationArguements = {\n",
    "    'gpt2':         dict(numberOfLayers=12, numberOfHeads=12, numberOfEmbeddingDimensions=768),  # 124M parameters\n",
    "    'gpt2-medium':  dict(numberOfLayers=24, numberOfHeads=16, numberOfEmbeddingDimensions=1024), # 350M parameters\n",
    "    'gpt2-large':   dict(numberOfLayers=36, numberOfHeads=20, numberOfEmbeddingDimensions=1280), # 774M parameters\n",
    "    'gpt2-xl':      dict(numberOfLayers=48, numberOfHeads=25, numberOfEmbeddingDimensions=1600), # 1558M parameters\n",
    "}[modelType]\n",
    "```\n",
    "\n",
    "And some of these numbers are quite easier to fix than the other...\n",
    "\n",
    "And the way to fix these numbers, is that we increase the numbers such that they reach the nearest powers of `2`..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's fix the biggest number first (the **vocabulary size**) which is `50257`...\n",
    "\n",
    "And the nicer number for this would be `50304`.\n",
    "\n",
    "And why is that?\n",
    "\n",
    "It's because `50304` can be divided by:\n",
    "$$50304 / 8 = 6288 \\\\ 50304 / 16 = 3144 \\\\ 50304 / 32 = 1572 \\\\ 50304 / 64 = 786 \\\\ 50304 / 128 = 393$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, how do we increase the `vocabularySize` now?\n",
    "\n",
    "Remember how `GPTConfiguration` is a `@dataclass`?\n",
    "\n",
    "Previously we initialized our model like this:\n",
    "```python\n",
    "# Constructing Model\n",
    "model = GPTModel(GPTConfiguration())\n",
    "```\n",
    "Now we can simply override the external arguement of this configuration like:\n",
    "```python\n",
    "# Constructing Model\n",
    "model = GPTModel(GPTConfiguration(vocabularySize=50304))\n",
    "```\n",
    "And this operation is like we are adding **fake tokens** to it. But it also does increase the amount of computation that we will be doing (more `FLOPS`) and we still have to think through if this breaks anything or not...\n",
    "\n",
    "Let's run this first and I get:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Time Difference: 23098.87ms, Tokens/Second: 650.79\n",
    "Step: 1, Loss: 9.640329360961914, Time Difference: 108.31ms, Tokens/Second: 162949.59\n",
    "Step: 2, Loss: 9.391995429992676, Time Difference: 92.82ms, Tokens/Second: 176679.51\n",
    "Step: 3, Loss: 8.106719017028809, Time Difference: 92.86ms, Tokens/Second: 176984.08\n",
    "Step: 4, Loss: 8.977012634277344, Time Difference: 92.09ms, Tokens/Second: 176984.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see why this didn't break anything in our code...\n",
    "\n",
    "Well we are using `vocabularySize` during the **Embedding Table** (bottom) and the **Classifier Layer** (top)...\n",
    "\n",
    "And these `tokens` are never used and we will never index into the rows of what we have added."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all, by adding **fake tokens** to our vocabulary, we have made our **Embedding Table** larger, meaning that we will never index into the rows of the **Embedding Table** for the tokens that are never used, but we are also wasting space, which is also not fully correct because we are using the **weight sharing scheme**, which again ends up being used in the **Classifier Layer**, which is predicting additional dimensions at the layer, and ends up predicting probabilities for tokens that will never be present in the training set, meaning that the network will have to learn that these probabilities have to be driven to `0`, which is no different from the tokens that are already in our dataset. So, functionally nothing breaks, but we are using a bit more extra memory and is a harmless operation..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is also something that `torch.compile()` doesn't find for us..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `GPT-3` Hyper-parameters & Techniques"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will also turn to the <a href=\"https://arxiv.org/abs/2005.14165\">GPT-3 paper</a>, because there is much less information about the hyperparameters in the <a href=\"https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf\">GPT-2 paper</a> and the <a href=\"https://github.com/openai/gpt-2/blob/master/src/model.py\">code</a>.\n",
    "\n",
    "And `GPT-3` papers consider a more detailed approach to training but sadly, the `GPT-3` models were never released.\n",
    "\n",
    "Meaning, for `GPT-2` we have the weights and no details, `GPT-3` we have lots of details but no weights...\n",
    "\n",
    "But, `GPT-2` and `GPT-3` architectures are very very similar and there are a very few changes.\n",
    "\n",
    "The context length was extended from `1024` to `2048`, and some of the hyperparameters around the network has changed. Otherwise, they're pretty much the same model, it's just `GPT-3` was trained for a lot longer on a bigger dataset and has a lot more thorough evaluations. And the model itself is `175B` instead of `1.6B`.\n",
    "\n",
    "So, let's go through the paper to follow along with some of the hyper-parameters..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And if we find the section **Details of Model Training** within the paper, we'd see the line:\n",
    "> To train all versions of GPT-3, we use Adam with $\\beta_{1} = 0.9$,$\\beta_{2} = 0.95$, and $\\epsilon = 10^{−8}$, we clip the global norm of the gradient at 1.0, and we use cosine decay for learning rate down to 10% of its value, over 260 billion tokens (after 260 billion tokens, training continues at 10% of the original learning rate). There is a linear LR warmup over the first 375 million tokens. We also gradually increase the batch size linearly from a small value (32k tokens) to the full value over the first 4-12 billion tokens of training, depending on the model size. Data are sampled without replacement during training (until an epoch boundary is reached) to minimize overfitting. All models use weight decay of 0.1 to provide a small amount of regularization.\\\n",
    "During training we always train on sequences of the full $n_{ctx} = 2048$ token context window, packing multiple documents into a single sequence when documents are shorter than 2048, in order to increase computational efficiency. Sequences with multiple documents are not masked in any special way but instead documents within a sequence are delimited with a special end of text token, giving the language model the information necessary to infer that context separated by the end of text token is unrelated. This allows for efficient training without need for any special sequence-specific masking.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer - (`AdamW`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we see that they use Adam with $\\beta_{1} = 0.9$,$\\beta_{2} = 0.95$. \n",
    "\n",
    "And we can now specify these hyperparameters for the `optimizer` using the `betas` parameter which defaults to `(0.9, 0.999)`, like this:\n",
    "```python\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4, betas=(0.9, 0.95))\n",
    "```\n",
    "\n",
    "We also see that they use the $\\epsilon = 10^{−8}$ for the `optimizer` as well.\n",
    "\n",
    "And we can now explicitly specify the `eps` parameter within the `optimizer` to set it as well, like this:\n",
    "```python\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4, betas=(0.9, 0.95), eps=1e-8)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradient Clipping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see the part:\n",
    ">\"we clip the global norm of the gradient at 1.0\"\n",
    "\n",
    "And what this is referring to, is once we calculate the gradients right after `loss.backward()`, what people like to do is **clip them to have some kind of a maximum normalization**.\n",
    "\n",
    "And this is fairly easy to do with PyTorch and is basically one line of code, which we will insert right after `loss.backward()`:\n",
    "```python\n",
    "normalization = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1.0)\n",
    "```\n",
    "\n",
    "And what this utility function is doing, is that it's calculating the global normalization of the parameters, so every single gradient of the parameters are squared and added up and the square root of that is taken into account which ends up being the normalization of the parameter vector (it's the length of it). And we are basically making sure that it's length is no more than `1.0` and we are going to clip it.\n",
    "\n",
    "And the reason that people like to use this is because, sometimes we can get unlucky during the optimization. Maybe if someone gets unlucky in a batch, they get a really high loss and the really high loss can lead to a really high gradient which could basically *shock* the model and the optimization. \n",
    "\n",
    "So, people like to use **gradient normalization clipping** to prevent the model from basically getting too big of *shocks* in terms of the gradient magnitude they're upper bound in the above way..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's also visualize this normalization by printing it in the same line like this:\n",
    "```python\n",
    "print(f\"Step: {epoch}, Loss: {loss.item()}, Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```\n",
    "\n",
    "And we get:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Normalization: 28.9463, Time Difference: 24098.87ms, Tokens/Second: 675.79\n",
    "Step: 1, Loss: 9.640329360961914, Normalization: 6.9656, Time Difference: 108.31ms, Tokens/Second: 152949.59\n",
    "Step: 2, Loss: 9.391995429992676, Normalization: 2.9491, Time Difference: 107.56ms, Tokens/Second: 156679.51\n",
    "Step: 3, Loss: 8.106719017028809, Normalization: 2.1653, Time Difference: 94.16ms, Tokens/Second: 174984.08\n",
    "Step: 4, Loss: 8.977012634277344, Normalization: 1.8498, Time Difference: 93.09ms, Tokens/Second: 174984.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the normalization is really high in the beginning but then it stabilizes...\n",
    "\n",
    "Which is fine, and we can now move on to the next part..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning Rate Scheduler (Warmup + Cosine Decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now notice the line:\n",
    ">\"we use cosine decay for learning rate down to 10% of its value, over 260 billion tokens (after 260 billion tokens, training continues at 10% of the original learning rate)\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the Learning Rate Scheduler that's been used in the `GPT-3` is called a cosine decay learning schedule with warmup.\n",
    "\n",
    "And it looks something like this:\\\n",
    "![Cosine_Learning_Rate_Decay](ExplanationMedia/Images/Cosine_Learning_Rate_Decay.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, the learning rate starts at `0`, linearly ramps up over some amount of time and comes down in like a cosine sort of form to a minimum learning rate that's upto us...\n",
    "\n",
    "And we can now start implementing that..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And PyTorch actually has learning rate schedulers and you can use them as well.\n",
    "\n",
    "But I don't really love using that because it's honestly just 4-5 lines of code and I fully understand what's happening inside these lines and I don't like to use abstractions and fully know what is doing what...\n",
    "\n",
    "And you might also remember that previously we used the learning rate of `3e-4`, but we see that in `GPT-3` paper we have this table:\n",
    "\n",
    "| Model Name     | $n_{params}$ | $n_{layers}$ | $d_{model}$ | $n_{heads}$ | $d_{head}$ | Batch Size | Learning Rate    |\n",
    "|----------------|---------|---------|--------|--------|-------|------------|------------------|\n",
    "| GPT-3 Small    | 125M    | 12      | 768    | 12     | 64    | 0.5M       | $6.0 \\times 10^{-4}$       |\n",
    "| GPT-3 Medium   | 350M    | 24      | 1024   | 16     | 64    | 0.5M       | $3.0 \\times 10^{-4}$       |\n",
    "| GPT-3 Large    | 760M    | 24      | 1536   | 16     | 96    | 0.5M       | $2.5 \\times 10^{-4}$       |\n",
    "| GPT-3 XL       | 1.3B    | 24      | 2048   | 24     | 128   | 1M         | $2.0 \\times 10^{-4}$       |\n",
    "| GPT-3 2.7B     | 2.7B    | 32      | 2560   | 32     | 80    | 1M         | $1.6 \\times 10^{-4}$       |\n",
    "| GPT-3 6.7B     | 6.7B    | 32      | 4096   | 32     | 128   | 2M         | $1.2 \\times 10^{-4}$       |\n",
    "| GPT-3 13B      | 13.0B   | 40      | 5140   | 40     | 128   | 2M         | $1.0 \\times 10^{-4}$       |\n",
    "| GPT-3 175B     | 175.0B  | 96      | 12288  | 96     | 128   | 3.2M       | $0.6 \\times 10^{-4}$       |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we see that they used $6.0 \\times 10^{-4}$ for a `GPT-3 125M` model which is similar to our `GPT-2 124M` model, and we can just go ahead and change the learning rate to `6e-4` which is the maximum learning rate..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now set our own minimum learning rate to `10%` of the maximum learning rate as per the description in the paper.\n",
    "\n",
    "And we will also specify some number of `warmupEpochs` to warmup our `epochs` over...\n",
    "\n",
    "Let's now change the learning rate at each iteration by introducing this code:\n",
    "```python\n",
    "...\n",
    "learningRate = getLearningRate(epoch=epoch)\n",
    "for parameterGroup in optimizer.param_groups:\n",
    "    parameterGroup['lr'] = learningRate\n",
    "optimizer.step()\n",
    "...\n",
    "```\n",
    "\n",
    "You see how `getLearningRate()` function is not defined?\n",
    "\n",
    "Let's implement that now..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, that we have the markers (`warmupEpochs` and `epochs`), we can go ahead and select the area between them and specify which section we want to work on based on the `epoch`:\n",
    "```python\n",
    "def getLearningRate(epoch):\n",
    "    if epoch < warmupEpochs:\n",
    "        return maximumLearningRate * (epoch + 1) / warmupEpochs\n",
    "    if epoch > epochs:\n",
    "        return minimumLearningRate\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, the section within the `warmupEpochs` is gradually increased overtime...\n",
    "\n",
    "And the second checker checks if the `epoch` exceeds the epochs to continue returning the `minimumLearningRate` just in case...\n",
    "\n",
    "And we can now implement logic for the cosine decay in between..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And it looks like this:\n",
    "```python\n",
    "def getLearningRate(epoch):\n",
    "    if epoch < warmupEpochs:\n",
    "        return maximumLearningRate * (epoch + 1) / warmupEpochs\n",
    "    if epoch > epochs:\n",
    "        return minimumLearningRate\n",
    "    decayRatio = (epoch - warmupEpochs) / (epochs - warmupEpochs)\n",
    "    assert 0 <= decayRatio <= 1\n",
    "    coefficient = 0.5 * (1.0 + math.cos(math.pi * decayRatio))\n",
    "    return minimumLearningRate + coefficient * (maximumLearningRate - minimumLearningRate)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can go through this code but, we are just determining the coefficient of the curve using the decay ratio at each epoch and returning the state of the learning rate at that epoch..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also print the Learning Rate:\n",
    "```python\n",
    "print(f\"Step: {epoch}, Loss: {loss.item()}, Learning Rate: {learningRate:.4e},  Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So our optimization looks like this:\n",
    "```python\n",
    "# Optimization\n",
    "maximumLearningRate = 6e-4\n",
    "minimumLearningRate = maximumLearningRate * 0.1\n",
    "warmupEpochs = 10\n",
    "epochs = 50\n",
    "def getLearningRate(epoch):\n",
    "    if epoch < warmupEpochs:\n",
    "        return maximumLearningRate * (epoch + 1) / warmupEpochs\n",
    "    if epoch > epochs:\n",
    "        return minimumLearningRate\n",
    "    decayRatio = (epoch - warmupEpochs) / (epochs - warmupEpochs)\n",
    "    assert 0 <= decayRatio <= 1\n",
    "    coefficient = 0.5 * (1.0 + math.cos(math.pi * decayRatio))\n",
    "    return minimumLearningRate + coefficient * (maximumLearningRate - minimumLearningRate)\n",
    "\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4, betas=(0.9, 0.95), eps=1e-8)\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    inputs, labels = trainingLoader.nextBatch()\n",
    "    inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "    optimizer.zero_grad()\n",
    "    with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "        logits, loss = model(inputs, labels)\n",
    "    loss.backward()\n",
    "    normalization = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1.0)\n",
    "    learningRate = getLearningRate(epoch=epoch)\n",
    "    for parameterGroup in optimizer.param_groups:\n",
    "        parameterGroup['lr'] = learningRate\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) / (endTime - startTime)\n",
    "    print(f\"Step: {epoch}, Loss: {loss.item()}, Learning Rate: {learningRate:.4e},  Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's now run it:\n",
    "```python\n",
    "Using Device: cuda\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 128 Batches\n",
    "Step: 0, Loss: 10.832972526550293, Learning Rate: 3.0000e-05, Normalization: 28.9463, Time Difference: 24098.87ms, Tokens/Second: 675.79\n",
    "Step: 1, Loss: 9.640329360961914, Learning Rate: 6.0000e-05, Normalization: 13.9656, Time Difference: 108.31ms, Tokens/Second: 142949.59\n",
    "Step: 2, Loss: 9.391995429992676, Learning Rate: 9.0000e-05, Normalization: 12.9491, Time Difference: 107.56ms, Tokens/Second: 151679.51\n",
    "Step: 3, Loss: 8.106719017028809, Learning Rate: 1.2000e-05, Normalization: 9.1653, Time Difference: 94.16ms, Tokens/Second: 174984.08\n",
    "Step: 4, Loss: 8.977012634277344, Learning Rate: 1.5000e-05, Normalization: 4.8498, Time Difference: 93.09ms, Tokens/Second: 174984.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we don't start the learning rate exactly at `0` because that wouldn't be useful.\n",
    "\n",
    "Then we linearly ramp up to the maximum learning rate (`6e-4`) and then comes down to the minimum learning rate (`6e-5` or `10% of 6e-4`)...\n",
    "\n",
    "And we are not exactly following the paper because, honestly I don't think this makes too big of a difference..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batch Scheduling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, the paper talks about gradual batch size increase:\n",
    "> We also gradually increase the batch size linearly from a small value (32k tokens) to the full value over\n",
    "the first 4-12 billion tokens of training, depending on the model size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are not going to use this, but I would like to discuss this here as well...\n",
    "\n",
    "And the reason we are not using this is because this complicates a lot of the arithmetic that we already have in our code. And also to my understanding this does not make that much of a difference...\n",
    "\n",
    "Basically here, we start with a small batch size and we gradually ramp up to a bigger batch size over time...\n",
    "\n",
    "And this is totally to improve the performance of the system...\n",
    "\n",
    "In the early stages of optimization the model is mostly learning to ignore the tokens that don't come up in our training set, and learning very simple biases and so the gradients of every single example that we throw into the network are extremely correlated and they all look roughly the same.\n",
    "\n",
    "And because they are extremely correlated, we basically get the exact same gradient in the early stages and doesn't require complete batches to determine the gradients that much...\n",
    "\n",
    "And that is the reason we use batch scheduling to schedule the batches in a way that their sizes ramp up over time..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weight Decay & `FusedAdamW`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we come to this line:\n",
    "> All models use weight decay of 0.1 to provide a\n",
    "small amount of regularization.\n",
    "\n",
    "Let's now implement this..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember how our `optimizer` object was just hanging around in the open like this:\n",
    "```python\n",
    "optimizer = torch.optim.AdamW(params=model.parameters(), lr=3e-4, betas=(0.9, 0.95), eps=1e-8)\n",
    "```\n",
    "\n",
    "Let's now create a method with the name `configureOptimizers()` inside of our model to move this code inside our `GPTModel` class and return the `optimizer` object once we call it...\n",
    "\n",
    "And this method will take the `weightDecayRate`, `learningRate` and `device` as arguements...\n",
    "\n",
    "So our `GPTModel` class looks like this for now:\n",
    "```python\n",
    "class GPTModel(torch.nn.Module):\n",
    "    def __init__(self, configuration):\n",
    "        ...\n",
    "\n",
    "    def forward(self, indeces, labels=None):\n",
    "        ...\n",
    "\n",
    "    def _initializeParameters(self, module):\n",
    "        ...\n",
    "\n",
    "    # Method to transfer weights from Hugging Face GPT-2\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, modelType):\n",
    "        ...\n",
    "    \n",
    "    def configureOptimizers(self, weightDecayRate, learningRate, device):\n",
    "        pass\n",
    "```\n",
    "\n",
    "Let's now implement this method..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we have to have some way of specifying the parameters that:\n",
    "1. Should be weight decayed\n",
    "2. Should not be weight decayed\n",
    "\n",
    "Don't worry I've got you covered before the next question comes to mind...\n",
    "\n",
    "It is common to **not** weight decay `biases` and any other `1D` tensors (layer normalizations, scales & biases).\n",
    "\n",
    "It is common to weight decay the tensors that participate in matrix multiplications and the embeddings...\n",
    "\n",
    "And we covered this already in our previous notebooks that we can think of this process as a bit of a regularization, because when we are pulling down all the weights we are forcing the optimization to use more of the weights and we are not allowing any weight individually to be way to large and we are forcing the network to distribute the work across more channels because there is a pull of gravity of some sort..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we map the `decayParameters` and `nonDecayParameters` like this:\n",
    "```python\n",
    "parameterDictionary = {parameterNumber: parameter for parameterNumber, parameter in self.named_parameters()}\n",
    "parameterDictionary = {parameterNumber: parameter for parameterNumber, parameter in parameterDictionary.items() if parameter.requires_grad}\n",
    "\n",
    "decayParameters = [parameter for parameterNumber, parameter in parameterDictionary.items() if parameter.dim() >= 2]\n",
    "nonDecayParameters = [parameter for parameterNumber, parameter in parameterDictionary.items() if parameter.dim() < 2]\n",
    "```\n",
    "\n",
    "And map the `weightDecayRate` into the groups of parameters like this:\n",
    "```python\n",
    "optimizerGroups = [\n",
    "    {\n",
    "        'params': decayParameters, 'weight_decay': weightDecayRate\n",
    "    },\n",
    "    {\n",
    "        'params': nonDecayParameters, 'weight_decay': 0.0\n",
    "    }\n",
    "]\n",
    "```\n",
    "\n",
    "We can then print to check the number of decayed and non-decayed parameters like this:\n",
    "```python\n",
    "numberOfDecayParameters = sum(parameter.numel() for parameter in decayParameters)\n",
    "numberOfNonDecayParameters = sum(parameter.numel() for parameter in nonDecayParameters)\n",
    "\n",
    "print(f\"Number of decayed parameter tensors: {len(decayParameters)}, with {numberOfDecayParameters} parameters\")\n",
    "print(f\"Number of non-decayed parameter tensors: {len(nonDecayParameters)}, with {numberOfNonDecayParameters} parameters\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now introduce ourselves with the idea of `Fused AdamW`...\n",
    "\n",
    "We will **import** something called the <a href=\"https://docs.python.org/3/library/inspect.html\">`inspect`</a> module from Python which let's us inspect if the `'fused'` is available in our version of PyTorch, because this concept is fairly new and is not available in the earlier versions of PyTorch.\n",
    "\n",
    "And what this is doing is, instead of iterating in a loop over all the parameter tensors and updating them (which would launch a lot of kernels), all the kernels are ***fused*** into a single kernel and call something that updates them...\n",
    "\n",
    "So, once again it's a **kernel fusion** operation for the `AdamW` updates...\n",
    "\n",
    "And we can now use the `Fused AdamW` like this and return it:\n",
    "```python\n",
    "isFusedAvailable = 'fused' in inspect.signature(torch.optim.AdamW).parameters\n",
    "useFused = isFusedAvailable and device == \"cuda\"\n",
    "print(f\"Using Fused AdamW: {useFused}\")\n",
    "optimizer = torch.optim.AdamW(params=optimizerGroups, lr=learningRate, betas=(0.9, 0.95), eps=1e-8, fused=useFused)\n",
    "return optimizer\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And so the entire `configureOptimizers()` method looks like this:\n",
    "```python\n",
    "def configureOptimizers(self, weightDecayRate, learningRate, device):\n",
    "    parameterDictionary = {parameterNumber: parameter for parameterNumber, parameter in self.named_parameters()}\n",
    "    parameterDictionary = {parameterNumber: parameter for parameterNumber, parameter in parameterDictionary.items() if parameter.requires_grad}\n",
    "\n",
    "    decayParameters = [parameter for parameterNumber, parameter in parameterDictionary.items() if parameter.dim() >= 2]\n",
    "    nonDecayParameters = [parameter for parameterNumber, parameter in parameterDictionary.items() if parameter.dim() < 2]\n",
    "\n",
    "    optimizerGroups = [\n",
    "        {\n",
    "            'params': decayParameters, 'weight_decay': weightDecayRate\n",
    "        },\n",
    "        {\n",
    "            'params': nonDecayParameters, 'weight_decay': 0.0\n",
    "        }\n",
    "    ]\n",
    "    numberOfDecayParameters = sum(parameter.numel() for parameter in decayParameters)\n",
    "    numberOfNonDecayParameters = sum(parameter.numel() for parameter in nonDecayParameters)\n",
    "\n",
    "    print(f\"Number of decayed parameter tensors: {len(decayParameters)}, with {numberOfDecayParameters} parameters\")\n",
    "    print(f\"Number of non-decayed parameter tensors: {len(nonDecayParameters)}, with {numberOfNonDecayParameters} parameters\")\n",
    "\n",
    "    isFusedAvailable = 'fused' in inspect.signature(torch.optim.AdamW).parameters\n",
    "    useFused = isFusedAvailable and device == \"cuda\"\n",
    "    print(f\"Using Fused AdamW: {useFused}\")\n",
    "    optimizer = torch.optim.AdamW(params=optimizerGroups, lr=learningRate, betas=(0.9, 0.95), eps=1e-8, fused=useFused)\n",
    "    return optimizer\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we run our code by calling the optimizer like this:\n",
    "```python\n",
    "optimizer = model.configureOptimizers(weightDecayRate=0.1, learningRate=maximumLearningRate, device=device)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we see a start like this:\n",
    "```python\n",
    "Using Device: cpu\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 16409 Batches\n",
    "Number of decayed parameter tensors: 50, with 124354560 parameters\n",
    "Number of non-decayed parameter tensors: 98, with 121344 parameters\n",
    "Using Fused AdamW: True\n",
    "...\n",
    "Step: 49, Loss: 5.861645634277344, Learning Rate: 6.0831-05, Normalization: 0.7612, Time Difference: 90.09ms, Tokens/Second: 184984.71\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I would also like to point out that all the things that we have discussed in this entire section are the things that have a very complicated mathematical relationship and for now we will simply copy the settings that OpenAI have used. Because explaining the relationships alone would take their own notebooks by themselves..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradient Accumulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now go back to the table\n",
    "\n",
    "| Model Name     | $n_{params}$ | $n_{layers}$ | $d_{model}$ | $n_{heads}$ | $d_{head}$ | Batch Size | Learning Rate    |\n",
    "|----------------|---------|---------|--------|--------|-------|------------|------------------|\n",
    "| GPT-3 Small    | 125M    | 12      | 768    | 12     | 64    | 0.5M       | $6.0 \\times 10^{-4}$       |\n",
    "| GPT-3 Medium   | 350M    | 24      | 1024   | 16     | 64    | 0.5M       | $3.0 \\times 10^{-4}$       |\n",
    "| GPT-3 Large    | 760M    | 24      | 1536   | 16     | 96    | 0.5M       | $2.5 \\times 10^{-4}$       |\n",
    "| GPT-3 XL       | 1.3B    | 24      | 2048   | 24     | 128   | 1M         | $2.0 \\times 10^{-4}$       |\n",
    "| GPT-3 2.7B     | 2.7B    | 32      | 2560   | 32     | 80    | 1M         | $1.6 \\times 10^{-4}$       |\n",
    "| GPT-3 6.7B     | 6.7B    | 32      | 4096   | 32     | 128   | 2M         | $1.2 \\times 10^{-4}$       |\n",
    "| GPT-3 13B      | 13.0B   | 40      | 5140   | 40     | 128   | 2M         | $1.0 \\times 10^{-4}$       |\n",
    "| GPT-3 175B     | 175.0B  | 96      | 12288  | 96     | 128   | 3.2M       | $0.6 \\times 10^{-4}$       |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see this kind of a change as the model size increases:\n",
    "| Model Size $\\uparrow$ | $n_{params}\\uparrow$ | $n_{layers}\\uparrow$ | $d_{model}\\uparrow$ | $n_{heads}\\uparrow$ | $d_{head}\\uparrow$ | Batch Size $\\uparrow$ | Learning Rate $\\downarrow$ |\n",
    "|-----------------------|----------------------|----------------------|--------------------------|---------------------|--------------------|-----------------------|------------------|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we also see that the `124M` parameter model uses a **batch size** of `0.5M`... \n",
    "\n",
    "But this is the **batch size** in the number of tokens...\n",
    "\n",
    "And remember how we used this code:\n",
    "```python\n",
    "# Data-Loader\n",
    "Batch, Time = 16, 1024\n",
    "trainingLoader = DataLoaderLite(Batch=Batch, Time=Time)\n",
    "```\n",
    "\n",
    "Here, every single one of our rows is `1024` tokens and there are `16` rows in total...\n",
    "\n",
    "To now come to the original **batch size** of `0.5M`, in our code (keeping in mind that we want to keep the original size of `1024` tokens in a row) we can find out the number of rows for our model training batch, and what it should be:\n",
    "$$\\text{Since:}500,000 = 0.5 \\times e^6$$\n",
    "$$0.5 \\times e^6 / 1024 \\approx 488$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the problem remains that we cannot just come in and set this to something like `488` because my `GPUs` would explode...\n",
    "\n",
    "But we still want to use this **batch size**, because the **batch size** is correlated to all the hyperparameters and we want to have a faithful representation of it...\n",
    "\n",
    "And how do we do that?\n",
    "\n",
    "For that we have something called the **gradient accumulation**, and it allows us to simulate a size in a serial way of any arbitrary value that we set..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basically the idea is, instead of doing gradient calculation each time the batch passes through the model for training, we are going to do multiple forward passes and backward passes, and do the updation a single time over a number of steps, and that number of steps will be defined by dividing over the size of a *micro-batch*...\n",
    "\n",
    "I used a nice number for the batch of `0.5M`: `524288`, which is a nice number and corresponds to a power of `2`:\n",
    "$$2^{19} = 524288$$\n",
    "\n",
    "And the code looks like this:\n",
    "```python\n",
    "# Data-Loader\n",
    "totalBatchSize = 524288\n",
    "Batch, Time = 16, 1024\n",
    "assert totalBatchSize % (Batch * Time) == 0, f\"Make sure totalBatchSize is divisible by (Batch * Time)\"\n",
    "gradientAccumulationSteps = totalBatchSize // (Batch * Time)\n",
    "print(f\"Total Desired Batch Size: {totalBatchSize}\")\n",
    "print(f\"Calculated Gradient Accumulation Steps: {gradientAccumulationSteps}\")\n",
    "trainingLoader = DataLoaderLite(Batch=Batch, Time=Time)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, for our case we will be doing\n",
    "$$ 2^{19} / (16 * 1024) = 32 \\text{ steps}$$\n",
    "\n",
    "Or in other words, we will be doing $32$ forward and backward passes before doing a weight update..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the numbers calculated, we can head over and implement the **gradient accumulation** itself...\n",
    "\n",
    "Remember previously we implemented the forward pass like this:\n",
    "```python\n",
    "...\n",
    "inputs, labels = trainingLoader.nextBatch()\n",
    "inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "...\n",
    "with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "    logits, loss = model(inputs, labels)\n",
    "loss.backward()\n",
    "...\n",
    "```\n",
    "\n",
    "We have to now repeat this $32$ times, before we do everything else that follows...\n",
    "\n",
    "And because we will have to load the batch every single time, we will be moving it down...\n",
    "\n",
    "So now our loop looks like this:\n",
    "```python\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    optimizer.zero_grad()\n",
    "    for microStep in range(gradientAccumulationSteps):\n",
    "        inputs, labels = trainingLoader.nextBatch()\n",
    "        inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "        with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "            logits, loss = model(inputs, labels)\n",
    "        loss.backward()\n",
    "    normalization = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1.0)\n",
    "    ...\n",
    "```\n",
    "\n",
    "And remember how `loss.backward()` always deposits gradients (adds up) which is fine, because the gradients are followed with a normalization clipping...\n",
    "\n",
    "But this implementation is potentially incorrect... And there's a very subtle but a very deep issue here..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To demonstrate the issue, let's take a small *toy-example*.\n",
    "\n",
    "Let's say that our neural network takes `16` inputs and predicts a single (`1`) output. We feed the network `4` examples and do a **mean-squared error** loss over those examples.\n",
    "\n",
    "```python\n",
    "import torch\n",
    "\n",
    "network = torch.nn.Sequential(\n",
    "    torch.nn.Linear(16, 32),\n",
    "    torch.nn.GELU(),\n",
    "    torch.nn.Linear(32, 1)\n",
    ")\n",
    "torch.random.manual_seed(69)\n",
    "inputs = torch.randn(4, 16)\n",
    "labels = torch.randn(4, 1)\n",
    "network.zero_grad()\n",
    "predictions = network(inputs)\n",
    "loss = torch.nn.functional.mse_loss(predictions, labels)\n",
    "loss.backward()\n",
    "print(\"Gradient values for the first linear layer: \", network[0].weight.grad.view(-1)[:10])\n",
    "```\n",
    "\n",
    "Which means that we are doing a simple regression over those `4` examples...\n",
    "\n",
    "And if we run this code we get:\n",
    "```python\n",
    "Gradient values for the first linear layer: tensor([-0.0051, -0.0124, 0.0227, -0.0177, 0.0248, 0.0152, -0.0276, -0.0198, -0.0016, -0.0300])\n",
    "```\n",
    "\n",
    "And if we check the docs of `mse_loss()` method, we will see that the default parameter for `reduction` has been set to `mean`:\n",
    "```python\n",
    "mse_loss(input, target, size_average=None, reduce=None, reduction='mean') -> Tensor\n",
    "```\n",
    "\n",
    "Which means that our network is potentially trying to perform this operation:\n",
    "```python\n",
    "# Objective: reduction='mean'\n",
    "# Loss = 1/4 * [\n",
    "#     (labels[0] - predictions[0])**2 *\n",
    "#     (labels[1] - predictions[1])**2 *\n",
    "#     (labels[2] - predictions[2])**2 *\n",
    "#     (labels[3] - predictions[3])**2\n",
    "# ]\n",
    "```\n",
    "In other words, we are trying to calculate the **squared errors** of each examples, and then we are trying to take the **mean** of it to get a single `loss` value."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's convert this code into the **gradient accumulation version**:\n",
    "```python\n",
    "import torch\n",
    "\n",
    "network = torch.nn.Sequential(\n",
    "    torch.nn.Linear(16, 32),\n",
    "    torch.nn.GELU(),\n",
    "    torch.nn.Linear(32, 1)\n",
    ")\n",
    "torch.random.manual_seed(69)\n",
    "inputs = torch.randn(4, 16)\n",
    "labels = torch.randn(4, 1)\n",
    "network.zero_grad()\n",
    "for microStep in range(4):\n",
    "    prediction = network(inputs[microStep])\n",
    "    loss = torch.nn.functional.mse_loss(prediction, labels[microStep])\n",
    "    loss.backward()\n",
    "print(\"Gradient values for the first linear layer: \", network[0].weight.grad.view(-1)[:10])\n",
    "```\n",
    "Previously, we did a single batch of inputs and calculated `loss.backward()` a single time. But, now we are trying to process each example individually instead and trying to call `loss.backward()` on them many times...\n",
    "\n",
    "And if we run it now, we get:\n",
    "```python\n",
    "Gradient values for the first linear layer: tensor([0.1694, -0.1119, -0.1118, 0.1299, 0.1124, 0.0392, 0.1042, 0.1846, -0.1740, -0.1214])\n",
    "```\n",
    "\n",
    "Immediately we notice that the **gradients do not match**. \n",
    "\n",
    "And the reason that they're not the same is because this **mean** for the **mean squared error** gets lost...\n",
    "\n",
    "Meaning, now the loss objective becomes:\n",
    "```python\n",
    "# Loss Objective gets lost\n",
    "# Loss0 = (labels[0] - predictions[0])**2\n",
    "# Loss1 = (labels[1] - predictions[1])**2\n",
    "# Loss2 = (labels[2] - predictions[2])**2\n",
    "# Loss3 = (labels[3] - predictions[3])**2\n",
    "# Loss = Loss0 * Loss1 * Loss2 * Loss3 # Backward\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now fix this...\n",
    "\n",
    "To fix this we can simply introduce the **mean** before `loss.backward()` like this:\n",
    "```python\n",
    "import torch\n",
    "\n",
    "network = torch.nn.Sequential(\n",
    "    torch.nn.Linear(16, 32),\n",
    "    torch.nn.GELU(),\n",
    "    torch.nn.Linear(32, 1)\n",
    ")\n",
    "torch.random.manual_seed(42)\n",
    "inputs = torch.randn(4, 16)\n",
    "labels = torch.randn(4, 1)\n",
    "network.zero_grad()\n",
    "for microStep in range(4):\n",
    "    prediction = network(inputs[microStep])\n",
    "    loss = torch.nn.functional.mse_loss(prediction, labels[microStep])\n",
    "    loss = loss / 4\n",
    "    loss.backward()\n",
    "print(\"Gradient values for the first linear layer: \", network[0].weight.grad.view(-1)[:10])\n",
    "```\n",
    "\n",
    "And what happens now, is that we end up scaling our individual losses with the mean...\n",
    "\n",
    "So the objective now becomes:\n",
    "```python\n",
    "# Loss Objective gets lost\n",
    "# Loss0 = 1/4 * (labels[0] - predictions[0])**2\n",
    "# Loss1 = 1/4 * (labels[1] - predictions[1])**2\n",
    "# Loss2 = 1/4 * (labels[2] - predictions[2])**2\n",
    "# Loss3 = 1/4 * (labels[3] - predictions[3])**2\n",
    "# Loss = 1/4 * Loss0 * 1/4 * Loss1 * 1/4 * Loss2 * 1/4 * Loss3 # Backward\n",
    "```\n",
    "\n",
    "And when we run this, we see that the gradients are now identical:\n",
    "```python\n",
    "Gradient values for the first linear layer: tensor([-0.0051, -0.0124, 0.0227, -0.0177, 0.0248, 0.0152, -0.0276, -0.0198, -0.0016, -0.0300])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, long story short the `cross_entropy()` in the loss calculation also has a `reduction` parameter which defaults to `'mean'`. And we need to fix this, the same way we did for the *toy-example*....\n",
    "\n",
    "And so our training loop looks like:\n",
    "```python\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    optimizer.zero_grad()\n",
    "    for microStep in range(gradientAccumulationSteps):\n",
    "        inputs, labels = trainingLoader.nextBatch()\n",
    "        inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "        with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "            logits, loss = model(inputs, labels)\n",
    "        loss = loss / gradientAccumulationSteps\n",
    "        loss.backward()\n",
    "    normalization = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1.0)\n",
    "    learningRate = getLearningRate(epoch=epoch)\n",
    "    for parameterGroup in optimizer.param_groups:\n",
    "        parameterGroup['lr'] = learningRate\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) / (endTime - startTime)\n",
    "    print(f\"Step: {epoch}, Loss: {loss.item()}, Learning Rate: {learningRate:.4e},  Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```\n",
    "\n",
    "And now, because we also want to print things nicely, we want a `lossAccumulator` kind of thing, because we will only be printing the loss at the **final micro step**.\n",
    "\n",
    "So we can initialize `lossAccumulator` at `0.0`, and we can `detach()` this tensor to **detach it from current graph** of which it keeps track of for gradient calculation, and we are just trying to keep track of the values inside of this tensor and we can sum these up and print only the `lossAccumulator` instead of the original `loss`, like this:\n",
    "```python\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    optimizer.zero_grad()\n",
    "    lossAccumulator = 0.0\n",
    "    for microStep in range(gradientAccumulationSteps):\n",
    "        inputs, labels = trainingLoader.nextBatch()\n",
    "        inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "        with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "            logits, loss = model(inputs, labels)\n",
    "        loss = loss / gradientAccumulationSteps\n",
    "        lossAccumulator += loss.detach()\n",
    "        loss.backward()\n",
    "    normalization = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1.0)\n",
    "    learningRate = getLearningRate(epoch=epoch)\n",
    "    for parameterGroup in optimizer.param_groups:\n",
    "        parameterGroup['lr'] = learningRate\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) / (endTime - startTime)\n",
    "    print(f\"Step: {epoch}, Loss: {lossAccumulator.item():.6f}, Learning Rate: {learningRate:.4e},  Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```\n",
    "\n",
    "But we still have to change one thing, that is the `tokensPerSecond` throughput. Because we are doing the calculation as per `gradientAccumulationSteps`. So our code changes from `tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) / (endTime - startTime)` to:\n",
    "```python\n",
    "tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) * gradientAccumulationSteps / (endTime - startTime)\n",
    "```\n",
    "\n",
    "So now our code becomes:\n",
    "```python\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    optimizer.zero_grad()\n",
    "    lossAccumulator = 0.0\n",
    "    for microStep in range(gradientAccumulationSteps):\n",
    "        inputs, labels = trainingLoader.nextBatch()\n",
    "        inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "        with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "            logits, loss = model(inputs, labels)\n",
    "        loss = loss / gradientAccumulationSteps\n",
    "        lossAccumulator += loss.detach()\n",
    "        loss.backward()\n",
    "    normalization = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1.0)\n",
    "    learningRate = getLearningRate(epoch=epoch)\n",
    "    for parameterGroup in optimizer.param_groups:\n",
    "        parameterGroup['lr'] = learningRate\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time) * gradientAccumulationSteps / (endTime - startTime)\n",
    "    print(f\"Step: {epoch}, Loss: {lossAccumulator.item():.6f}, Learning Rate: {learningRate:.4e},  Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now if we run our code we get:\n",
    "```python\n",
    "Using Device: cpu\n",
    "Total Desired Batch Size: 524288\n",
    "Calculated Gradient Accumulation Steps: 4096\n",
    "Loaded 2100390 Tokens\n",
    "1 Epoch = 16409 Batches\n",
    "Number of decayed parameter tensors: 50, with 124354560 parameters\n",
    "Number of non-decayed parameter tensors: 98, with 121344 parameters\n",
    "Using Fused AdamW: True\n",
    "Step: 0, Loss: 10.832972, Learning Rate: 6.0000e-05, Normalization: 27.9463, Time Difference: 26981.87ms, Tokens/Second: 19849.79\n",
    "Step: 1, Loss: 9.640329, Learning Rate: 1.2000e-04, Normalization: 9.9656, Time Difference: 2849.31ms, Tokens/Second: 182949.59\n",
    "Step: 2, Loss: 9.391995, Learning Rate: 1.0000e-04, Normalization: 5.9491, Time Difference: 2894.56ms, Tokens/Second: 181679.51\n",
    "Step: 3, Loss: 8.106719, Learning Rate: 2.4000e-04, Normalization: 8.1653, Time Difference: 2891.16ms, Tokens/Second: 184984.08\n",
    "Step: 4, Loss: 7.977012, Learning Rate: 3.0000e-04, Normalization: 4.8498, Time Difference: 2828.09ms, Tokens/Second: 184984.71\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we were able to fit our `GPU` with the desired **batch size** with the help of **gradient accumulation**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a href=\"https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html\">Distributed Data Parallel (DDP)</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding DDP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, it's time to bring out the heavy weapons...\n",
    "\n",
    "And you've probably noticed that we've only been using a single `GPU` for our tasks:\n",
    "```bash\n",
    "+---------------------------------------------------------------------------------+\n",
    "| NVIDIA-SMI 535.129.03       Driver Version: 535.129.03    CUDA Version: 12.2    |\n",
    "|-----------------------------------+----------------------+----------------------+\n",
    "| GPU  Name            Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
    "| Fan  Temp  Perf      Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
    "|                                   |                      |               MIG M. |\n",
    "|===================================+======================+======================|\n",
    "|   0  NVIDIA A100-SXM4-80GB    On  | 00000000:1E:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        64W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   1  NVIDIA A100-SXM4-80GB    On  | 00000000:1F:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        66W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   2  NVIDIA A100-SXM4-80GB    On  | 00000000:20:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        64W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   3  NVIDIA A100-SXM4-80GB    On  | 00000000:21:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        66W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   4  NVIDIA A100-SXM4-80GB    On  | 00000000:22:00.0 Off |                    0 |\n",
    "| N/A   32C    P0        65W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   5  NVIDIA A100-SXM4-80GB    On  | 00000000:23:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        66W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   6  NVIDIA A100-SXM4-80GB    On  | 00000000:24:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        61W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "|   7  NVIDIA A100-SXM4-80GB    On  | 00000000:25:00.0 Off |                    0 |\n",
    "| N/A   31C    P0        63W / 400W |      4MiB / 81920MiB |      0%      Default |\n",
    "|                                   |                      |             Disabled |\n",
    "+-----------------------------------+----------------------+----------------------+\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And so, we will be putting all of them to work...\n",
    "\n",
    "In particular we want all of them to collaborate and optimize over tokens at the same time...\n",
    "\n",
    "And for this we will be using the <a href=\"https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html\">**Distributed Data Parallel (DDP)**</a> from PyTorch...\n",
    "\n",
    "There's also a legacy <a href=\"https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html#torch.nn.DataParallel\">torch.nn.DataParallel</a> which I will **not** recommend using...\n",
    "\n",
    "The **Distributed Data Parallel (DDP)** works in a very simple way.\n",
    "\n",
    "Since we have `8 GPUs` we are going to launch `8 processes` and each process is going to be assigned to a `GPU`. And for for each process the training loop and everything that we've worked on so far is going to look pretty much the same, but now secretly, because they are `8` of them, they are **going to process different parts of the data**. And we are going to add one more part, where once they all calculate their gradients we are going to do an **average of those gradients**.\n",
    "\n",
    "And to do so, we are not going to be launching our script anymore with:\n",
    "```bash\n",
    "python GPT_v2.5.py\n",
    "```\n",
    "Instead we are going to use a special command called <a href=\"https://pytorch.org/docs/stable/elastic/run.html\">torchrun</a>.\n",
    "\n",
    "And when `torchrun` runs a python script will actually make sure to run `8` of them in parallel, and it creates some environmental variables such as <a href=\"https://pytorch.org/docs/stable/elastic/run.html#environment-variables\">`RANK`, `LOCAL_RANK` & `WORLD_SIZE`</a>, where each of these processes can look up and determine which one of the processes it is..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, what are these environmental variables?\n",
    "\n",
    "Well, for our case `WORLD_SIZE` represents the number of `GPUs` we have which is `8` (or the total number of processes running).\n",
    "\n",
    "Each process will run the exact same code at roughly the exact same time, but the only difference between these processes is that they all have a different **DDP Rank** which will be represented as `RANK`.\n",
    "\n",
    "For example (`RANK`):\n",
    "1. `GPU 0` $\\rightarrow$ `RANK 0`\n",
    "1. `GPU 1` $\\rightarrow$ `RANK 1`\n",
    "3. `GPU 2` $\\rightarrow$ `RANK 2`\n",
    "4. `GPU 3` $\\rightarrow$ `RANK 3`\n",
    "\n",
    "And this `RANK` will be used to determine that these processes for example don't run on the same data.\n",
    "\n",
    "Now, `LOCAL_RANK` is something that is used in a **multi-node** setting. Right now we only have a single node with `8 GPUs`, and the `LOCAL_RANK` is the rank of the `GPU` on the single node (for us it can range from `0` to `7`).\n",
    "\n",
    "But, for us, we will be running it on a single box, and the most important things that we care about are `RANK` and `WORLD_SIZE`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And lastly, the thing to understand is, we want to use **Distributed Data Parallel**, but, we also want to revert back to our own code, if we do not have **Distributed Data Parallel** available...\n",
    "\n",
    "And as we start changing our code, we have to imagine and think forward of having `8` different python interpreters running the exact same code, and the only different thing between them is that they have a different **Distributed Data Parallel Rank**..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initializing Python Variables with Environmental Variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let's check if the `DDP` is available to run for us by checking if any of these environmental variables are there after importing `DDP` like this:\n",
    "```python\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP\n",
    "import os\n",
    "\n",
    "ddp = int(os.environ.get('RANK', -1)) != -1\n",
    "```\n",
    "Here we are getting the rank's value as `integer` and if it cannot find a value then the value defaults to `-1`. Then we are converting this to a boolean variable checker by checking if it's **not** `-1`, so that we would be able to use it in an `if` statement...\n",
    "\n",
    "Then we can use this checker to initialize these environmental variables into python variables for later use if the `DDP` is available, like this:\n",
    "```python\n",
    "import os\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP\n",
    "from torch.distributed import init_process_group, destroy_process_group\n",
    "...\n",
    "# Device Auto-Detection\n",
    "ddp = int(os.environ.get('RANK', -1)) != -1\n",
    "if ddp:\n",
    "    assert torch.cuda.is_available(), \"We need CUDA for DDP, because DDP for CPUs makes less sense\"\n",
    "    init_process_group(backend=\"nccl\")\n",
    "    DDPRank = int(os.environ['RANK'])\n",
    "    DDPLocalRank = int(os.environ['LOCAL_RANK'])\n",
    "    DDPWorldSize = int(os.environ['WORLD_SIZE'])\n",
    "    device = f'cuda:{DDPLocalRank}'\n",
    "    torch.cuda.set_device(device=device)\n",
    "    masterProcess = DDPRank == 0\n",
    "else:\n",
    "    DDPRank = 0\n",
    "    DDPLocalRank = 0\n",
    "    DDPWorldSize = 1\n",
    "    masterProcess = True\n",
    "    device = \"cpu\"\n",
    "    if torch.cuda.is_available():\n",
    "        device = \"cuda\"\n",
    "    elif hasattr(torch.backends, \"mps\") and torch.backends.mps.is_available():\n",
    "        device = \"mps\"\n",
    "    print(f\"Using Device: {device}\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, you will see that we are using something like `init_process_group()` which initializes the process group for distributed training or computation and `nccl` is NVIDIA Collective Communications Library developed by NVIDIA that provides highly optimized communication primitives for NVIDIA `GPUs`. And since I have NVIDIA `GPUs`, I will be using this backend.\n",
    "\n",
    "And finally we are also using a `boolean` variable, that I like to use called the `masterProcess` which is arbitrarily process number `0` which will have some additional work to do (logging, checkpointing, etc.) and other processes will do just the forward and backward passes.\n",
    "\n",
    "And when DDP is not available, we revert back to the original process along with the initialization of the new python variables that we are going to use..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, in the case where we will run this using `torchrun` (which we'll see in a bit), there's going to be `8` copies running in parallel with each one having a different rank. And now we have to make sure that everything happens correctly afterwards...\n",
    "\n",
    "And now we have to adjust all the calculations that take into account that there are these `DDPWorldSize` and `DDPRank` variables.\n",
    "\n",
    "And the first thing that we encounter is this snippet:\n",
    "```python\n",
    "gradientAccumulationSteps = totalBatchSize // (Batch * Time)\n",
    "```\n",
    "And because there are `8` processes now, our steps would decrease by a factor of `8` which is the `DDPWorldSize`, and so the snippet now becomes:\n",
    "```python\n",
    "gradientAccumulationSteps = totalBatchSize // (Batch * Time * DDPWorldSize)\n",
    "```\n",
    "For our case, since the `Batch` is `16` and the `Time` is `1024` will have:\n",
    "$$ 524288 / (16 \\times 1024 \\times 8) = 4 \\text{ steps.}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding `Master Process`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And this is where we start coming into issues:\n",
    "```python\n",
    "print(f\"Total Desired Batch Size: {totalBatchSize}\")\n",
    "print(f\"Calculated Gradient Accumulation Steps: {gradientAccumulationSteps}\")\n",
    "```\n",
    "And since we are going to have `8` processes in parallel, they are going to print the same thing again and again.\n",
    "\n",
    "This is where the `masterProcess` variable comes in handy and we say that if it is just the main process then only print this otherwise don't, like this:\n",
    "```python\n",
    "if masterProcess:\n",
    "    print(f\"Total Desired Batch Size: {totalBatchSize}\")\n",
    "    print(f\"Calculated Gradient Accumulation Steps: {gradientAccumulationSteps}\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modifying `DataLoaderLite`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modifying `__init__()`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the next thing we encounter this line:\n",
    "```python\n",
    "trainingLoader = DataLoaderLite(Batch=Batch, Time=Time)\n",
    "```\n",
    "Which means that we have to make our `DataLoaderLite` context aware that we might use this multi-process setting, because in those cases, we don't want the processes to load the exact same data, and want every process to get it's own chunk of data so that they're all working on different parts of the dataset.\n",
    "\n",
    "So let's adjust that...\n",
    "\n",
    "For now our `DataLoaderLite` looks like this:\n",
    "```python\n",
    "# Data-Loader\n",
    "class DataLoaderLite:\n",
    "    def __init__(self, Batch, Time):\n",
    "        self.Batch = Batch\n",
    "        self.Time = Time\n",
    "        with open(\"Datasets/Harry_Potter_Books.txt\", \"r\", encoding=\"UTF-8\") as file:\n",
    "            text = file.read()\n",
    "        encoder = tiktoken.get_encoding('gpt2')\n",
    "        encodedDataTokens = encoder.encode(text)\n",
    "        self.encodedDataTokens = torch.tensor(encodedDataTokens)\n",
    "        print(f\"Loaded {len(self.encodedDataTokens)} Tokens\")\n",
    "        print(f\"1 Epoch = {len(self.encodedDataTokens) // (Batch * Time)} Batches\")\n",
    "\n",
    "        # State\n",
    "        self.currentPosition = 0\n",
    "        \n",
    "    def nextBatch(self):\n",
    "        Batch, Time = self.Batch, self.Time\n",
    "        buffer = self.encodedDataTokens[self.currentPosition : self.currentPosition + Batch*Time + 1]\n",
    "        inputs = buffer[:-1].view(Batch, Time)\n",
    "        labels = buffer[1:].view(Batch, Time)\n",
    "        self.currentPosition += Batch * Time\n",
    "        if self.currentPosition + (Batch * Time + 1) > len(self.encodedDataTokens):\n",
    "            self.currentPosition = 0\n",
    "        return inputs, labels\n",
    "```\n",
    "\n",
    "And one simple way to implement this multi-process setting is by passing `DDPRank` as a `processRank` and `DDPWorldSize` as `numberOfProcesses` into the `DataLoaderLite` constructor to save the data and use it later like this:\n",
    "```python\n",
    "trainingLoader = DataLoaderLite(Batch=Batch, Time=Time, processRank=DDPRank, numberOfProcesses=DDPWorldSize)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's modify the constructor first...\n",
    "```python\n",
    "def __init__(self, Batch, Time, processRank, numberOfProcesses):\n",
    "    self.Batch = Batch\n",
    "    self.Time = Time\n",
    "    self.processRank = processRank\n",
    "    self.numberOfProcesses = numberOfProcesses\n",
    "    with open(\"Datasets/Harry_Potter_Books.txt\", \"r\", encoding=\"UTF-8\") as file:\n",
    "        text = file.read()\n",
    "    encoder = tiktoken.get_encoding('gpt2')\n",
    "    encodedDataTokens = encoder.encode(text)\n",
    "    self.encodedDataTokens = torch.tensor(encodedDataTokens)\n",
    "    print(f\"Loaded {len(self.encodedDataTokens)} Tokens\")\n",
    "    print(f\"1 Epoch = {len(self.encodedDataTokens) // (Batch * Time)} Batches\")\n",
    "\n",
    "    # State\n",
    "    self.currentPosition = self.Batch * self.Time * self.processRank\n",
    "```\n",
    "And you will also see that the starting position (`currentPosition`) is set to `(Batch * Time * processRank)`. And that is because we want to process the dataset in terms of batch chunks.\n",
    "\n",
    "So they will start at for example:\n",
    "1. $\\text{Batch}\\times\\text{Time}\\times\\text{(processRank=0)} = 0$\n",
    "2. $\\text{Batch}\\times\\text{Time}\\times\\text{(processRank=1)} = \\text{Batch}\\times\\text{Time}$\n",
    "3. $\\text{Batch}\\times\\text{Time}\\times\\text{(processRank=2)} = 2\\times\\text{Batch}\\times\\text{Time}$\n",
    "4. $\\text{Batch}\\times\\text{Time}\\times\\text{(processRank=3)} = 3\\times\\text{Batch}\\times\\text{Time}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modifying `nextBatch()`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's modify the advancing code to `nextBatch()` method, which for now looks like this:\n",
    "```python\n",
    "def nextBatch(self):\n",
    "    Batch, Time = self.Batch, self.Time\n",
    "    buffer = self.encodedDataTokens[self.currentPosition : self.currentPosition + Batch*Time + 1]\n",
    "    inputs = buffer[:-1].view(Batch, Time)\n",
    "    labels = buffer[1:].view(Batch, Time)\n",
    "    self.currentPosition += Batch * Time\n",
    "    if self.currentPosition + (Batch * Time + 1) > len(self.encodedDataTokens):\n",
    "        self.currentPosition = 0\n",
    "    return inputs, labels\n",
    "```\n",
    "\n",
    "Now, everything works almost absolutely the same, but when we advance in this line:\n",
    "```python\n",
    "self.currentPosition += Batch * Time\n",
    "```\n",
    "The position has to advance by the entire chunk of the dataset being devoured by all the processes.\n",
    "So the line now becomes:\n",
    "```python\n",
    "self.currentPosition += Batch * Time * self.numberOfProcesses\n",
    "```\n",
    "We also see that the checker for the out of bounds needs to be modified as well such that it could loop through the dataset again:\n",
    "```python\n",
    "if self.currentPosition + (Batch * Time + 1) > len(self.encodedDataTokens):\n",
    "    self.currentPosition = 0\n",
    "```\n",
    "And the modified snippet looks like this:\n",
    "```python\n",
    "if self.currentPosition + (Batch * Time * self.numberOfProcesses + 1) > len(self.encodedDataTokens):\n",
    "    self.currentPosition = self.Batch * self.Time * self.processRank\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So our entire `DataLoaderLite` now looks like this:\n",
    "```python\n",
    "# Data-Loader\n",
    "class DataLoaderLite:\n",
    "    def __init__(self, Batch, Time, processRank, numberOfProcesses):\n",
    "        self.Batch = Batch\n",
    "        self.Time = Time\n",
    "        self.processRank = processRank\n",
    "        self.numberOfProcesses = numberOfProcesses\n",
    "        with open(\"Datasets/Harry_Potter_Books.txt\", \"r\", encoding=\"UTF-8\") as file:\n",
    "            text = file.read()\n",
    "        encoder = tiktoken.get_encoding('gpt2')\n",
    "        encodedDataTokens = encoder.encode(text)\n",
    "        self.encodedDataTokens = torch.tensor(encodedDataTokens)\n",
    "        print(f\"Loaded {len(self.encodedDataTokens)} Tokens\")\n",
    "        print(f\"1 Epoch = {len(self.encodedDataTokens) // (Batch * Time)} Batches\")\n",
    "\n",
    "        # State\n",
    "        self.currentPosition = self.Batch * self.Time * self.processRank\n",
    "        \n",
    "    def nextBatch(self):\n",
    "        Batch, Time = self.Batch, self.Time\n",
    "        buffer = self.encodedDataTokens[self.currentPosition : self.currentPosition + Batch*Time + 1]\n",
    "        inputs = buffer[:-1].view(Batch, Time)\n",
    "        labels = buffer[1:].view(Batch, Time)\n",
    "        self.currentPosition += Batch * Time * self.numberOfProcesses\n",
    "        if self.currentPosition + (Batch * Time * self.numberOfProcesses + 1) > len(self.encodedDataTokens):\n",
    "            self.currentPosition = self.Batch * self.Time * self.processRank\n",
    "        return inputs, labels\n",
    "```\n",
    "And this is the simplest change that I could find, and you'll also notice that if we have `processRank=0 & numberOfProcesses=1`, we'd have the exact same thing that we had implemented before, along with the logic of the newest code..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making `GPTModel` use DDP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now comes the snippet:\n",
    "```python\n",
    "# Constructing Model\n",
    "model = GPTModel(GPTConfiguration(vocabularySize=50304))\n",
    "\n",
    "model.eval()\n",
    "model.to(device=device)\n",
    "model = torch.compile(model)\n",
    "```\n",
    "Now, because we have seeds set for the model construction all the processes are going to construct the same model which is fine...\n",
    "\n",
    "But we also want our model to use DDP if it is available...\n",
    "\n",
    "To do that we have to wrap our model using the `DDP` wrapper, and pass in the correct arguement for `device_ids` parameter. And you can check the <a href=\"https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html\">documentation</a> for this, but it is not very clear. And to my understanding this `device_ids` takes the `DDPLocalRank` as the arguement and **not** the `DDPRank`, wrapped in a list.\n",
    "\n",
    "And this is how it's done:\n",
    "```python\n",
    "# Constructing Model\n",
    "model = GPTModel(GPTConfiguration(vocabularySize=50304))\n",
    "\n",
    "model.eval()\n",
    "model.to(device=device)\n",
    "model = torch.compile(model)\n",
    "if ddp:\n",
    "    model = DDP(model, device_ids=[DDPLocalRank])\n",
    "```\n",
    "And we have to be careful because everything complexifies very very quickly when multiple processes are involved and there's a lot of things to be careful with..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And in particular what `DDP` does for us is, in a forward pass it actually behaves identically (my understanding is, nothing should be changed in the forward pass) but in the backward pass, as we proceed and once it is over in each independent `GPU`, each independent `GPU` has a gradient of all the parameters and it will come something called the <a href=\"https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_reduce\">`torch.distributed.all_reduce()`</a> and basically does an **average** across all the ranks of their gradients, and then it will deposit that average on every single rank such that every single rank will end up with the average on it. And this is the communication and synchronization that `DDP` offers.\n",
    "\n",
    "Now, `DDP` is a little bit more involved than that because as we are doing the backward pass through the layers of the transformer, it actually can dispact communications for the gradient while backward pass is still happening. So there's overlap of the communication of the gradients and their synchronization and the backward pass. And this is just more efficient to do it that way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DDP in Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing DDP Optimization Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's go through the optimization loop and think through what we have to modify to support the multi-process setting...\n",
    "\n",
    "And as I mentioned, as we do the `loss.backward()`, it will do the backward pass and then it will synchronize the gradients...\n",
    "\n",
    "Now, the problem here is, because of our inner gradient accumulation loop, we will be synchronizing it every time, and we don't want that, because we are just depositing gradients making it extremely wasteful of the resources.\n",
    "\n",
    "So, we want to add them up and then on the very last step of the gradient accumulation loop we actually want to use the <a href=\"https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_reduce\">`torch.distributed.all_reduce()`</a> to average up the gradients...\n",
    "\n",
    "And the official way to do this is by using the <a href=\"https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync\">`no_sync()`</a> context manager where within this context, gradients will be accumulated and there will be no communication.\n",
    "\n",
    "And they are telling us to do something like this:\n",
    "```python\n",
    "ddp = torch.nn.parallel.DistributedDataParallel(model, pg)\n",
    "with ddp.no_sync():\n",
    "    for input in inputs:\n",
    "        ddp(input).backward()  # no synchronization, accumulate grads\n",
    "ddp(another_input).backward()  # synchronize grads\n",
    "```\n",
    "And I absolutely hate this way and it looks very ugly and the fact that we have to copy paste our code...\n",
    "\n",
    "So when we go to the <a href=\"https://pytorch.org/docs/stable/_modules/torch/nn/parallel/distributed.html#DistributedDataParallel.no_sync\">source code</a> of `no_sync()` we see that we simply toggle a variable called `require_backward_grad_sync` and is being toggled around and changed like this:\n",
    "```python\n",
    "self.require_backward_grad_sync = False\n",
    "```\n",
    "And I actually like to use that directly...\n",
    "\n",
    "Confusingly, `model.require_backward_grad_sync` is actually used by both the forward and backward pass.\n",
    "\n",
    "And we only want to synchronize and want this variable to be `True` when it is the final iteration of the gradient accumulation loop. So, I just toggle it like this:\n",
    "```python\n",
    "for epoch in range(epochs):\n",
    "    ...\n",
    "    for microStep in range(gradientAccumulationSteps):\n",
    "        inputs, labels = trainingLoader.nextBatch()\n",
    "        inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "        if ddp:\n",
    "            model.require_backward_grad_sync = (microStep == gradientAccumulationSteps - 1)\n",
    "        with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "            logits, loss = model(inputs, labels)\n",
    "        loss = loss / gradientAccumulationSteps\n",
    "        lossAccumulator += loss.detach()\n",
    "        loss.backward()\n",
    "    ...\n",
    "```\n",
    "And this is a naughty thing to do, because if they change the `DDP`, this variable might go away and our code might not work in the future. But for now I believe that this works absolutely fine...\n",
    "\n",
    "And once this gradient accumulation loop is over, and we come out of the loop, every single rank will suddenly have the average of all the gradients that were stored in all the ranks.\n",
    "\n",
    "And now we have to think through whether that's what we want and how it works with the `loss` and `lossAccumulator` as well...\n",
    "\n",
    "So, let's now think through this..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The problem we now get at is that, we have averaged up the gradients but the `lossAccumulator` as not been impacted by the newer code and is outside of the `DDP` container and it is not being averaged and so, when we are printing the `lossAccumulator` we are only going to be printing on the `masterProcess` and it's just going to be printing the losses that it saw on it's process and instead we want it to print the `lossAccumulator` over all the processes and the average of that `lossAccumulator` (because we did average of gradients and we want average of `lossAccumulator` as well)...\n",
    "\n",
    "And to do that we can explicitly call `all_reduce()` on the loss and specify the operation to be an average after we come out of the gradient accumulator loop, like this:\n",
    "```python\n",
    "if ddp:\n",
    "    torch.distributed.all_reduce(lossAccumulator, op=torch.distributed.ReduceOp.AVG)\n",
    "```\n",
    "And so this `lossAccumulator` tensor first exists on all the processes, and when we call `all_reduce()` of average, it creates the average of those numbers and it deposits that average on all the processes. In other words, after this line, all the process will contain the `lossAccumulator` averaged up. So when we print on the process it will have the `lossAccumulator` identical on the other ranks as well...\n",
    "\n",
    "And we don't want multiple same prints, we only want it to print the `lossAccumulator` if it's from the `masterProcess`\n",
    "So our print statement now becomes:\n",
    "```python\n",
    "if masterProcess:\n",
    "    print(f\"Step: {epoch}, Loss: {lossAccumulator.item():.6f}, Learning Rate: {learningRate:.4e},  Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "```\n",
    "\n",
    "But we have to be careful because we are now processing more tokens now, so our `tokensPerSecond` becomes:\n",
    "```python\n",
    "tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time * gradientAccumulationSteps * DDPWorldSize) / (endTime - startTime)\n",
    "```\n",
    "\n",
    "And finally, the only other thing to be careful with is when we are trying to be done with training and when we don't need multiple `GPUs` anymore, we want to free the system from the extra things from happening by using `destroy_process_group()` like this after the optimization loop:\n",
    "```python\n",
    "if ddp:\n",
    "    destroy_process_group()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixing Optimizer and Model Use"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One last thing that we need to fix is this line that we have:\n",
    "```python\n",
    "optimizer = model.configureOptimizers(weightDecayRate=0.1, learningRate=maximumLearningRate, device=device)\n",
    "```\n",
    "\n",
    "And this won't work with the new code because now our model is a `DDP` model and instead we will be using a `rawModel`, where raw model is something that I create right after the `DDP` wrapper that will always contain the base model that we have already created:\n",
    "```python\n",
    "rawModel = model.module if ddp else model\n",
    "```\n",
    "And our optimizer line becomes:\n",
    "```python\n",
    "optimizer = rawModel.configureOptimizers(weightDecayRate=0.1, learningRate=maximumLearningRate, device=device)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And so the entire optimization code looks like this:\n",
    "```python\n",
    "optimizer = rawModel.configureOptimizers(weightDecayRate=0.1, learningRate=maximumLearningRate, device=device)\n",
    "for epoch in range(epochs):\n",
    "    startTime = time.time()\n",
    "    optimizer.zero_grad()\n",
    "    lossAccumulator = 0.0\n",
    "    for microStep in range(gradientAccumulationSteps):\n",
    "        inputs, labels = trainingLoader.nextBatch()\n",
    "        inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "        if ddp:\n",
    "            model.require_backward_grad_sync = (microStep == gradientAccumulationSteps - 1)\n",
    "        with torch.autocast(device_type=device, dtype=torch.bfloat16):\n",
    "            logits, loss = model(inputs, labels)\n",
    "        loss = loss / gradientAccumulationSteps\n",
    "        lossAccumulator += loss.detach()\n",
    "        loss.backward()\n",
    "    if ddp:\n",
    "        torch.distributed.all_reduce(lossAccumulator, op=torch.distributed.ReduceOp.AVG)\n",
    "    normalization = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=1.0)\n",
    "    learningRate = getLearningRate(epoch=epoch)\n",
    "    for parameterGroup in optimizer.param_groups:\n",
    "        parameterGroup['lr'] = learningRate\n",
    "    optimizer.step()\n",
    "    torch.cuda.synchronize()\n",
    "    endTime = time.time()\n",
    "    timeDifference = (endTime - startTime) * 1000\n",
    "    tokensPerSecond = (trainingLoader.Batch * trainingLoader.Time * gradientAccumulationSteps * DDPWorldSize) / (endTime - startTime)\n",
    "    if masterProcess:\n",
    "        print(f\"Step: {epoch}, Loss: {lossAccumulator.item():.6f}, Learning Rate: {learningRate:.4e},  Normalization: {normalization:.4f}, Time Difference: {timeDifference:.2f}ms, Tokens/Second: {tokensPerSecond:.2f}tokens/sec\")\n",
    "if ddp:\n",
    "    destroy_process_group()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fixing `device_type`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `torch.autocast()` function takes an arguement `device_type` to which we tried to stubbornly just pass device hoping it works okay, but PyTorch actually really wants just the type and creates errors in some version of PyTorch.\n",
    "\n",
    "So we will change all the `device` to `device_type` to the following areas:\n",
    "1. `configureOptimizers()` inside the `GPTModel` for fused kernel\n",
    "2. Adding a `device_type` variable into the global scope itself\n",
    "3. `optimizer` object\n",
    "4. `torch.autocast()` function\n",
    "5. wrapping `torch.cuda.synchronize()` into a checker if `cuda` is available"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Changing Actual Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}