{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9a5936bd-af17-4a7e-a4d2-e910411708ea",
   "metadata": {},
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"http://mng.bz/orYv\">Build a Large Language Model From Scratch</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af53bcb1-ff9d-49c7-a0bc-5b8d32ff975b",
   "metadata": {},
   "source": [
    "## Appendix D: Adding Bells and Whistles to the Training Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f58c142-9434-49af-b33a-356b80a45b86",
   "metadata": {},
   "source": [
    "- In this appendix, we add a few more advanced features to the training function, which are used in typical pretraining and finetuning; finetuning is covered in chapters 6 and 7\n",
    "- The next three sections below discuss learning rate warmup, cosine decay, and gradient clipping\n",
    "- The final section adds these techniques to the training function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "744def4f-c03f-42ee-97bb-5d7d5b89b723",
   "metadata": {},
   "source": [
    "- We start by initializing a model reusing the code from chapter 5:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8755bd5e-bc06-4e6e-9e63-c7c82b816cbe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch version: 2.9.0\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "import torch\n",
    "\n",
    "print(\"torch version:\", version(\"torch\"))\n",
    "\n",
    "\n",
    "from previous_chapters import GPTModel\n",
    "# If the `previous_chapters.py` file is not available locally,\n",
    "# you can import it from the `llms-from-scratch` PyPI package.\n",
    "# For details, see: https://github.com/rasbt/LLMs-from-scratch/tree/main/pkg\n",
    "# E.g.,\n",
    "# from llms_from_scratch.ch04 import GPTModel\n",
    "\n",
    "GPT_CONFIG_124M = {\n",
    "    \"vocab_size\": 50257,   # Vocabulary size\n",
    "    \"context_length\": 256, # Shortened context length (orig: 1024)\n",
    "    \"emb_dim\": 768,        # Embedding dimension\n",
    "    \"n_heads\": 12,         # Number of attention heads\n",
    "    \"n_layers\": 12,        # Number of layers\n",
    "    \"drop_rate\": 0.1,      # Dropout rate\n",
    "    \"qkv_bias\": False      # Query-key-value bias\n",
    "}\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    # Use PyTorch 2.9 or newer for stable mps results\n",
    "    major, minor = map(int, torch.__version__.split(\".\")[:2])\n",
    "    if (major, minor) >= (2, 9):\n",
    "        device = torch.device(\"mps\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "print(\"Device:\", device)\n",
    "\n",
    "torch.manual_seed(123)\n",
    "model = GPTModel(GPT_CONFIG_124M)\n",
    "model.eval();  # Disable dropout during inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51574e57-a098-412c-83e8-66dafa5a0b99",
   "metadata": {},
   "source": [
    "- Next, using the same code we used in chapter 5, we initialize the data loaders:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "386ca110-2bb4-42f1-bd54-8836df80acaa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nimport os\\nimport urllib.request\\n\\nif not os.path.exists(file_path):\\n    with urllib.request.urlopen(url) as response:\\n        text_data = response.read().decode(\\'utf-8\\')\\n    with open(file_path, \"w\", encoding=\"utf-8\") as file:\\n        file.write(text_data)\\nelse:\\n    with open(file_path, \"r\", encoding=\"utf-8\") as file:\\n        text_data = file.read()\\n'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import requests\n",
    "\n",
    "file_path = \"the-verdict.txt\"\n",
    "url = \"https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch02/01_main-chapter-code/the-verdict.txt\"\n",
    "\n",
    "if not os.path.exists(file_path):\n",
    "    response = requests.get(url, timeout=30)\n",
    "    response.raise_for_status()\n",
    "    text_data = response.text\n",
    "    with open(file_path, \"w\", encoding=\"utf-8\") as file:\n",
    "        file.write(text_data)\n",
    "else:\n",
    "    with open(file_path, \"r\", encoding=\"utf-8\") as file:\n",
    "        text_data = file.read()\n",
    "\n",
    "# The book originally used the following code below\n",
    "# However, urllib uses older protocol settings that\n",
    "# can cause problems for some readers using a VPN.\n",
    "# The `requests` version above is more robust\n",
    "# in that regard.\n",
    "\n",
    "\"\"\"\n",
    "import os\n",
    "import urllib.request\n",
    "\n",
    "if not os.path.exists(file_path):\n",
    "    with urllib.request.urlopen(url) as response:\n",
    "        text_data = response.read().decode('utf-8')\n",
    "    with open(file_path, \"w\", encoding=\"utf-8\") as file:\n",
    "        file.write(text_data)\n",
    "else:\n",
    "    with open(file_path, \"r\", encoding=\"utf-8\") as file:\n",
    "        text_data = file.read()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ae96992b-536a-4684-a924-658b9ffb7e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from previous_chapters import create_dataloader_v1\n",
    "# Alternatively:\n",
    "# from llms_from_scratch.ch02 import create_dataloader_v1\n",
    "\n",
    "\n",
    "# Train/validation ratio\n",
    "train_ratio = 0.90\n",
    "split_idx = int(train_ratio * len(text_data))\n",
    "\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "train_loader = create_dataloader_v1(\n",
    "    text_data[:split_idx],\n",
    "    batch_size=2,\n",
    "    max_length=GPT_CONFIG_124M[\"context_length\"],\n",
    "    stride=GPT_CONFIG_124M[\"context_length\"],\n",
    "    drop_last=True,\n",
    "    shuffle=True,\n",
    "    num_workers=0\n",
    ")\n",
    "\n",
    "val_loader = create_dataloader_v1(\n",
    "    text_data[split_idx:],\n",
    "    batch_size=2,\n",
    "    max_length=GPT_CONFIG_124M[\"context_length\"],\n",
    "    stride=GPT_CONFIG_124M[\"context_length\"],\n",
    "    drop_last=False,\n",
    "    shuffle=False,\n",
    "    num_workers=0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "939c08d8-257a-41c6-b842-019f7897ac74",
   "metadata": {},
   "source": [
    "## D.1 Learning rate warmup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fafcd30-ddf7-4a9f-bcf4-b13c052b3133",
   "metadata": {},
   "source": [
    "- When training complex models like LLMs, implementing learning rate warmup can help stabilize the training\n",
    "- In learning rate warmup, we gradually increase the learning rate from a very low value (`initial_lr`) to a user-specified maximum (`peak_lr`)\n",
    "- This way, the model will start the training with small weight updates, which helps decrease the risk of large destabilizing updates during the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2bb4790b-b8b6-4e9e-adf4-704a04b31ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 15\n",
    "initial_lr = 0.0001\n",
    "peak_lr = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bf3a8da-abc4-4b80-a5d8-f1cc1c7cc5f3",
   "metadata": {},
   "source": [
    "- Typically, the number of warmup steps is between 0.1% to 20% of the total number of steps\n",
    "- We can compute the increment as the difference between the `peak_lr` and `initial_lr` divided by the number of warmup steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5f6d083f-1b25-4c23-b46d-ef7783446690",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "27\n"
     ]
    }
   ],
   "source": [
    "total_steps = len(train_loader) * n_epochs\n",
    "warmup_steps = int(0.2 * total_steps) # 20% warmup\n",
    "print(warmup_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b6bbdc8-0104-459e-a7ed-b08be8578709",
   "metadata": {},
   "source": [
    "- Note that the print book accidentally includes a leftover code line, `warmup_steps = 20`, which is not used and can be safely ignored"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e075f80e-a398-4809-be1d-8019e1d31c90",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_increment = (peak_lr - initial_lr) / warmup_steps\n",
    "\n",
    "global_step = -1\n",
    "track_lrs = []\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), weight_decay=0.1)\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    for input_batch, target_batch in train_loader:\n",
    "        optimizer.zero_grad()\n",
    "        global_step += 1\n",
    "    \n",
    "        if global_step < warmup_steps:\n",
    "            lr = initial_lr + global_step * lr_increment\n",
    "        else:\n",
    "            lr = peak_lr\n",
    "        \n",
    "        # Apply the calculated learning rate to the optimizer\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group[\"lr\"] = lr\n",
    "        track_lrs.append(optimizer.param_groups[0][\"lr\"])\n",
    "    \n",
    "        # Calculate loss and update weights\n",
    "        # ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cb6da121-eeed-4023-bdd8-3666c594b4ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18312 characters]",
      "text/plain": [
       "<Figure size 500x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(5, 3))\n",
    "plt.ylabel(\"Learning rate\")\n",
    "plt.xlabel(\"Step\")\n",
    "total_training_steps = len(train_loader) * n_epochs\n",
    "plt.plot(range(total_training_steps), track_lrs)\n",
    "plt.tight_layout(); plt.savefig(\"1.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b3996b6-3f7a-420a-8584-c5760249f3d8",
   "metadata": {},
   "source": [
    "## D.2 Cosine decay"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5216214-de79-40cf-a733-b1049a73023c",
   "metadata": {},
   "source": [
    "- Another popular technique for training complex deep neural networks is cosine decay, which also adjusts the learning rate across training epochs\n",
    "- In cosine decay, the learning rate follows a cosine curve, decreasing from its initial value to near zero following a half-cosine cycle\n",
    "- This gradual reduction is designed to slow the pace of learning as the model begins to improve its weights; it reduces the risk of overshooting minima as the training progresses,  which is crucial for stabilizing the training in its later stages\n",
    "- Cosine decay is often preferred over linear decay for its smoother transition in learning rate adjustments, but linear decay is also used in practice (for example, [OLMo: Accelerating the Science of Language Models](https://arxiv.org/abs/2402.00838))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4e8d2068-a057-4abf-b478-f02cc37191f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "min_lr = 0.1 * initial_lr\n",
    "track_lrs = []\n",
    "\n",
    "lr_increment = (peak_lr - initial_lr) / warmup_steps\n",
    "global_step = -1\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    for input_batch, target_batch in train_loader:\n",
    "        optimizer.zero_grad()\n",
    "        global_step += 1\n",
    "    \n",
    "        # Adjust the learning rate based on the current phase (warmup or cosine annealing)\n",
    "        if global_step < warmup_steps:\n",
    "            # Linear warmup\n",
    "            lr = initial_lr + global_step * lr_increment  \n",
    "        else:\n",
    "            # Cosine annealing after warmup\n",
    "            progress = ((global_step - warmup_steps) / \n",
    "                        (total_training_steps - warmup_steps))\n",
    "            lr = min_lr + (peak_lr - min_lr) * 0.5 * (\n",
    "                1 + math.cos(math.pi * progress))\n",
    "        \n",
    "        # Apply the calculated learning rate to the optimizer\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group[\"lr\"] = lr\n",
    "        track_lrs.append(optimizer.param_groups[0][\"lr\"])\n",
    "    \n",
    "        # Calculate loss and update weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0e779e33-8a44-4984-bb23-be0603dc4158",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25372 characters]",
      "text/plain": [
       "<Figure size 500x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(5, 3))\n",
    "plt.ylabel(\"Learning rate\")\n",
    "plt.xlabel(\"Step\")\n",
    "plt.plot(range(total_training_steps), track_lrs)\n",
    "plt.tight_layout(); plt.savefig(\"2.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7512808-b48d-4146-86a1-5931b1e3aec1",
   "metadata": {},
   "source": [
    "## D.3 Gradient clipping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0a74f76-8d2b-4974-a03c-d645445cdc21",
   "metadata": {},
   "source": [
    "- Gradient clipping is yet another technique used to stabilize the training when training LLMs\n",
    "- By setting a threshold, gradients exceeding this limit are scaled down to a maximum magnitude to ensure that the updates to the model's parameters during backpropagation remain within a manageable range\n",
    "- For instance, using the `max_norm=1.0` setting in PyTorch's `clip_grad_norm_` method means that the norm of the gradients is clipped such that their maximum norm does not exceed 1.0\n",
    "- the \"norm\" refers to a measure of the gradient vector's length (or magnitude) in the parameter space of the model\n",
    "- Specifically, it's the L2 norm, also known as the Euclidean norm\n",
    "- Mathematically, for a vector $\\mathbf{v}$ with components $\\mathbf{v} = [v_1, v_2, \\ldots, v_n]$, the L2 norm is defined as:\n",
    "$$\n",
    "\\| \\mathbf{v} \\|_2 = \\sqrt{v_1^2 + v_2^2 + \\ldots + v_n^2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d44838a6-4322-47b2-a935-c00d3a88355f",
   "metadata": {},
   "source": [
    "- The L2 norm is calculated similarly for matrices.\n",
    "- Let's assume our gradient matrix is:\n",
    "$$\n",
    "G = \\begin{bmatrix}\n",
    "1 & 2 \\\\\n",
    "2 & 4\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "- And we want to clip these gradients with a `max_norm` of 1.\n",
    "\n",
    "- First, we calculate the L2 norm of these gradients:\n",
    "$$\n",
    "\\|G\\|_2 = \\sqrt{1^2 + 2^2 + 2^2 + 4^2} = \\sqrt{25} = 5\n",
    "$$\n",
    "\n",
    "- Since $\\|G\\|_2 = 5$ is greater than our `max_norm` of 1, we need to scale down the gradients so that their norm is exactly 1. The scaling factor is calculated as $\\frac{max\\_norm}{\\|G\\|_2} = \\frac{1}{5}$.\n",
    "\n",
    "- Therefore, the scaled gradient matrix $G'$ will be as follows:\n",
    "$$\n",
    "G' = \\frac{1}{5} \\times G = \\begin{bmatrix}\n",
    "\\frac{1}{5} & \\frac{2}{5} \\\\\n",
    "\\frac{2}{5} & \\frac{4}{5}\n",
    "\\end{bmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeb0c3c1-2cff-46f5-8127-24412184428c",
   "metadata": {},
   "source": [
    "- Let's see this in action\n",
    "- First, we initialize a new model and calculate the loss for a training batch like we would do in the regular training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e199e1ff-58c4-413a-855e-5edbe9292649",
   "metadata": {},
   "outputs": [],
   "source": [
    "from previous_chapters import calc_loss_batch\n",
    "# Alternatively:\n",
    "# from llms_from_scratch.ch05 import calc_loss_batch\n",
    "\n",
    "\n",
    "torch.manual_seed(123)\n",
    "model = GPTModel(GPT_CONFIG_124M)\n",
    "model.to(device)\n",
    "\n",
    "loss = calc_loss_batch(input_batch, target_batch, model, device)\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76b60f3a-15ec-4846-838d-fdef3df99899",
   "metadata": {},
   "source": [
    "- If we call `.backward()`, PyTorch will calculate the gradients and store them in a `.grad` attribute for each weight (parameter) matrix\n",
    "- Let's define a utility function to calculate the highest gradient based on all model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e70729a3-24d1-411d-a002-2529cd3a8a9e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0446, device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "def find_highest_gradient(model):\n",
    "    max_grad = None\n",
    "    for param in model.parameters():\n",
    "        if param.grad is not None:\n",
    "            grad_values = param.grad.data.flatten()\n",
    "            max_grad_param = grad_values.max()\n",
    "            if max_grad is None or max_grad_param > max_grad:\n",
    "                max_grad = max_grad_param\n",
    "    return max_grad\n",
    "\n",
    "print(find_highest_gradient(model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "734f30e6-6b24-4d4b-ae91-e9a4b871113f",
   "metadata": {},
   "source": [
    "- Applying gradient clipping, we can see that the largest gradient is now substantially smaller:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fa81ef8b-4280-400f-a93e-5210f3e62ff0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0201, device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "print(find_highest_gradient(model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b62c2af0-dac3-4742-be4b-4292c6753099",
   "metadata": {},
   "source": [
    "## D.4 The modified training function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76715332-94ec-4185-922a-75cb420819d5",
   "metadata": {},
   "source": [
    "- Now let's add the three concepts covered above (learning rate warmup, cosine decay, and gradient clipping) to the `train_model_simple` function covered in chapter 5 to create the more sophisticated `train_model` function below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "46eb9c84-a293-4016-a523-7ad726e171e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from previous_chapters import evaluate_model, generate_and_print_sample\n",
    "# Alternatively:\n",
    "# from llms_from_scratch.ch05 import evaluate_model, generate_and_print_samplee\n",
    "\n",
    "\n",
    "ORIG_BOOK_VERSION = False\n",
    "\n",
    "\n",
    "def train_model(model, train_loader, val_loader, optimizer, device,\n",
    "                n_epochs, eval_freq, eval_iter, start_context, tokenizer,\n",
    "                warmup_steps, initial_lr=3e-05, min_lr=1e-6):\n",
    "\n",
    "    train_losses, val_losses, track_tokens_seen, track_lrs = [], [], [], []\n",
    "    tokens_seen, global_step = 0, -1\n",
    "\n",
    "    # Retrieve the maximum learning rate from the optimizer\n",
    "    peak_lr = optimizer.param_groups[0][\"lr\"]\n",
    "\n",
    "    # Calculate the total number of iterations in the training process\n",
    "    total_training_steps = len(train_loader) * n_epochs\n",
    "\n",
    "    # Calculate the learning rate increment during the warmup phase\n",
    "    lr_increment = (peak_lr - initial_lr) / warmup_steps\n",
    "\n",
    "    for epoch in range(n_epochs):\n",
    "        model.train()\n",
    "        for input_batch, target_batch in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "            global_step += 1\n",
    "\n",
    "            # Adjust the learning rate based on the current phase (warmup or cosine annealing)\n",
    "            if global_step < warmup_steps:\n",
    "                # Linear warmup\n",
    "                lr = initial_lr + global_step * lr_increment  \n",
    "            else:\n",
    "                # Cosine annealing after warmup\n",
    "                progress = ((global_step - warmup_steps) / \n",
    "                            (total_training_steps - warmup_steps))\n",
    "                lr = min_lr + (peak_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))\n",
    "\n",
    "            # Apply the calculated learning rate to the optimizer\n",
    "            for param_group in optimizer.param_groups:\n",
    "                param_group[\"lr\"] = lr\n",
    "            track_lrs.append(lr)  # Store the current learning rate\n",
    "\n",
    "            # Calculate and backpropagate the loss\n",
    "            loss = calc_loss_batch(input_batch, target_batch, model, device)\n",
    "            loss.backward()\n",
    "\n",
    "            # Apply gradient clipping after the warmup phase to avoid exploding gradients\n",
    "            if ORIG_BOOK_VERSION:\n",
    "                if global_step > warmup_steps:\n",
    "                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  \n",
    "            else:\n",
    "                if global_step >= warmup_steps:  # the book originally used global_step > warmup_steps, which led to a skipped clipping step after warmup\n",
    "                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "                \n",
    "            optimizer.step()\n",
    "            tokens_seen += input_batch.numel()\n",
    "\n",
    "            # Periodically evaluate the model on the training and validation sets\n",
    "            if global_step % eval_freq == 0:\n",
    "                train_loss, val_loss = evaluate_model(\n",
    "                    model, train_loader, val_loader,\n",
    "                    device, eval_iter\n",
    "                )\n",
    "                train_losses.append(train_loss)\n",
    "                val_losses.append(val_loss)\n",
    "                track_tokens_seen.append(tokens_seen)\n",
    "                # Print the current losses\n",
    "                print(f\"Ep {epoch+1} (Iter {global_step:06d}): \"\n",
    "                      f\"Train loss {train_loss:.3f}, \"\n",
    "                      f\"Val loss {val_loss:.3f}\"\n",
    "                )\n",
    "\n",
    "        # Generate and print a sample from the model to monitor progress\n",
    "        generate_and_print_sample(\n",
    "            model, tokenizer, device, start_context\n",
    "        )\n",
    "\n",
    "    return train_losses, val_losses, track_tokens_seen, track_lrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "55fcd247-ba9d-4b93-a757-0f7ce04fee41",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 1 (Iter 000000): Train loss 10.969, Val loss 10.938\n",
      "Ep 1 (Iter 000005): Train loss 9.337, Val loss 9.460\n",
      "Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Ep 2 (Iter 000010): Train loss 7.823, Val loss 8.186\n",
      "Ep 2 (Iter 000015): Train loss 6.343, Val loss 6.890\n",
      "Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Ep 3 (Iter 000020): Train loss 6.083, Val loss 6.594\n",
      "Ep 3 (Iter 000025): Train loss 5.677, Val loss 6.786\n",
      "Every effort moves you, the of the of the of the of the, the of the of the, the, the, the, the, the of the, the, the, the, the, the, the, the, the, the, the, the\n",
      "Ep 4 (Iter 000030): Train loss 5.837, Val loss 6.822\n",
      "Ep 4 (Iter 000035): Train loss 6.136, Val loss 7.203\n",
      "Every effort moves youisisis a                                              \n",
      "Ep 5 (Iter 000040): Train loss 5.767, Val loss 7.003\n",
      "Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Ep 6 (Iter 000045): Train loss 5.957, Val loss 6.874\n",
      "Ep 6 (Iter 000050): Train loss 5.633, Val loss 6.839\n",
      "Every effort moves you, the, the, the,, the, the, the, I, the, the,, I, I, the, the, the, the, the, the, the, the, the, the, the, the, the\n",
      "Ep 7 (Iter 000055): Train loss 5.699, Val loss 6.878\n",
      "Ep 7 (Iter 000060): Train loss 5.066, Val loss 6.622\n",
      "Every effort moves you, I had, I had a, I had a a a a of the a, I had, I had been, I I had a of the of the of the of the of the of the of the of the of the of the of\n",
      "Ep 8 (Iter 000065): Train loss 4.593, Val loss 6.548\n",
      "Ep 8 (Iter 000070): Train loss 4.278, Val loss 6.424\n",
      "Every effort moves you.                                                 \n",
      "Ep 9 (Iter 000075): Train loss 4.107, Val loss 6.279\n",
      "Ep 9 (Iter 000080): Train loss 3.146, Val loss 6.215\n",
      "Every effort moves you, in the picture\"I felt--I had been \" it.                                    \n",
      "Ep 10 (Iter 000085): Train loss 2.827, Val loss 6.177\n",
      "Every effort moves you know, and in a little of the Riv--I had a little of the Riv, and I was, and, and in the Riv of the of the fact of the fact--as of the, and, and--because.   \n",
      "Ep 11 (Iter 000090): Train loss 2.579, Val loss 6.135\n",
      "Ep 11 (Iter 000095): Train loss 2.277, Val loss 6.122\n",
      "Every effort moves you know the inevitable of the picture to have been--I turned. Gisburn's an--and it--I  \" to see a smile behind his pictures that he had married her. I had the fact, and it. Gisburn\n",
      "Ep 12 (Iter 000100): Train loss 2.231, Val loss 6.141\n",
      "Ep 12 (Iter 000105): Train loss 1.552, Val loss 6.148\n",
      "Every effort moves you know the inevitable garlanded frame. The mere, and Mrs.            \" to the picture; and I had been the donkey.  \"I had been the bull--and I had\n",
      "Ep 13 (Iter 000110): Train loss 1.800, Val loss 6.157\n",
      "Ep 13 (Iter 000115): Train loss 1.448, Val loss 6.145\n",
      "Every effort moves you know.\" \"I looked up his pictures--I had a little of a flash that he was, and to me to have to see a smile behind his pictures--I had married her--the, I had been the bull--that I had\n",
      "Ep 14 (Iter 000120): Train loss 1.557, Val loss 6.165\n",
      "Ep 14 (Iter 000125): Train loss 1.190, Val loss 6.164\n",
      "Every effort moves you know.\" \"I looked up his pictures--I had a single one in the house.\"        \"I must, and I looked up at my elbow and as I had been the man of the hour. \n",
      "Ep 15 (Iter 000130): Train loss 1.312, Val loss 6.158\n",
      "Every effort moves you know.\" \"I looked up his pictures--I had a single one in the house.\"        \"--it was his close grayish beard--as if he had the donkey. \"There were days when I\n"
     ]
    }
   ],
   "source": [
    "import tiktoken\n",
    "\n",
    "# Note:\n",
    "# Uncomment the following code to calculate the execution time\n",
    "# import time\n",
    "# start_time = time.time()\n",
    "\n",
    "torch.manual_seed(123)\n",
    "model = GPTModel(GPT_CONFIG_124M)\n",
    "model.to(device)\n",
    "\n",
    "peak_lr = 0.001  # this was originally set to 5e-4 in the book by mistake\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=peak_lr, weight_decay=0.1)  # the book accidentally omitted the lr assignment\n",
    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "\n",
    "n_epochs = 15\n",
    "train_losses, val_losses, tokens_seen, lrs = train_model(\n",
    "    model, train_loader, val_loader, optimizer, device, n_epochs=n_epochs,\n",
    "    eval_freq=5, eval_iter=1, start_context=\"Every effort moves you\",\n",
    "    tokenizer=tokenizer, warmup_steps=warmup_steps, \n",
    "    initial_lr=1e-5, min_lr=1e-5\n",
    ")\n",
    "\n",
    "# Note:\n",
    "# Uncomment the following code to show the execution time\n",
    "# end_time = time.time()\n",
    "# execution_time_minutes = (end_time - start_time) / 60\n",
    "# print(f\"Training completed in {execution_time_minutes:.2f} minutes.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "827e8d5e-0872-4b90-98ac-200c80ee2d53",
   "metadata": {},
   "source": [
    "- Looking at the results above, we can see that the model starts out generating incomprehensible strings of words, whereas, towards the end, it's able to produce grammatically more or less correct sentences\n",
    "- If we were to check a few passages it writes towards the end, we would find that they are contained in the training set verbatim -- it simply memorizes the training data\n",
    "- Note that the overfitting here occurs because we have a very, very small training set, and we iterate over it so many times\n",
    "  - The LLM training here primarily serves educational purposes; we mainly want to see that the model can learn to produce coherent text\n",
    "  - Instead of spending weeks or months on training this model on vast amounts of expensive hardware, we load the pretrained weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9decec45-4fdf-4ff6-85a7-1806613f8af7",
   "metadata": {},
   "source": [
    "- A quick check that the learning rate behaves as intended"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d8ebb8d2-8308-4a83-a2a6-730c3bf84452",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 26952 characters]",
      "text/plain": [
       "<Figure size 500x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(5, 3))\n",
    "plt.plot(range(len(lrs)), lrs)\n",
    "plt.ylabel(\"Learning rate\")\n",
    "plt.xlabel(\"Steps\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2f85b01-859b-4454-a3a3-c7ef593735a6",
   "metadata": {},
   "source": [
    "- And a quick look at the loss curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "445d8155-6eae-4b50-a381-d0820ebc27cc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28128 characters]",
      "text/plain": [
       "<Figure size 500x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from previous_chapters import plot_losses\n",
    "# Alternatively:\n",
    "# from llms_from_scratch.ch05 import plot_losses\n",
    "\n",
    "\n",
    "epochs_tensor = torch.linspace(1, n_epochs, len(train_losses))\n",
    "plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)\n",
    "plt.tight_layout(); plt.savefig(\"3.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c16fa614-67e1-4254-8b7e-c3e2f690c29c",
   "metadata": {},
   "source": [
    "- Note that the model is overfitting here because the dataset is kept very small for educational purposes (so that the code can be executed on a laptop computer)\n",
    "- For a longer pretraining run on a much larger dataset, see [../../ch05/03_bonus_pretraining_on_gutenberg](../../ch05/03_bonus_pretraining_on_gutenberg)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
