7_computational_optimizations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gpt2 import GPT, GPTConfig # our GPT class\n",
    "import time\n",
    "import tiktoken\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = tiktoken.get_encoding('gpt2')\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # dynamic device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(13) # for reproducibility\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(13)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Loader Lite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoaderLite:\n",
    "\n",
    "    def __init__(self, B, T):\n",
    "\n",
    "        self.B, self.T = B, T\n",
    "\n",
    "        with open('data/input.txt', 'r') as file:\n",
    "            text = file.read().replace('\\n', '')\n",
    "        \n",
    "        enc = tiktoken.get_encoding('gpt2')\n",
    "        tokens = enc.encode(text)\n",
    "        self.tokens = torch.tensor(tokens, dtype=torch.long, device=device)\n",
    "\n",
    "        self.current_batch = 0\n",
    "        self.number_of_batches = len(self.tokens) // (B * T)\n",
    "\n",
    "        print(f'Loaded {len(self.tokens)} tokens, {self.number_of_batches} batches of size {B}x{T}')\n",
    "\n",
    "    \n",
    "    def next_batch(self):\n",
    "\n",
    "        B, T = self.B, self.T\n",
    "\n",
    "        buf = self.tokens[self.current_batch * B * T : (self.current_batch + 1) * B * T + 1]\n",
    "        x = buf[:-1].view(B, T)\n",
    "        y = buf[1:].view(B, T)\n",
    "\n",
    "        self.current_batch += 1\n",
    "        if self.current_batch >= self.number_of_batches:\n",
    "            self.current_batch = 0\n",
    "        \n",
    "        return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and Timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 297884 tokens, 72 batches of size 4x1024\n",
      "Step 0 | Loss 10.9441 | 21268.0 ms | 192.59 tok/s\n",
      "Step 1 | Loss 9.9717 | 227.2 ms | 18029.13 tok/s\n",
      "Step 2 | Loss 9.4953 | 227.1 ms | 18036.61 tok/s\n",
      "Step 3 | Loss 9.2852 | 222.1 ms | 18445.96 tok/s\n",
      "Step 4 | Loss 8.9665 | 220.7 ms | 18562.16 tok/s\n",
      "Step 5 | Loss 8.6727 | 221.2 ms | 18516.34 tok/s\n",
      "Step 6 | Loss 8.4885 | 222.2 ms | 18430.19 tok/s\n",
      "Step 7 | Loss 8.3961 | 221.4 ms | 18503.44 tok/s\n",
      "Step 8 | Loss 8.0739 | 228.8 ms | 17903.68 tok/s\n",
      "Step 9 | Loss 7.9872 | 225.1 ms | 18197.09 tok/s\n",
      "Step 10 | Loss 8.3108 | 224.9 ms | 18211.13 tok/s\n",
      "Step 11 | Loss 8.0791 | 226.4 ms | 18094.07 tok/s\n",
      "Step 12 | Loss 7.8551 | 225.1 ms | 18193.56 tok/s\n",
      "Step 13 | Loss 7.6549 | 224.8 ms | 18218.95 tok/s\n",
      "Step 14 | Loss 7.5592 | 225.7 ms | 18144.10 tok/s\n",
      "Step 15 | Loss 7.3684 | 225.6 ms | 18156.09 tok/s\n",
      "Step 16 | Loss 7.3472 | 225.8 ms | 18137.72 tok/s\n",
      "Step 17 | Loss 7.5623 | 226.3 ms | 18098.82 tok/s\n",
      "Step 18 | Loss 7.5903 | 226.7 ms | 18071.88 tok/s\n",
      "Step 19 | Loss 7.4158 | 225.3 ms | 18179.52 tok/s\n",
      "Step 20 | Loss 7.3637 | 225.3 ms | 18178.08 tok/s\n",
      "Step 21 | Loss 7.3745 | 225.3 ms | 18178.01 tok/s\n",
      "Step 22 | Loss 7.4077 | 227.4 ms | 18013.31 tok/s\n",
      "Step 23 | Loss 7.2600 | 226.8 ms | 18059.08 tok/s\n",
      "Step 24 | Loss 7.3014 | 227.9 ms | 17974.82 tok/s\n"
     ]
    }
   ],
   "source": [
    "# torch.set_float32_matmul_precision('high') # This won't help with RTX 2070\n",
    "scaler = torch.cuda.amp.GradScaler() # create a gradient scaler\n",
    "\n",
    "model = GPT(GPTConfig(vocab_size=50304)) # make the vocab size a \"nice number\" (50304 = 2**7 × 3 × 131)\n",
    "model = model.to(device).train()\n",
    "model = torch.compile(model)\n",
    "\n",
    "B, T = 4, 1024\n",
    "data_loader = DataLoaderLite(B, T)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)\n",
    "for i in range(25):\n",
    "\n",
    "    t0 = time.time()\n",
    "    x, y = data_loader.next_batch()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    with torch.autocast(device_type=device.type, dtype=torch.float16):\n",
    "        logits, loss = model(x, y)\n",
    "\n",
    "    scaler.scale(loss).backward()\n",
    "    scaler.step(optimizer)\n",
    "    scaler.update()\n",
    "    \n",
    "    torch.cuda.synchronize() # wait for GPU to finish work\n",
    "    t1 = time.time()\n",
    "    dt = (t1 - t0) * 1000 # time difference in milliseconds\n",
    "    thoughput = (B * T) / (t1 - t0) # tokens per second\n",
    "    print(f\"Step {i} | Loss {loss.item():.4f} | {dt:.1f} ms | {thoughput:.2f} tok/s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Observations:\n",
    "1. Initial throughput around 4000 tok/sec\n",
    "1. Using mixed-precision matmuls with tensor-core via `torch.set_float32_matmul_precision('high')` **did not** increase the throughput. This is because my GPU (RTX 2070) does not support `bfloat16`.\n",
    "1. Using mixed-precision via `torch.autocast` did improve the throughput, as expected. The throughput is now 8000 tok/sec (doubled).\n",
    "1. Unfortunately, my device (RTX 2090) does not have a \"device capability\" above 8 so I cannot use `bfloat16` efficiently and have to use `float16` with gradient scalers (to prevent overflow) instead\n",
    "1. Using `torch.compile` further improved throughput by about 37% (~11000 tok/sec) - why are my performance improvements different from Andrej who got 2.3x after adding `torch.compile`?\n",
    "1. Using Flash Attention further improves the throughput by about 23% (~13500 tok/sec)\n",
    "1. Using \"nice numbers\" by **increasing** the vocab size further improves performance by about 4% (14100 tok/sec)\n",
    "1. Since we are using mixed-precision via `torch.autocast`, we can increase the batch size to 4 making the throughput 18000 tok/sec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7, 5)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.get_device_capability() # only 8 and above can use dtype=torch.bfloat16 :("
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}