{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM4-L3: Assemble a GPT model in PyTorch\n",
    "\n",
    "You will build a decoder-only GPT model from its parts: layer normalisation, a GELU feed-forward block, a Transformer block with residual connections, and the full model with\n",
    "embeddings and an output head. Then you will count its parameters (and match the 163,009,536 of GPT-2 small without allocating its memory), run a forward pass, and generate token IDs\n",
    "from the **untrained** model.\n",
    "\n",
    "**What this notebook does not do:** it does not train anything. The weights are random, so the generated IDs are meaningless; the point is that the machinery is correct.\n",
    "Training starts in Module 5.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, a few minutes. The GPT-2 sized model is created on PyTorch's `meta` device, which records shapes without allocating memory.\n",
    "\n",
    "Cells marked `TODO` raise `NotImplementedError` until you write the code. Each is followed by a self-check (`assert`). Theory: Lessons LM4-T1 to LM4-T5."
   ]
  },
  {
   "cell_type": "code",
   "id": "c002",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(123)\n",
    "print(\"torch\", torch.__version__)\n",
    "\n",
    "GPT2_SMALL = dict(vocab_size=50257, context_length=1024, emb_dim=768, n_heads=12, n_layers=12, drop_rate=0.1, qkv_bias=False)\n",
    "TINY = dict(vocab_size=100, context_length=32, emb_dim=32, n_heads=4, n_layers=2, drop_rate=0.0, qkv_bias=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. Given: causal multi-head attention\n",
    "\n",
    "This is the module you wrote in LM3-L3, provided so that you can concentrate on the new parts."
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class CausalMultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_in, d_out, n_heads, context_length, dropout, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        assert d_out % n_heads == 0\n",
    "        self.n_heads, self.head_dim, self.d_out = n_heads, d_out // n_heads, d_out\n",
    "        self.W_q = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_k = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_v = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.out_proj = nn.Linear(d_out, d_out)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.register_buffer(\"mask\", torch.triu(torch.ones(context_length, context_length), diagonal=1).bool())\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, n, _ = x.shape\n",
    "        H, hd = self.n_heads, self.head_dim\n",
    "        q = self.W_q(x).view(B, n, H, hd).transpose(1, 2)\n",
    "        k = self.W_k(x).view(B, n, H, hd).transpose(1, 2)\n",
    "        v = self.W_v(x).view(B, n, H, hd).transpose(1, 2)\n",
    "        scores = (q @ k.transpose(2, 3)) / hd ** 0.5\n",
    "        scores = scores.masked_fill(self.mask[:n, :n], float(\"-inf\"))\n",
    "        w = self.dropout(torch.softmax(scores, dim=-1))\n",
    "        return self.out_proj((w @ v).transpose(1, 2).contiguous().view(B, n, self.d_out))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. TODO: layer normalisation\n",
    "\n",
    "Normalise over the **last** dimension: `y = scale * (x - mean) / sqrt(var + eps) + shift`, with the **biased** variance (`unbiased=False`) and `eps = 1e-5`.\n",
    "`scale` starts as ones and `shift` as zeros; both are `nn.Parameter`s of size `emb_dim`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, emb_dim):\n",
    "        super().__init__()\n",
    "        self.eps = 1e-5\n",
    "        self.scale = nn.Parameter(torch.ones(emb_dim))\n",
    "        self.shift = nn.Parameter(torch.zeros(emb_dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(dim=-1, keepdim=True)\n",
    "        var = x.var(dim=-1, keepdim=True, unbiased=False)\n",
    "        return self.scale * (x - mean) / torch.sqrt(var + self.eps) + self.shift"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "### Self-check 1"
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ln = LayerNorm(8)\n",
    "x = torch.randn(4, 5, 8) * 3 + 7\n",
    "y = ln(x)\n",
    "assert y.shape == x.shape\n",
    "assert torch.allclose(y.mean(-1), torch.zeros(4, 5), atol=1e-5), \"each vector should have mean 0 after normalisation\"\n",
    "assert torch.allclose(y.var(-1, unbiased=False), torch.ones(4, 5), atol=1e-3), \"each vector should have (biased) variance 1\"\n",
    "assert torch.allclose(y, F.layer_norm(x, (8,), eps=1e-5), atol=1e-5), \"should equal torch.nn.functional.layer_norm\"\n",
    "print(\"self-check 1 passed: layer norm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "## 3. TODO: GELU and the feed-forward block\n",
    "\n",
    "GELU (tanh approximation used by GPT-2): `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x**3)))`.\n",
    "The feed-forward block applies, to each token separately: `Linear(d, 4d)`, GELU, `Linear(4d, d)`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class GELU(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x ** 3)))\n",
    "\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, emb_dim):\n",
    "        super().__init__()\n",
    "        self.layers = nn.Sequential(nn.Linear(emb_dim, 4 * emb_dim), GELU(), nn.Linear(4 * emb_dim, emb_dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "### Self-check 2"
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "g = GELU()\n",
    "t = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0])\n",
    "assert torch.allclose(g(t), F.gelu(t, approximate=\"tanh\"), atol=1e-6), \"GELU should match F.gelu(approximate='tanh')\"\n",
    "assert abs(g(torch.tensor(1.0)).item() - 0.8412) < 1e-3\n",
    "ff = FeedForward(16)\n",
    "assert ff(torch.randn(2, 3, 16)).shape == (2, 3, 16)\n",
    "assert sum(p.numel() for p in ff.parameters()) == 16 * 64 + 64 + 64 * 16 + 16, \"expected d*4d + 4d + 4d*d + d parameters\"\n",
    "print(\"self-check 2 passed: GELU and feed-forward\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "## 4. TODO: the Transformer block\n",
    "\n",
    "The pre-LN block used by GPT-2, with two residual connections:\n",
    "\n",
    "```\n",
    "x = x + dropout( attention( norm1(x) ) )\n",
    "x = x + dropout( feedforward( norm2(x) ) )\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        d = cfg[\"emb_dim\"]\n",
    "        self.att = CausalMultiHeadAttention(d, d, cfg[\"n_heads\"], cfg[\"context_length\"], cfg[\"drop_rate\"], cfg[\"qkv_bias\"])\n",
    "        self.ff = FeedForward(d)\n",
    "        self.norm1, self.norm2 = LayerNorm(d), LayerNorm(d)\n",
    "        self.drop = nn.Dropout(cfg[\"drop_rate\"])\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.drop(self.att(self.norm1(x)))\n",
    "        x = x + self.drop(self.ff(self.norm2(x)))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "### Self-check 3: shape, residual path and causality"
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "blk = TransformerBlock(TINY).eval()\n",
    "h = torch.randn(2, 10, 32)\n",
    "assert blk(h).shape == h.shape, \"a block keeps the shape (B, n, d)\"\n",
    "h2 = h.clone(); h2[:, -1] += 5.0\n",
    "assert torch.allclose(blk(h)[:, :-1], blk(h2)[:, :-1], atol=1e-5), \"the block leaks information from a later token\"\n",
    "zero_blk = TransformerBlock(TINY).eval()\n",
    "for p in list(zero_blk.att.out_proj.parameters()) + list(zero_blk.ff.layers[2].parameters()):\n",
    "    nn.init.zeros_(p)\n",
    "assert torch.allclose(zero_blk(h), h, atol=1e-6), \"with both sublayers outputting zero, a residual block must return its input\"\n",
    "print(\"self-check 3 passed: Transformer block\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "## 5. TODO: the GPT model\n",
    "\n",
    "Token embedding `(vocab, d)`, position embedding `(context, d)`, dropout, `n_layers` blocks, a final layer norm and an output head `Linear(d, vocab, bias=False)`.\n",
    "`forward(idx)` takes integer token IDs of shape `(B, n)` and returns logits of shape `(B, n, vocab)`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class GPTModel(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"])\n",
    "        self.pos_emb = nn.Embedding(cfg[\"context_length\"], cfg[\"emb_dim\"])\n",
    "        self.drop_emb = nn.Dropout(cfg[\"drop_rate\"])\n",
    "        self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])])\n",
    "        self.final_norm = LayerNorm(cfg[\"emb_dim\"])\n",
    "        self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False)\n",
    "\n",
    "    def forward(self, idx):\n",
    "        B, n = idx.shape\n",
    "        x = self.tok_emb(idx) + self.pos_emb(torch.arange(n, device=idx.device))\n",
    "        x = self.drop_emb(x)\n",
    "        x = self.trf_blocks(x)\n",
    "        x = self.final_norm(x)\n",
    "        return self.out_head(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "### Self-check 4: forward pass and parameter counts"
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "model = GPTModel(TINY).eval()\n",
    "idx = torch.randint(0, TINY[\"vocab_size\"], (3, 12))\n",
    "logits = model(idx)\n",
    "assert logits.shape == (3, 12, TINY[\"vocab_size\"]), f\"logits should be (B, n, vocab), got {tuple(logits.shape)}\"\n",
    "\n",
    "def formula(cfg, tied=False):\n",
    "    d, V, L, n = cfg[\"emb_dim\"], cfg[\"vocab_size\"], cfg[\"n_layers\"], cfg[\"context_length\"]\n",
    "    attn = 3 * d * d + (3 * d if cfg[\"qkv_bias\"] else 0) + d * d + d\n",
    "    ffn = d * 4 * d + 4 * d + 4 * d * d + d\n",
    "    return V * d + n * d + L * (attn + ffn + 4 * d) + 2 * d + (0 if tied else V * d)\n",
    "\n",
    "n_tiny = sum(p.numel() for p in model.parameters())\n",
    "assert n_tiny == formula(TINY), f\"tiny model has {n_tiny} parameters, the formula says {formula(TINY)}\"\n",
    "\n",
    "with torch.device(\"meta\"):                 # shapes only, no memory is allocated\n",
    "    big = GPTModel(GPT2_SMALL)\n",
    "n_big = sum(p.numel() for p in big.parameters())\n",
    "assert n_big == 163_009_536, f\"GPT-2 small with a separate output head has 163,009,536 parameters, got {n_big:,}\"\n",
    "n_tied = n_big - big.out_head.weight.numel()\n",
    "assert n_tied == 124_412_160, \"sharing the output head with the token embedding gives 124,412,160 (the 124M model)\"\n",
    "print(f\"self-check 4 passed: tiny model {n_tiny:,} parameters; GPT-2 small {n_big:,} ({n_tied:,} with weight sharing)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## 6. TODO: generate token IDs (untrained model)\n",
    "\n",
    "Greedy generation: repeat `max_new_tokens` times: crop the sequence to the last `context_size` tokens, run the model, take the logits of the **last** position, pick the arg-max token and append it."
   ]
  },
  {
   "cell_type": "code",
   "id": "c022",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def generate_simple(model, idx, max_new_tokens, context_size):\n",
    "    for _ in range(max_new_tokens):\n",
    "        cond = idx[:, -context_size:]\n",
    "        with torch.no_grad():\n",
    "            logits = model(cond)\n",
    "        nxt = logits[:, -1, :].argmax(dim=-1, keepdim=True)\n",
    "        idx = torch.cat([idx, nxt], dim=1)\n",
    "    return idx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "### Self-check 5"
   ]
  },
  {
   "cell_type": "code",
   "id": "c024",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "start = torch.tensor([[1, 2, 3]])\n",
    "out = generate_simple(model, start, max_new_tokens=40, context_size=TINY[\"context_length\"])\n",
    "assert out.shape == (1, 43), f\"expected shape (1, 43), got {tuple(out.shape)}\"\n",
    "assert torch.equal(out[:, :3], start), \"the prompt must be kept at the start\"\n",
    "assert torch.equal(out, generate_simple(model, start, 40, TINY[\"context_length\"])), \"greedy decoding of an eval-mode model must be deterministic\"\n",
    "assert out.max().item() < TINY[\"vocab_size\"]\n",
    "print(\"generated IDs (meaningless: the weights are random):\", out[0, 3:13].tolist(), \"...\")\n",
    "print(\"self-check 5 passed: generation loop, including cropping to the context length (43 > 32)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c025",
   "metadata": {},
   "source": [
    "## 7. Look at an untrained model\n",
    "\n",
    "The cross-entropy of a model that knows nothing is about `ln(vocab)`. This is a useful sanity check for any new model before training. No assertion: read the numbers."
   ]
  },
  {
   "cell_type": "code",
   "id": "c026",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "targets = torch.randint(0, TINY[\"vocab_size\"], (3, 12))\n",
    "loss = F.cross_entropy(model(idx).flatten(0, 1), targets.flatten())\n",
    "print(f\"initial loss {loss.item():.3f}   ln(vocab) = {math.log(TINY['vocab_size']):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c027",
   "metadata": {},
   "source": [
    "## 8. Reflect\n",
    "\n",
    "1. The logits tensor for GPT-2 small with a batch of 8 sequences of 1,024 tokens has how many entries, and how many megabytes in float32?\n",
    "2. Where are most of the parameters of the model, and what changes if you share the output head with the token embedding?\n",
    "3. Why does the block return its input exactly when both sublayers output zero, and why does that matter for training deep models?\n",
    "4. Why is the text generated here meaningless, and what must happen before it becomes meaningful?\n",
    "\n",
    "**Done when** all five self-checks print \"passed\"."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}