{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM3-L3: Implement causal multi-head attention in PyTorch\n",
    "\n",
    "In this notebook you write the core operation of a GPT-style model: **causal self-attention**, first for one head and then for several heads in parallel.\n",
    "You will check your code three ways: tensor shapes, a leakage test (a later token must not change an earlier output), and a comparison with PyTorch's\n",
    "own `scaled_dot_product_attention`.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, a few minutes. No downloads; all data are random tensors.\n",
    "\n",
    "**How it works:** cells marked `TODO` raise `NotImplementedError` until you write the code. Each is followed by a self-check (`assert`).\n",
    "Run the cells from the top. Theory: Lessons LM3-T3 to LM3-T6 on this site."
   ]
  },
  {
   "cell_type": "code",
   "id": "c002",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(0)\n",
    "print(\"torch\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. Inputs: a batch of embedded tokens\n",
    "\n",
    "A batch of `B = 2` sequences, each with `T = 6` tokens, each token an embedding vector of width `d = 8`. The tensor `x` has shape `(B, T, d)`.\n",
    "(In a real model it would be the sum of a token embedding and a position embedding; here it is random.)"
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "B, T, d = 2, 6, 8\n",
    "x = torch.randn(B, T, d)\n",
    "print(x.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. TODO: single-head causal attention\n",
    "\n",
    "Given `x` of shape `(B, T, d_in)` and three projection matrices `Wq, Wk, Wv` of shape `(d_in, d_k)`:\n",
    "\n",
    "1. `q = x @ Wq`, `k = x @ Wk`, `v = x @ Wv`, each of shape `(B, T, d_k)`.\n",
    "2. `scores = q @ k^T / sqrt(d_k)`, shape `(B, T, T)`. Transpose the last two dimensions of `k`.\n",
    "3. Set the scores of future positions (column index greater than row index) to `-inf` **before** the softmax.\n",
    "4. `weights = softmax(scores, dim=-1)`, shape `(B, T, T)`; then `out = weights @ v`, shape `(B, T, d_k)`.\n",
    "\n",
    "Return `(out, weights)`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def causal_attention(x, Wq, Wk, Wv):\n",
    "    raise NotImplementedError(\"TODO: implement single-head causal attention (steps 1 to 4 above)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "### Self-check 1: shapes, normalisation, mask, and agreement with PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "d_k = 4\n",
    "Wq, Wk, Wv = (torch.randn(d, d_k) / d ** 0.5 for _ in range(3))\n",
    "out, w = causal_attention(x, Wq, Wk, Wv)\n",
    "assert out.shape == (B, T, d_k), f\"out should have shape {(B, T, d_k)}, got {tuple(out.shape)}\"\n",
    "assert w.shape == (B, T, T), f\"weights should have shape {(B, T, T)}, got {tuple(w.shape)}\"\n",
    "assert torch.allclose(w.sum(-1), torch.ones(B, T), atol=1e-5), \"each row of weights must sum to 1: mask BEFORE the softmax\"\n",
    "assert torch.all(torch.triu(w, diagonal=1) == 0), \"weights above the diagonal (future tokens) must be exactly 0\"\n",
    "assert torch.allclose(w[:, 0, 0], torch.ones(B)), \"the first token can only attend to itself\"\n",
    "ref = F.scaled_dot_product_attention(x @ Wq, x @ Wk, x @ Wv, is_causal=True)\n",
    "assert torch.allclose(out, ref, atol=1e-5), \"output differs from F.scaled_dot_product_attention(..., is_causal=True)\"\n",
    "print(\"self-check 1 passed: single-head causal attention matches PyTorch\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "## 3. A leakage test\n",
    "\n",
    "The most important property of a causal model: **changing token `j` must not change the outputs at positions before `j`.**\n",
    "The helper below changes the last token of every sequence and returns the largest change in the earlier outputs.\n",
    "Use it to test your function, and see that it detects a broken implementation."
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def leak(fn, x):\n",
    "    \"\"\"Largest change in outputs at positions 0..T-2 when the LAST token is replaced. Zero for a correct causal layer.\"\"\"\n",
    "    x2 = x.clone()\n",
    "    x2[:, -1, :] = torch.randn_like(x2[:, -1, :]) * 3\n",
    "    a, b = fn(x), fn(x2)\n",
    "    return (a[:, :-1] - b[:, :-1]).abs().max().item()\n",
    "\n",
    "def broken_attention(x, Wq, Wk, Wv):\n",
    "    \"\"\"BROKEN on purpose: no mask at all.\"\"\"\n",
    "    q, k, v = x @ Wq, x @ Wk, x @ Wv\n",
    "    w = torch.softmax(q @ k.transpose(-2, -1) / k.shape[-1] ** 0.5, dim=-1)\n",
    "    return w @ v, w\n",
    "\n",
    "bad = leak(lambda z: broken_attention(z, Wq, Wk, Wv)[0], x)\n",
    "good = leak(lambda z: causal_attention(z, Wq, Wk, Wv)[0], x)\n",
    "print(f\"broken: {bad:.4f}   yours: {good:.2e}\")\n",
    "assert bad > 1e-3, \"the leakage test should detect the unmasked layer\"\n",
    "assert good < 1e-6, \"your causal attention leaks: a later token changed an earlier output\"\n",
    "print(\"self-check 2 passed: no information flows from the future\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "## 4. TODO: multi-head attention\n",
    "\n",
    "One head can attend by one notion of relevance. Several heads, each with a smaller width, run in parallel and their results are concatenated.\n",
    "With `n_heads = H` and output width `d_out`, each head has width `head_dim = d_out / H`.\n",
    "\n",
    "The constructor is given. In `forward`, with input `x` of shape `(B, T, d_in)`:\n",
    "\n",
    "1. Project: `q, k, v = self.W_q(x), self.W_k(x), self.W_v(x)`, each `(B, T, d_out)`.\n",
    "2. Split heads: `view(B, T, H, head_dim)`, then `transpose(1, 2)` to get `(B, H, T, head_dim)`.\n",
    "3. Scores `(B, H, T, T)` = `q @ k.transpose(2, 3) / sqrt(head_dim)`; mask the future with `-inf` using `self.mask[:T, :T]`; softmax over the last dimension; dropout on the weights.\n",
    "4. Mix: `weights @ v` has shape `(B, H, T, head_dim)`. Transpose back to `(B, T, H, head_dim)` and merge heads with `contiguous().view(B, T, d_out)`.\n",
    "5. Apply `self.out_proj` and return the result, shape `(B, T, d_out)`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class CausalMultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_in, d_out, n_heads, context_length, dropout=0.0):\n",
    "        super().__init__()\n",
    "        assert d_out % n_heads == 0, \"d_out must be divisible by n_heads\"\n",
    "        self.d_out, self.n_heads, self.head_dim = d_out, n_heads, d_out // n_heads\n",
    "        self.W_q = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.W_k = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.W_v = nn.Linear(d_in, d_out, bias=False)\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",
    "        raise NotImplementedError(\"TODO: implement the multi-head forward pass (steps 1 to 5 above)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "### Self-check 3: shapes, parameters, leakage, and agreement with PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(1)\n",
    "mha = CausalMultiHeadAttention(d_in=d, d_out=8, n_heads=2, context_length=16)\n",
    "y = mha(x)\n",
    "assert y.shape == (B, T, 8), f\"output should be {(B, T, 8)}, got {tuple(y.shape)}\"\n",
    "assert leak(lambda z: mha(z), x) < 1e-6, \"your multi-head attention leaks information from future tokens\"\n",
    "\n",
    "# the same weights through PyTorch's fused function must give the same result\n",
    "H, hd = mha.n_heads, mha.head_dim\n",
    "q, k, v = (proj(x).view(B, T, H, hd).transpose(1, 2) for proj in (mha.W_q, mha.W_k, mha.W_v))\n",
    "ref = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(B, T, 8)\n",
    "ref = mha.out_proj(ref)\n",
    "assert torch.allclose(y, ref, atol=1e-5), \"output differs from the reference built with F.scaled_dot_product_attention\"\n",
    "\n",
    "# the number of heads splits the width; it does not add parameters\n",
    "n_params = lambda m: sum(p.numel() for p in m.parameters())\n",
    "assert n_params(CausalMultiHeadAttention(d, 8, 1, 16)) == n_params(CausalMultiHeadAttention(d, 8, 4, 16)), \"heads must not change the parameter count\"\n",
    "assert n_params(mha) == 3 * d * 8 + 8 * 8 + 8, \"expected 3 * d_in * d_out (Q, K, V) + d_out * d_out + d_out (output projection with bias) parameters\"\n",
    "print(\"self-check 3 passed: multi-head attention is correct;\", n_params(mha), \"parameters\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "## 5. TODO: track the tensor shapes\n",
    "\n",
    "Shape errors are the most common bug in attention code. For a batch of `B` sequences of `T` tokens, with `H` heads of width `hd` (so `d_out = H * hd`),\n",
    "fill in the shapes of the intermediate tensors of the multi-head forward pass as tuples."
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def expected_shapes(B, T, H, hd):\n",
    "    raise NotImplementedError(\"TODO: return the shape tuples for each stage\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "### Self-check 4: your answer against the real tensors"
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "shapes = expected_shapes(2, 6, 2, 4)\n",
    "Bn, Tn, Hn, hdn = 2, 6, 2, 4\n",
    "real_q = mha.W_q(x)\n",
    "assert shapes[\"projected_q\"] == tuple(real_q.shape), f\"projected_q: expected {tuple(real_q.shape)}\"\n",
    "assert shapes[\"q_split_into_heads\"] == tuple(real_q.view(Bn, Tn, Hn, hdn).shape)\n",
    "assert shapes[\"q_per_head\"] == tuple(real_q.view(Bn, Tn, Hn, hdn).transpose(1, 2).shape)\n",
    "qh = real_q.view(Bn, Tn, Hn, hdn).transpose(1, 2)\n",
    "assert shapes[\"scores\"] == tuple((qh @ qh.transpose(2, 3)).shape), \"scores should be (B, H, T, T)\"\n",
    "assert shapes[\"weights\"] == shapes[\"scores\"], \"the softmax does not change the shape\"\n",
    "assert shapes[\"context_per_head\"] == tuple(((qh @ qh.transpose(2, 3)) @ qh).shape)\n",
    "assert shapes[\"context_merged\"] == tuple(y.shape)\n",
    "print(\"self-check 4 passed: every shape matches\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "## 6. Experiment: what do the heads cost?\n",
    "\n",
    "The attention scores have shape `(B, H, T, T)`: the memory for them grows with the **square** of the sequence length.\n",
    "Run the cell and read the numbers; they need no checks."
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "for n_tokens in (64, 128, 256, 512, 1024):\n",
    "    n_scores = 1 * 12 * n_tokens * n_tokens   # one sequence, 12 heads\n",
    "    print(f\"T = {n_tokens:5d}: {n_scores:>12,} attention scores per layer = {n_scores * 4 / 1e6:8.1f} MB in float32\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## 7. Reflect\n",
    "\n",
    "Write short answers in your own notes, then compare with Lesson LM3-T5 and LM3-T6.\n",
    "\n",
    "1. Why must the mask be applied **before** the softmax?\n",
    "2. You double the context length. By what factor does the memory for the attention scores grow? What about the memory for `q`, `k` and `v`?\n",
    "3. Why does `contiguous()` appear before `view` when merging heads?\n",
    "4. With `d_out = 768` and 12 heads, what is `head_dim`? How many parameters do the Q, K and V projections have together?\n",
    "\n",
    "**Done when** all four self-checks print \"passed\"."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}