{ "cells": [ { "cell_type": "markdown", "id": "c9ec2700", "metadata": {}, "source": [ "# 07 - PyTorch and autograd: the same ideas, far less code\n", "\n", "We built autograd by hand in notebook 04 and a Transformer's forward pass by hand in notebook 06. Working out the backward pass through all of that by hand would be a nightmare. **PyTorch** removes the drudgery with three things:\n", "\n", "1. **tensors**, which are like NumPy's grids of numbers but can also live on the GPU and quietly remember how they were computed, so they can be differentiated,\n", "2. **autograd**, which is exactly notebook 04 generalized from single numbers to whole tensors: every operation records itself, and `.backward()` runs the chain rule for you,\n", "3. **`nn` modules and optimizers**, which are ready-made layers and update rules, so you no longer hand-write them.\n", "\n", "Everything you have learned still applies; you simply stop writing the gradient code yourself." ] }, { "cell_type": "code", "execution_count": null, "id": "colab-setup-07", "metadata": {}, "outputs": [], "source": [ "# Colab setup -- fetch the files this notebook needs.\n", "# (Does nothing when run locally in the course folder.)\n", "import os, urllib.request\n", "BASE = (\"https://raw.githubusercontent.com/waze\"\n", " \"emlabs/llm-book-code/main/\")\n", "for f in ['data/input.txt']:\n", " if not os.path.exists(f):\n", " d = os.path.dirname(f)\n", " if d: os.makedirs(d, exist_ok=True)\n", " urllib.request.urlretrieve(BASE + f, f)\n", " print(\"downloaded\", f)\n" ] }, { "cell_type": "markdown", "id": "colab-setup-md-07", "metadata": {}, "source": [ "
Line by line: what each line does\n\n
" ] }, { "cell_type": "code", "execution_count": 1, "id": "e644b564", "metadata": { "execution": { "iopub.execute_input": "2026-07-01T09:18:24.464961Z", "iopub.status.busy": "2026-07-01T09:18:24.464743Z", "iopub.status.idle": "2026-07-01T09:18:24.910938Z", "shell.execute_reply": "2026-07-01T09:18:24.910494Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "torch 2.12.0 | device: mps\n" ] } ], "source": [ "import torch\n", "import torch.nn as nn\n", "from torch.nn import functional as F\n", "torch.manual_seed(1337)\n", "\n", "device = \"mps\" if torch.backends.mps.is_available() else (\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "print(\"torch\", torch.__version__, \"| device:\", device)" ] }, { "cell_type": "markdown", "id": "5e8708dc", "metadata": {}, "source": [ "
Line by line: what each line does\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "2b463f74", "metadata": {}, "source": [ "## Autograd: notebook 04, now built in\n", "\n", "Remember `L = (a*b + b**2) * tanh(c)` from notebook 04? We worked out its gradients by hand, and then our small engine reproduced them. Watch PyTorch produce the exact same numbers with a single call to `.backward()`.\n", "\n", "The one new ingredient is `requires_grad=True`. It marks an input as something we want gradients for, so PyTorch quietly records every operation performed on it. Calling `L.backward()` then walks that record backward and fills in `a.grad`, `b.grad`, and `c.grad`. It is the same chain rule you built, giving the same answers (`-1.386`, `-1.848`, `2.359`), only handled for you." ] }, { "cell_type": "code", "execution_count": 2, "id": "282a2245", "metadata": { "execution": { "iopub.execute_input": "2026-07-01T09:18:24.912135Z", "iopub.status.busy": "2026-07-01T09:18:24.912045Z", "iopub.status.idle": "2026-07-01T09:18:24.917196Z", "shell.execute_reply": "2026-07-01T09:18:24.916881Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "a.grad = -1.38635 b.grad = -1.84847 c.grad = 2.35934\n", "(notebook 04 by hand: a=-1.38635 b=-1.84847 c=2.35934)\n" ] } ], "source": [ "a = torch.tensor(2.0, requires_grad=True)\n", "b = torch.tensor(-3.0, requires_grad=True)\n", "c = torch.tensor(0.5, requires_grad=True)\n", "\n", "L = (a*b + b**2) * torch.tanh(c)\n", "L.backward() # autograd walks the graph for us\n", "print(\"a.grad = %.5f b.grad = %.5f c.grad = %.5f\" % (a.grad, b.grad, c.grad))\n", "print(\"(notebook 04 by hand: a=-1.38635 b=-1.84847 c=2.35934)\")" ] }, { "cell_type": "markdown", "id": "657726cf", "metadata": {}, "source": [ "
Line by line: what each line does\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "m7c1a", "metadata": {}, "source": [ "### The graph PyTorch built for you\n", "\n", "In notebook 04 you built the expression graph by hand and walked it in reverse. PyTorch did the same thing the moment you wrote the formula: every operation left a `grad_fn` pointing back at its inputs. Here is that recorded graph for `L` -- the same nodes as notebook 04's table, but you never wrote them." ] }, { "cell_type": "code", "execution_count": 3, "id": "m7c1b", "metadata": { "execution": { "iopub.execute_input": "2026-07-01T09:18:24.918233Z", "iopub.status.busy": "2026-07-01T09:18:24.918178Z", "iopub.status.idle": "2026-07-01T09:18:24.920414Z", "shell.execute_reply": "2026-07-01T09:18:24.920072Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "MulBackward0\n", " AddBackward0\n", " MulBackward0\n", " leaf a\n", " leaf b\n", " PowBackward0\n", " leaf b\n", " TanhBackward0\n", " leaf c\n", "\n", "the same graph as notebook 04: two *, one +, one **2, one tanh, leaves a/b/c\n" ] } ], "source": [ "# PyTorch recorded every operation in L.grad_fn. Walk it, exactly like notebook 04.\n", "leaf = {id(a): \"a\", id(b): \"b\", id(c): \"c\"}\n", "L = (a*b + b**2) * torch.tanh(c) # rebuild so the graph is intact to walk\n", "\n", "def show(fn, depth=0):\n", " if fn is None:\n", " return\n", " name = type(fn).__name__\n", " if name == \"AccumulateGrad\": # a leaf input (a, b, or c)\n", " print(\" \" * depth + f\"leaf {leaf.get(id(fn.variable), '?')}\")\n", " else:\n", " print(\" \" * depth + name) # e.g. MulBackward0, TanhBackward0\n", " for nxt, _ in fn.next_functions:\n", " show(nxt, depth + 1)\n", "\n", "show(L.grad_fn)\n", "print(\"\\nthe same graph as notebook 04: two *, one +, one **2, one tanh, leaves a/b/c\")" ] }, { "cell_type": "markdown", "id": "m7c1c", "metadata": {}, "source": [ "
Line by line: what each line does\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "e4c2a476", "metadata": {}, "source": [ "## The training loop, in PyTorch form\n", "\n", "Two new helpers keep the loop tidy. An `nn.Module` is a container that bundles up a model's learnable numbers, and an **optimizer** is the part that actually takes the downhill step. We will use **AdamW**, which you can think of as a smarter version of plain gradient descent: it carries a little momentum and automatically tunes the step size for each number, so it trains faster and needs less hand-holding. (Plain gradient descent -- step every number downhill by a fixed learning rate -- is the simplest optimizer, `torch.optim.SGD`. The name is short for *stochastic gradient descent*: \"stochastic\" because each step's gradient is measured on a fresh random mini-batch, not the whole dataset, exactly as `get_batch` draws one here; notebook 04 unpacks it.) The loop is always the same four lines:\n", "\n", "```\n", "logits, loss = model(x, y) # forward\n", "optimizer.zero_grad() # clear old gradients\n", "loss.backward() # autograd: compute new gradients\n", "optimizer.step() # nudge every parameter downhill\n", "```\n", "\n", "The line that trips people up is `zero_grad()`. Recall from notebook 04 that gradients add up. PyTorch does the same, so if you do not reset them to zero each step, this step's gradients pile on top of the last step's and the update is meaningless. So the order is: clear, compute, step.\n", "\n", "Below we rebuild the bigram from notebook 03 in this style. `nn.Embedding` is just the lookup table `W` from back then, now packaged as a layer, and notice that we write no gradient code at all. It lands at the same ceiling of about 2.45 that we computed in notebook 03, which is exactly the point: that ceiling belongs to the model's one-character memory, not to how cleverly you train it." ] }, { "cell_type": "code", "execution_count": 4, "id": "9acac002", "metadata": { "execution": { "iopub.execute_input": "2026-07-01T09:18:24.921324Z", "iopub.status.busy": "2026-07-01T09:18:24.921273Z", "iopub.status.idle": "2026-07-01T09:18:25.008592Z", "shell.execute_reply": "2026-07-01T09:18:25.008149Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "vocab: 65 | one batch: (32, 8)\n" ] } ], "source": [ "from pathlib import Path\n", "# --- data + tokenizer (notebook 02) ---\n", "text = Path(\"data/input.txt\").read_text()\n", "chars = sorted(set(text)); vocab_size = len(chars)\n", "stoi = {c:i for i,c in enumerate(chars)}; itos = {i:c for c,i in stoi.items()}\n", "encode = lambda s: [stoi[c] for c in s]; decode = lambda l: \"\".join(itos[i] for i in l)\n", "data = torch.tensor(encode(text), dtype=torch.long)\n", "n = int(0.9*len(data)); train_data, val_data = data[:n], data[n:]\n", "\n", "block_size, batch_size = 8, 32\n", "def get_batch(split):\n", " d = train_data if split==\"train\" else val_data\n", " ix = torch.randint(len(d)-block_size, (batch_size,))\n", " x = torch.stack([d[i:i+block_size] for i in ix])\n", " y = torch.stack([d[i+1:i+1+block_size] for i in ix])\n", " return x.to(device), y.to(device)\n", "print(\"vocab:\", vocab_size, \"| one batch:\", tuple(get_batch(\"train\")[0].shape))" ] }, { "cell_type": "markdown", "id": "0003843f", "metadata": {}, "source": [ "
Line by line: what each line does\n", "\n", "
" ] }, { "cell_type": "code", "execution_count": 5, "id": "7e360c39", "metadata": { "execution": { "iopub.execute_input": "2026-07-01T09:18:25.009680Z", "iopub.status.busy": "2026-07-01T09:18:25.009611Z", "iopub.status.idle": "2026-07-01T09:18:29.667269Z", "shell.execute_reply": "2026-07-01T09:18:29.666842Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "final loss (avg over 200 batches): 2.446\n", "right at the ~2.45 bigram ceiling we computed in notebook 03 -- same model, zero gradient code\n" ] } ], "source": [ "class BigramLM(nn.Module):\n", " def __init__(self, vocab_size):\n", " super().__init__()\n", " # each token id directly indexes a row of next-token logits\n", " self.token_emb = nn.Embedding(vocab_size, vocab_size)\n", " def forward(self, idx, targets=None):\n", " logits = self.token_emb(idx) # (B,T,vocab)\n", " loss = None\n", " if targets is not None:\n", " B,T,V = logits.shape\n", " loss = F.cross_entropy(logits.view(B*T, V), targets.view(B*T))\n", " return logits, loss\n", "\n", "model = BigramLM(vocab_size).to(device)\n", "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)\n", "\n", "for step in range(4000):\n", " x, y = get_batch(\"train\")\n", " _, loss = model(x, y)\n", " optimizer.zero_grad(set_to_none=True)\n", " loss.backward()\n", " optimizer.step()\n", "\n", "# a single batch is a noisy reading, so evaluate properly over many batches\n", "with torch.no_grad():\n", " eval_loss = torch.stack([model(*get_batch(\"train\"))[1] for _ in range(200)]).mean()\n", "print(f\"final loss (avg over 200 batches): {eval_loss.item():.3f}\")\n", "print(\"right at the ~2.45 bigram ceiling we computed in notebook 03 -- same model, zero gradient code\")" ] }, { "cell_type": "markdown", "id": "2bc5b54c", "metadata": {}, "source": [ "
Line by line: what each line does\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "d980b1ed", "metadata": {}, "source": [ "Same model, same result, but the gradient that took a careful derivation and `np.add.at` in notebook 03 is now handled entirely by `loss.backward()`. That is the whole point of a framework: you specify the forward computation, and the backward pass comes for free.\n", "\n", "## Attention in PyTorch\n", "\n", "And here is notebook 05's scaled, causal self-attention, written in PyTorch. Compare it line by line with the NumPy version; it is the same arithmetic. `nn.Linear` is just the learned projection (our `W_q`, `W_k`, `W_v`) packaged as a layer. `register_buffer` stores the causal mask so that it travels to the GPU along with the model but is not treated as a trainable number. Because autograd is watching every step, this module is ready to train as it stands." ] }, { "cell_type": "code", "execution_count": 6, "id": "b1b37399", "metadata": { "execution": { "iopub.execute_input": "2026-07-01T09:18:29.668428Z", "iopub.status.busy": "2026-07-01T09:18:29.668322Z", "iopub.status.idle": "2026-07-01T09:18:29.674698Z", "shell.execute_reply": "2026-07-01T09:18:29.674406Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "attention output: (4, 8, 16)\n", "causal (upper triangle zero): True\n" ] } ], "source": [ "class Head(nn.Module):\n", " def __init__(self, n_embd, head_size, block_size):\n", " super().__init__()\n", " self.key = nn.Linear(n_embd, head_size, bias=False)\n", " self.query = nn.Linear(n_embd, head_size, bias=False)\n", " self.value = nn.Linear(n_embd, head_size, bias=False)\n", " self.register_buffer(\"tril\", torch.tril(torch.ones(block_size, block_size)))\n", " def forward(self, x):\n", " B,T,C = x.shape\n", " k, q = self.key(x), self.query(x)\n", " wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 # scaled affinities\n", " wei = wei.masked_fill(self.tril[:T,:T]==0, float('-inf')) # causal mask\n", " wei = F.softmax(wei, dim=-1)\n", " return wei @ self.value(x)\n", "\n", "x = torch.randn(4, 8, 32)\n", "head = Head(n_embd=32, head_size=16, block_size=8)\n", "out = head(x)\n", "print(\"attention output:\", tuple(out.shape))\n", "# verify causality on the attention weights\n", "with torch.no_grad():\n", " k,q = head.key(x), head.query(x)\n", " wei = (q @ k.transpose(-2,-1) * k.shape[-1]**-0.5).masked_fill(head.tril[:8,:8]==0, float('-inf')).softmax(-1)\n", "print(\"causal (upper triangle zero):\", torch.allclose(torch.triu(wei[0], 1), torch.zeros(8,8)))" ] }, { "cell_type": "markdown", "id": "337af1d9", "metadata": {}, "source": [ "
Line by line: what each line does\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "m7c2a", "metadata": {}, "source": [ "### What `nn.Module` is holding\n", "\n", "An `nn.Module` quietly tracks two kinds of tensors: **parameters**, the weights the optimizer trains, and **buffers**, fixed helpers that ride along to the GPU but are never trained. Ask the `Head` what it holds -- the query, key, and value matrices are parameters; the causal `tril` mask is a buffer." ] }, { "cell_type": "code", "execution_count": 7, "id": "m7c2b", "metadata": { "execution": { "iopub.execute_input": "2026-07-01T09:18:29.675744Z", "iopub.status.busy": "2026-07-01T09:18:29.675685Z", "iopub.status.idle": "2026-07-01T09:18:29.677547Z", "shell.execute_reply": "2026-07-01T09:18:29.677257Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "parameters (the optimizer trains these):\n", " key.weight (16, 32) 512 numbers\n", " query.weight (16, 32) 512 numbers\n", " value.weight (16, 32) 512 numbers\n", "buffers (fixed, moved to the device but never trained):\n", " tril (8, 8)\n", "\n", "total trainable numbers in this head: 1536\n", "model.parameters() is exactly what you hand the optimizer -- no manual list to keep.\n" ] } ], "source": [ "# nn.Module auto-registers its weights (parameters) and its fixed helpers (buffers).\n", "print(\"parameters (the optimizer trains these):\")\n", "for name, p in head.named_parameters():\n", " print(f\" {name:14s} {tuple(p.shape)} {p.numel()} numbers\")\n", "print(\"buffers (fixed, moved to the device but never trained):\")\n", "for name, buf in head.named_buffers():\n", " print(f\" {name:14s} {tuple(buf.shape)}\")\n", "print(f\"\\ntotal trainable numbers in this head: {sum(p.numel() for p in head.parameters())}\")\n", "print(\"model.parameters() is exactly what you hand the optimizer -- no manual list to keep.\")" ] }, { "cell_type": "markdown", "id": "m7c2c", "metadata": {}, "source": [ "
Line by line: what each line does\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "0deedccf", "metadata": {}, "source": [ "## Recap\n", "\n", "PyTorch is, in short, tensors plus autograd plus the `nn` building blocks. The bigram and the attention head that you worked through by hand now take only a few lines and train themselves, and you know exactly what every line is doing, because you built it the hard way first. You now have every component of a GPT.\n", "\n", "Next: notebook 08 assembles the full **GPT** and inspects it before training." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (.venv)", "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.12.12" } }, "nbformat": 4, "nbformat_minor": 5 }