{ "cells": [ { "cell_type": "markdown", "id": "ef460528", "metadata": {}, "source": [ "# Let's build GPT — from scratch\n", "\n", "Tiny GPT on *tiny Shakespeare* — **concept -> code -> Your turn** each step.\n", "\n", "**Red thread:** every section has **-> Training:** plus a **goal stack**\n", "so you always know *why* this step exists and what it unlocks next.\n", "\n", "**Kernel:** `gpt_from_scratch/.venv`. Run cells top to bottom.\n", "\n", "**Ultimate goal:** weight matrices that score the **true** next character\n", "highly on Shakespeare — so generation looks like language, not noise.\n", "\n", "**Read the chain bottom-up.** Each line needs the line below it first:\n", "```\n", "GOAL better weights -> better text generation\n", " ^ step 4 optimizer.step() uses gradients to nudge weights\n", " ^ step 3 loss.backward() computes gradient per weight\n", " ^ step 2 loss = cross_entropy(...) one number: how wrong we were\n", " ^ step 1 forward: model(xb, yb) logits compared to answer key yb\n", " ^ step 0 get_batch -> xb, yb one practice batch from train_data\n", " ^ ... tokenize, tensor tape everything in sections 0-5\n", "```\n", "\n", "| Sec | You build | So that later... |\n", "|-----|-----------|------------------|\n", "| 0 | PyTorch + GPU | tensors & weights can run fast |\n", "| 1-2 | text -> ids | model eats numbers |\n", "| 3 | tensor tape + split | sample train/val batches |\n", "| 4-5 | x/y windows, batches | one `forward` has questions + answers |\n", "| 6 | logits, loss, weights | `loss.backward()` has something to fix |\n", "| 7 | generate (untrained) | baseline before learning |\n", "| 8 | forward/backward/step | **weights actually change** |\n", "| 9 | generate (trained) | see the goal — better patterns |\n" ] }, { "cell_type": "markdown", "id": "33e0cea0", "metadata": {}, "source": [ "## 0. Setup (overview)\n", "\n", "\n", "**You are here (section 0):** install the math engine (PyTorch + GPU).\n", "\n", "**Needs:** nothing in this notebook yet.\n", "**Unlocks:** tensors on `device` (section 3 batches, section 6 weights all live here).\n", "We use **PyTorch**: tensors (arrays of numbers) + `torch.nn` (neural nets).\n", "Training is heavy math on millions of numbers — a **GPU** speeds that up 10–100×.\n", "\n", "**→ Training:** Learning = huge amounts of matrix math. Setup ensures PyTorch, tensors, and the GPU are ready to run that math millions of times.\n" ] }, { "cell_type": "markdown", "id": "91321dde", "metadata": {}, "source": [ "### 0.1 Imports\n", "\n", "These three lines appear in almost every PyTorch project.\n", "\n", "**→ Training:** `torch` runs the math, `nn` holds the **weight matrices** we train, `F` computes **loss** (how wrong we are) from predictions.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d99d4439", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "from torch.nn import functional as F\n", "print(\"PyTorch version:\", torch.__version__)\n" ] }, { "cell_type": "markdown", "id": "57084056", "metadata": {}, "source": [ "### 0.2 CPU vs GPU (`device`)\n", "\n", "Tensors can live on **CPU** (always works) or **CUDA/GPU** (fast). We pick once\n", "and reuse `device` so every tensor and model lands in the same place.\n", "\n", "If `CUDA available` is False, you installed CPU-only PyTorch — training still\n", "works but will be slow.\n", "\n", "**→ Training:** Model **weights** and each **batch** must sit on the same device. Training on GPU is the same logic, just far faster.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c3f68853", "metadata": {}, "outputs": [], "source": [ "print(\"CUDA available:\", torch.cuda.is_available())\n", "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "print(\"using device:\", device)\n", "if device == \"cuda\":\n", " print(\"GPU:\", torch.cuda.get_device_name(0))\n" ] }, { "cell_type": "markdown", "id": "6126b5a6", "metadata": {}, "source": [ "### 0.3 What is a tensor?\n", "\n", "A **tensor** is an n-dimensional array. `shape (5,)` = 5 numbers in a row.\n", "`shape (2, 3)` = 2 rows × 3 columns. Token ids will be 1-D tensors.\n", "\n", "**→ Training:** Training data (`xb`, `yb`), model **weights**, and **gradients** are all tensors. Slicing `train_data` builds each practice batch.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2b37eccd", "metadata": {}, "outputs": [], "source": [ "row = torch.tensor([10, 20, 30])\n", "grid = torch.tensor([[1, 2], [3, 4]])\n", "print(\"row shape:\", row.shape, \" values:\", row.tolist())\n", "print(\"grid shape:\", grid.shape)\n" ] }, { "cell_type": "markdown", "id": "8a1bb533", "metadata": {}, "source": [ "**Your turn 0** — Create a 2×2 tensor of zeros on `device`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "25810c05", "metadata": {}, "outputs": [], "source": [ "practice = torch.zeros(2, 2, device=device)\n", "print(practice)\n", "print(\"on device:\", practice.device.type == device)\n" ] }, { "cell_type": "markdown", "id": "bad266c2", "metadata": {}, "source": [ "## 1. The data (overview)\n", "\n", "\n", "**You are here (section 1):** load the raw training material.\n", "\n", "**Needs:** section 0 (imports).\n", "**Unlocks:** `text` -> tokenization (section 2) -> the pool of next-token questions.\n", "A **language model** learns patterns in text by repeatedly asking: *given what\n", "came before, what comes next?* No separate label file — the text **is** the labels.\n", "\n", "We use ~1M characters of Shakespeare in `input.txt`.\n", "\n", "**→ Training:** The text is our only ingredient. Each training step asks: \"given chars before, predict the next\" — labels are **inside** the text.\n" ] }, { "cell_type": "markdown", "id": "f7039a6a", "metadata": {}, "source": [ "### 1.1 Load the file\n", "\n", "Plain Python: read the whole file into one Python string `text`.\n", "\n", "**→ Training:** `text` is the raw pool of examples. Nothing trains until we encode it and sample windows from it.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "9e4fa669", "metadata": {}, "outputs": [], "source": [ "with open(\"input.txt\", \"r\", encoding=\"utf-8\") as f:\n", " text = f.read()\n", "print(\"characters in file:\", len(text))\n", "print(\"---- first 200 chars ----\")\n", "print(text[:200])\n" ] }, { "cell_type": "markdown", "id": "83f3fde7", "metadata": {}, "source": [ "### 1.2 What's inside?\n", "\n", "The model sees **every** character: letters, spaces, punctuation, newlines.\n", "Newlines matter — they mark line breaks in the plays.\n", "\n", "**→ Training:** The model learns the **statistics** of these characters (spaces, newlines, letter pairs). It does not store the file verbatim in weights.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "31f45b4a", "metadata": {}, "outputs": [], "source": [ "print(\"newlines:\", text.count(\"\\n\"))\n", "print(\"unique character count:\", len(set(text)))\n", "print(\"first line:\", repr(text.split(\"\\n\")[0][:60]))\n" ] }, { "cell_type": "markdown", "id": "9bfde4e7", "metadata": {}, "source": [ "**Your turn 1** — Print the last 100 characters of `text`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "dc5dfe07", "metadata": {}, "outputs": [], "source": [ "print(repr(text[-100:]))\n" ] }, { "cell_type": "markdown", "id": "f9a73d64", "metadata": {}, "source": [ "## 2. Tokenization (overview)\n", "\n", "\n", "**You are here (section 2):** text -> integer token ids.\n", "\n", "**Needs:** section 1 (`text`).\n", "**Unlocks:** `encode` / `vocab_size` -> tensor tape (section 3) -> model input/output size (section 6).\n", "Neural nets need **integers**. **Tokenization** = chop text into **tokens** and\n", "assign each an id.\n", "\n", "| Approach | Token | Used by |\n", "|----------|-------|---------|\n", "| Character | one letter | us (simple) |\n", "| Sub-word (BPE) | pieces like \"ing\" | GPT, Llama |\n", "\n", "We use **one char = one token** — tiny vocab, longer sequences.\n", "\n", "**→ Training:** The network only sees integers. Tokenization turns Shakespeare into the id sequences that every forward pass and loss calculation uses.\n" ] }, { "cell_type": "markdown", "id": "78bd31aa", "metadata": {}, "source": [ "### 2.1 Vocabulary\n", "\n", "The **vocabulary** = all distinct characters in the training text. `vocab_size`\n", "is how many different tokens the model can output.\n", "\n", "**→ Training:** `vocab_size` fixes the width of the output scores. Each training step: model outputs `vocab_size` logits; truth is **one** of those ids.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "07af82cf", "metadata": {}, "outputs": [], "source": [ "chars = sorted(list(set(text)))\n", "vocab_size = len(chars)\n", "print(\"vocab_size:\", vocab_size)\n", "print(\"all chars:\", repr(\"\".join(chars)))\n" ] }, { "cell_type": "markdown", "id": "9f03ff92", "metadata": {}, "source": [ "### 2.2 Lookup tables: `stoi` and `itos`\n", "\n", "- `stoi` (**s**tring **to** **i**nt): character → id\n", "- `itos` (**i**nt **to** **s**tring): id → character\n", "\n", "**→ Training:** `stoi` converts batches of text to `xb`/`yb` tensors. `itos` is for us humans when inspecting generated output.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3e0b4759", "metadata": {}, "outputs": [], "source": [ "stoi = {ch: i for i, ch in enumerate(chars)}\n", "itos = {i: ch for i, ch in enumerate(chars)}\n", "print('stoi[h] =', stoi['h'])\n", "print('itos[0] =', repr(itos[0]))\n" ] }, { "cell_type": "markdown", "id": "8f5a3085", "metadata": {}, "source": [ "### 2.3 `encode` and `decode`\n", "\n", "Encode: string → list of ids. Decode: list of ids → string. Round-trip should match.\n", "\n", "**→ Training:** `encode(text)` built the master tape `data`. Every training batch is a slice of those ids — never raw strings inside the model.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "14c4e3a3", "metadata": {}, "outputs": [], "source": [ "encode = lambda s: [stoi[c] for c in s]\n", "decode = lambda l: \"\".join(itos[i] for i in l)\n", "demo = \"hello there\"\n", "ids = encode(demo)\n", "print(demo, \"->\", ids)\n", "print(\"back:\", decode(ids))\n" ] }, { "cell_type": "markdown", "id": "40dd1ac1", "metadata": {}, "source": [ "**Your turn 2** — Encode a short phrase. Flip one id and decode — see one wrong letter.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0b50c865", "metadata": {}, "outputs": [], "source": [ "phrase = \"First Citizen\"\n", "ids = encode(phrase)\n", "ids[3] = (ids[3] + 1) % vocab_size\n", "print(\"corrupted:\", decode(ids))\n" ] }, { "cell_type": "markdown", "id": "f90ea698", "metadata": {}, "source": [ "## 3. Tensor tape + train/val split (overview)\n", "\n", "\n", "**You are here (section 3):** whole book as one integer tape + train/val split.\n", "\n", "**Needs:** section 2 (`encode`, `vocab_size`).\n", "**Unlocks:** `train_data` to sample batches (section 5); `val_data` to check generalisation (section 8).\n", "Section 2 left us with Python lists of ids. Here we wrap the whole book in a\n", "**tensor** (3.1), then split it into train and validation (3.2–3.4).\n", "\n", "**→ Training:** `train_data` is what we sample from when updating weights. `val_data` checks whether those weights work on **held-out** text.\n" ] }, { "cell_type": "markdown", "id": "5a1caac1", "metadata": {}, "source": [ "### 3.1 What is a tensor?\n", "\n", "In section 2 we had a Python **list** of token ids: `[18, 47, 56, ...]`.\n", "PyTorch uses a **tensor** for the same idea — a block of numbers the library\n", "can slice, stack, and run on a GPU very fast.\n", "\n", "**Think of it like this:** a list is a handwritten shopping list. A tensor is\n", "the same list typed into a spreadsheet — same items, but the computer can do\n", "math on the whole column at once.\n", "\n", "#### How many dimensions? (`shape`)\n", "\n", "| Example | shape | Meaning |\n", "|---------|-------|---------|\n", "| `torch.tensor(5)` | `()` | one number (scalar) |\n", "| `torch.tensor([1,2,3])` | `(3,)` | one row of 3 ids — **our book tape** |\n", "| `torch.tensor([[1,2],[3,4]])` | `(2, 2)` | 2 rows × 2 cols — later `(B, T)` batches |\n", "\n", "The comma in `(3,)` means \"one dimension of length 3\" — not a tuple pair.\n", "\n", "#### `dtype` — what kind of number?\n", "\n", "- `torch.long` — **integers** (token ids). No decimals.\n", "- `torch.float32` — **decimals** (model weights, logits, loss).\n", "\n", "Token ids are always integers, so we use `dtype=torch.long`.\n", "\n", "#### Why not stay with Python lists?\n", "\n", "1. **Speed** — millions of multiply/adds per second on GPU.\n", "2. **Slicing** — `data[1000:1008]` gives another tensor, no copy of the whole book.\n", "3. **Gradients** — later, PyTorch tracks how each weight affects loss (training).\n", "\n", "Run the cell below: same ids as a list vs as a tensor.\n", "\n", "**→ Training:** We store the whole book as one tensor so we can grab random windows fast and (later) move batches to the GPU in one shot.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "21051ca8", "metadata": {}, "outputs": [], "source": [ "ids_list = encode(\"First\")\n", "ids_tensor = torch.tensor(ids_list, dtype=torch.long)\n", "\n", "print(\"Python list:\", ids_list)\n", "print(\"tensor: \", ids_tensor)\n", "print(\"shape:\", ids_tensor.shape)\n", "print(\"dtype:\", ids_tensor.dtype)\n", "print()\n", "print(\"slice [1:3]:\", ids_tensor[1:3])\n", "print(\"decode slice:\", decode(ids_tensor[1:3].tolist()))\n", "print(\"back to list:\", ids_tensor.tolist())\n" ] }, { "cell_type": "markdown", "id": "7b498e68", "metadata": {}, "source": [ "**Your turn 3.1** — Encode a word of your choice. Build a tensor, print `.shape`\n", "and `.dtype`. Slice the middle two ids and decode them.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c40b416b", "metadata": {}, "outputs": [], "source": [ "word = \"Citizen\" # edit me\n", "t = torch.tensor(encode(word), dtype=torch.long)\n", "print(\"shape:\", t.shape, \" dtype:\", t.dtype)\n", "mid = t[1:3]\n", "print(\"middle ids:\", mid.tolist(), \"->\", decode(mid.tolist()))\n" ] }, { "cell_type": "markdown", "id": "d1c84cec", "metadata": {}, "source": [ "### 3.2 String → tensor (the whole book)\n", "\n", "Now we wrap **every** token id in the dataset into one 1-D tensor `data`.\n", "Same idea as 3.1, but ~1 million ids long — our **tape**.\n", "\n", "`dtype=torch.long` = integers only. Shape `(N,)` = one long row of length `N`.\n", "\n", "**→ Training:** `data` / `train_data` never change during training. Only the **model weights** change. The tape is fixed questions; weights are learned answers.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d419690c", "metadata": {}, "outputs": [], "source": [ "data = torch.tensor(encode(text), dtype=torch.long)\n", "print(\"shape:\", data.shape, \" dtype:\", data.dtype)\n", "print(\"first 10 ids:\", data[:10].tolist())\n", "print(\"first 10 chars:\", decode(data[:10].tolist()))\n" ] }, { "cell_type": "markdown", "id": "5450cf66", "metadata": {}, "source": [ "### 3.3 Train / validation split\n", "\n", "- **train_data** — model learns from this\n", "- **val_data** — later we check loss on unseen tail of the book\n", "\n", "If train loss falls but val loss stays high → **memorising**, not generalising.\n", "\n", "**→ Training:** We **only** run `backward()` on train batches. Val loss is read-only: if train improves but val does not, weights are memorising.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "bbc77a95", "metadata": {}, "outputs": [], "source": [ "n = int(0.9 * len(data))\n", "train_data = data[:n]\n", "val_data = data[n:]\n", "print(\"train:\", len(train_data), \" val:\", len(val_data))\n", "print(\"sum check:\", len(train_data) + len(val_data) == len(data))\n" ] }, { "cell_type": "markdown", "id": "dcf68150", "metadata": {}, "source": [ "### 3.4 The split boundary in plain text\n", "\n", "Train ends and val begins at the same place in the original book — no gap, no overlap.\n", "\n", "**→ Training:** Confirms train and val are contiguous slices — so val loss is a fair \"unseen tail of the book\" check.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e9bc8a68", "metadata": {}, "outputs": [], "source": [ "print(\"...end of train:\", repr(decode(train_data[-40:].tolist())))\n", "print(\"start of val: \", repr(decode(val_data[:40].tolist())))\n" ] }, { "cell_type": "markdown", "id": "51b55dd8", "metadata": {}, "source": [ "**Your turn 3** — What fraction of the book is validation? *(~10%.)*\n" ] }, { "cell_type": "code", "execution_count": null, "id": "70e5f2ac", "metadata": {}, "outputs": [], "source": [ "frac_val = len(val_data) / len(data)\n", "print(f\"validation fraction: {frac_val:.1%}\")\n" ] }, { "cell_type": "markdown", "id": "fe71c7d8", "metadata": {}, "source": [ "## 4. Predicting the next token (overview)\n", "\n", "\n", "**You are here (section 4):** define ONE training question (context -> next token).\n", "\n", "**Needs:** section 3 (`train_data` tape).\n", "**Unlocks:** `x` / `y` / `block_size` -> batched training examples (section 5) -> loss targets (section 6).\n", "A language model is a **next-token guesser**. You show it some text; it scores\n", "every possible next token; you compare to the real answer and nudge its weights.\n", "\n", "Section 4 breaks that idea into small pieces. **Read one part, run the code\n", "right below it**, then move on. By the end you'll understand chunks, `block_size`,\n", "`x`/`y`, and the prefix trick — the foundation for batching in section 5.\n", "\n", "**→ Training:** The **objective** for the whole notebook: minimise surprise at the true next token. Everything in §4 defines one training example.\n" ] }, { "cell_type": "markdown", "id": "e7a8f7c0", "metadata": {}, "source": [ "### 4.1 The next-token game (one step)\n", "\n", "The training question is always the same:\n", "\n", "> Given what came before → what is the **single** next character?\n", "\n", "There is no separate labels file. The label is hiding **in the text itself**:\n", "the next character in the sequence.\n", "\n", "**Analogy:** a friend starts a sentence and you blurt out the next word. They\n", "tell you if you were right. That's one training step.\n", "\n", "**→ Training:** One row of the loss compares model scores to **one** correct next id. A batch contains `B × T` such comparisons at once.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6adf5e2f", "metadata": {}, "outputs": [], "source": [ "# tiny phrase from our vocabulary — watch one step at a time\n", "phrase = \"First Citizen\"\n", "ids = encode(phrase)\n", "print(\"text: \", phrase)\n", "print(\"ids: \", ids)\n", "print()\n", "for i in range(len(ids) - 1):\n", " context = decode(ids[: i + 1])\n", " nxt = itos[ids[i + 1]]\n", " print(f\" seen {context!r:20s} -> predict {nxt!r}\")\n" ] }, { "cell_type": "markdown", "id": "537d244e", "metadata": {}, "source": [ "**Your turn 4.1** — Change `phrase` to another short string (letters/spaces only).\n", "How many training steps does a string of length *L* give? *(Answer: L − 1.)*\n" ] }, { "cell_type": "code", "execution_count": null, "id": "49b69d41", "metadata": {}, "outputs": [], "source": [ "phrase = \"To be\" # edit me\n", "ids = encode(phrase)\n", "steps = len(ids) - 1\n", "print(f\"length {len(ids)} -> {steps} next-token questions\")\n", "for i in range(steps):\n", " print(f\" {decode(ids[:i+1])!r} -> {itos[ids[i+1]]!r}\")\n" ] }, { "cell_type": "markdown", "id": "a5039f5f", "metadata": {}, "source": [ "### 4.2 The book is one long tape\n", "\n", "After section 3, the whole dataset is a single 1-D tensor `train_data` — a tape\n", "of token ids with no breaks between chapters:\n", "\n", "```\n", "train_data: [18, 47, 56, 57, ...] ← millions of ids, one after another\n", "```\n", "\n", "Training never asks \"predict the whole book.\" It only ever asks about **local**\n", "stretches of that tape.\n", "\n", "**→ Training:** Every example is a contiguous window on this tape. Random `start` indices = seeing different parts of the book each step.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d217e75d", "metadata": {}, "outputs": [], "source": [ "print(\"train_data length:\", len(train_data))\n", "print(\"first 15 ids: \", train_data[:15].tolist())\n", "print(\"first 15 chars: \", decode(train_data[:15].tolist()))\n", "print(\"chars 15-30: \", decode(train_data[15:30].tolist()))\n", "print()\n", "print(\"The tape is continuous — char 15 follows char 14 in the real book.\")\n" ] }, { "cell_type": "markdown", "id": "0ecf904b", "metadata": {}, "source": [ "### 4.3 Why we use chunks (`block_size`)\n", "\n", "We could slide one character at a time, but in practice we grab a **window** of\n", "`block_size` characters at once. Three reasons:\n", "\n", "1. **Memory** — models store state per position; a 1M-wide window would not fit on a GPU.\n", "2. **Fixed context** — Transformers are built to look at at most `block_size` tokens.\n", "3. **Efficiency** — one forward pass on 8 positions trains 8 questions in parallel (next part).\n", "\n", "`block_size` is the **context window**: how far back the model is *allowed* to look.\n", "We use 8 here so the numbers fit on screen; GPT-4 uses thousands.\n", "\n", "**→ Training:** `block_size` = how many next-token questions per row (`T`). Also caps context length the architecture can use.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "7b330af9", "metadata": {}, "outputs": [], "source": [ "block_size = 8\n", "start = 1000 # where we cut the tape\n", "chunk = train_data[start : start + block_size]\n", "\n", "print(f\"block_size = {block_size}\")\n", "print(f\"slice from index {start} to {start + block_size - 1}\")\n", "print(\"chunk ids: \", chunk.tolist())\n", "print(\"chunk text: \", repr(decode(chunk.tolist())))\n", "print()\n", "print(\"book length:\", len(train_data), \" window:\", block_size,\n", " \" ratio:\", round(len(train_data) / block_size))\n" ] }, { "cell_type": "markdown", "id": "f19d6e14", "metadata": {}, "source": [ "**Your turn 4.3** — Change `start`. The chunk should always be exactly `block_size`\n", "characters (except if you pick a start too close to the end).\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8528eb17", "metadata": {}, "outputs": [], "source": [ "start = 25000 # try different values\n", "chunk = train_data[start : start + block_size]\n", "print(\"len(chunk):\", len(chunk), \" text:\", repr(decode(chunk.tolist())))\n" ] }, { "cell_type": "markdown", "id": "342c105f", "metadata": {}, "source": [ "### 4.4 Inputs `x` and targets `y` (shift by one)\n", "\n", "From one slice we build two equal-length tensors:\n", "\n", "- **`x`** = the window `train_data[start : start + block_size]`\n", "- **`y`** = the same window, shifted **one step forward** in time:\n", " `train_data[start + 1 : start + block_size + 1]`\n", "\n", "So `y` is always \"what comes next\" relative to `x`, position by position:\n", "\n", "```\n", "x: h e l l o _ t h\n", "y: e l l o _ t h ? ← ? is the 9th char (first char AFTER the window)\n", "```\n", "\n", "The model never sees `y` during the forward pass — only during loss calculation.\n", "\n", "#### Tie-in to the training goal\n", "\n", "**Ultimate goal:** nudge **weights** so the model scores the true next character highly.\n", "\n", "`x` and `y` are how we package **one practice round** for that goal:\n", "\n", "| Tensor | Role in training | Analogy |\n", "|--------|------------------|---------|\n", "| **`x`** (later `xb`) | **Question** — tokens the model reads to produce logits | exam paper |\n", "| **`y`** (later `yb`) | **Answer key** — the correct next token at each position | mark scheme |\n", "\n", "At column `t`: model reads `x[t]`, predicts a distribution; loss compares it to **`y[t]`** (the char that really followed on the tape). Wrong guess → higher loss → `backward()` pushes weights to do better next time.\n", "\n", "**In the chain (why this section exists):**\n", "```\n", "GOAL: better weights\n", " ^ needs backward()\n", " ^ needs loss (one number)\n", " ^ needs yb — the CORRECT next tokens ← y is the template for yb\n", " ^ needs forward on xb — model predictions ← x is the template for xb\n", "```\n", "\n", "Without the shift-by-one split you have text but **no labelled examples** — nothing to score, no loss, no training.\n", "\n", "**→ Training:** `get_batch` builds matrices `xb`/`yb` the same way as `x`/`y` here. Every training step is: predict from `xb`, grade against `yb`, update weights.\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0dec560b", "metadata": {}, "outputs": [], "source": [ "start = 1000\n", "x = train_data[start : start + block_size]\n", "y = train_data[start + 1 : start + block_size + 1]\n", "\n", "print(\"x:\", decode(x.tolist()))\n", "print(\"y:\", decode(y.tolist()))\n", "print()\n", "for i in range(block_size):\n", " print(f\" x[{i}]={repr(itos[x[i].item()])} y[{i}]={repr(itos[y[i].item()])}\")\n" ] }, { "cell_type": "markdown", "id": "70613e82", "metadata": {}, "source": [ "**Your turn 4.4** — Confirm `y[i]` is always the id **right after** `x[i]` on the tape.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5be4764d", "metadata": {}, "outputs": [], "source": [ "ok = all(y[t].item() == train_data[start + t + 1].item() for t in range(block_size))\n", "print(\"every y[t] is the next char after x[t] on the tape:\", ok)\n", "if not ok:\n", " for t in range(block_size):\n", " print(t, x[t].item(), y[t].item(), train_data[start + t + 1].item())\n" ] }, { "cell_type": "markdown", "id": "6fecc0fb", "metadata": {}, "source": [ "### 4.5 The prefix trick (one chunk → many lessons)\n", "\n", "Here is the key idea that confuses many beginners:\n", "\n", "An 8-character chunk is **not** one training example. It is **eight** examples\n", "packed together — one for each prefix length.\n", "\n", "| step `t` | context (prefix of `x`) | must predict `y[t]` |\n", "|----------|---------------------------|---------------------|\n", "| 0 | 1st char only | 2nd char |\n", "| 1 | 1st + 2nd | 3rd char |\n", "| 2 | first 3 chars | 4th char |\n", "| … | … | … |\n", "| 7 | all 8 chars | 9th char (just past the window) |\n", "\n", "**Analogy:** teaching spelling by asking \"after **h**?\", then \"after **he**?\",\n", "then \"after **hel**?\" — same word, harder each time.\n", "\n", "**→ Training:** One slice yields `block_size` supervised examples per forward pass — more learning signal per batch.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e4559d0f", "metadata": {}, "outputs": [], "source": [ "start = 1000\n", "x = train_data[start : start + block_size]\n", "y = train_data[start + 1 : start + block_size + 1]\n", "\n", "print(\"full chunk:\", repr(decode(x.tolist())))\n", "print()\n", "for t in range(block_size):\n", " prefix = decode(x[: t + 1].tolist())\n", " target = itos[y[t].item()]\n", " print(f\"t={t} prefix {prefix!r:12s} -> {target!r}\")\n" ] }, { "cell_type": "markdown", "id": "3be6fbb3", "metadata": {}, "source": [ "**Your turn 4.5** — At `t=0`, how many characters does the model see? At `t=7`?\n", "*(0: one char; 7: all eight.)*\n" ] }, { "cell_type": "code", "execution_count": null, "id": "898927dc", "metadata": {}, "outputs": [], "source": [ "for t in [0, 3, 7]:\n", " prefix_len = t + 1\n", " print(f\"t={t}: prefix length = {prefix_len}, target = {repr(itos[y[t].item()])}\")\n" ] }, { "cell_type": "markdown", "id": "36cde17a", "metadata": {}, "source": [ "### 4.6 Column alignment (what `x[t]` and `y[t]` mean together)\n", "\n", "When we later batch data into a matrix, **column `t`** means:\n", "\n", "- input token at that step: `x[t]`\n", "- correct next token: `y[t]`\n", "\n", "Important: `y[t]` is the character that follows **`x[t]`** in the original book,\n", "not \"the char after the whole prefix.\" (A Transformer uses the full prefix via\n", "attention; our Bigram model will only look at `x[t]` — more in section 6.)\n", "\n", "**→ Training:** At column `t`, loss asks: \"after token `xb[b,t]`, was your distribution close to `yb[b,t]`?\"\n" ] }, { "cell_type": "code", "execution_count": null, "id": "876ad4ce", "metadata": {}, "outputs": [], "source": [ "start = 1000\n", "x = train_data[start : start + block_size]\n", "y = train_data[start + 1 : start + block_size + 1]\n", "\n", "for t in range(block_size):\n", " on_tape = train_data[start + t].item()\n", " on_tape_next = train_data[start + t + 1].item()\n", " match = x[t].item() == on_tape and y[t].item() == on_tape_next\n", " print(f\"t={t} x={repr(itos[x[t].item()])} y={repr(itos[y[t].item()])} matches tape: {match}\")\n" ] }, { "cell_type": "markdown", "id": "78bef65a", "metadata": {}, "source": [ "### 4.7 Shape of one chunk\n", "\n", "Right now `x` and `y` are **1-D** tensors of shape `(block_size,)` — one chunk,\n", "one row. In section 5 we stack many chunks:\n", "\n", "- `(B, T)` where **B** = batch size, **T** = `block_size`\n", "- Later the model outputs `(B, T, C)` — **C** = `vocab_size` scores per step\n", "\n", "Same idea, just more rows for GPU parallelism.\n", "\n", "**→ Training:** `(B, T)` tells us we get **B×T** gradient signals per step. Later `(B, T, C)` is one score vector per position for cross-entropy.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8bb27b4f", "metadata": {}, "outputs": [], "source": [ "x = train_data[1000 : 1000 + block_size]\n", "y = train_data[1001 : 1001 + block_size]\n", "\n", "print(\"x shape:\", x.shape, \" y shape:\", y.shape)\n", "print(\"T (time steps) = block_size =\", block_size)\n", "print()\n", "# pretend we had batch_size=3 — same T, three independent rows\n", "B = 3\n", "print(f\"batched inputs would be shape ({B}, {block_size}) = (B, T)\")\n", "print(f\"model logits later: ({B}, {block_size}, {vocab_size}) = (B, T, C)\")\n" ] }, { "cell_type": "markdown", "id": "f86da52d", "metadata": {}, "source": [ "### 4.8 Put it together — build your own `x` and `y`\n", "\n", "Pick any valid `start` index, build `x` and `y`, decode them, and verify every\n", "target. This is exactly what `get_batch` does — but with random `start` values.\n", "\n", "**→ Training:** Same logic as `get_batch` — you're manually doing what the training loop automates thousands of times.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c2ee21e5", "metadata": {}, "outputs": [], "source": [ "start = 5000 # change me\n", "x = train_data[start : start + block_size]\n", "y = train_data[start + 1 : start + block_size + 1]\n", "\n", "print(\"start:\", start)\n", "print(\"x:\", decode(x.tolist()))\n", "print(\"y:\", decode(y.tolist()))\n", "checks = [y[t].item() == train_data[start + t + 1].item() for t in range(block_size)]\n", "print(\"all correct:\", all(checks))\n" ] }, { "cell_type": "markdown", "id": "c1d3c8ce", "metadata": {}, "source": [ "### 4.9 Why this matters for the Bigram model\n", "\n", "We feed 8 positions per row, and the loss averages over **all** of them — so one\n", "chunk gives 8 gradients per row. But the Bigram architecture only uses the\n", "**current** character at each position, not the full prefix. The data is ready for\n", "long context; the model is not. Fixing that is what **self-attention** does later.\n", "\n", "**Section 4 checklist**\n", "- [ ] Next-token = one step forward on the tape\n", "- [ ] `block_size` = window length / context limit\n", "- [ ] `y` is `x` shifted by one\n", "- [ ] One chunk = `block_size` prefix drills\n", "- [ ] Shapes: `(T,)` now → `(B, T)` in section 5\n", "\n", "**→ Training:** Data supplies full prefixes; Bigram weights only use the current char. Training still works, but capacity is limited until attention.\n" ] }, { "cell_type": "markdown", "id": "45dafc47", "metadata": {}, "source": [ "## 5. Batching (overview)\n", "\n", "\n", "**You are here (section 5):** many `x`/`y` windows at once as matrices.\n", "\n", "**Needs:** sections 3-4 (`train_data`, `block_size`, shift-by-one).\n", "**Unlocks:** `get_batch()` -> every `forward` in training (section 8) starts here.\n", "Section 4 used **one** chunk (shape `(T,)`). GPUs want **many** chunks at once.\n", "We stack `batch_size` independent windows → shape **`(B, T)`**.\n", "\n", "- **B** = `batch_size` (parallel snippets)\n", "- **T** = `block_size` (time steps / context length)\n", "\n", "**→ Training:** Each optimizer step uses one batch `(xb, yb)`. Larger `B` = more examples contributing to the **same** weight update.\n" ] }, { "cell_type": "markdown", "id": "ffc77ac6", "metadata": {}, "source": [ "### 5.1 Why batch?\n", "\n", "**Analogy:** grading one exam at a time vs a stack of 32 exams — same rules,\n", "more throughput. Random start indices each step = model sees different parts of\n", "the book every iteration.\n", "\n", "**→ Training:** Random windows each step ≈ shuffled exposure to the book. Parallel rows = faster steps on GPU.\n" ] }, { "cell_type": "markdown", "id": "0b02e589", "metadata": {}, "source": [ "### 5.2 Stack two chunks by hand\n", "\n", "Before `get_batch`, see how two slices become a 2×8 matrix.\n", "\n", "**→ Training:** `get_batch` does this with `batch_size` rows. One `optimizer.step()` updates weights using loss from **all** rows.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "11fc074b", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(0)\n", "ix = torch.randint(len(train_data) - block_size, (2,))\n", "chunk_a = train_data[ix[0] : ix[0] + block_size]\n", "chunk_b = train_data[ix[1] : ix[1] + block_size]\n", "stacked = torch.stack([chunk_a, chunk_b])\n", "print(\"stacked shape:\", stacked.shape, \" (= 2 rows, block_size cols)\")\n", "for row in range(2):\n", " print(f\" row {row}:\", decode(stacked[row].tolist()))\n" ] }, { "cell_type": "markdown", "id": "c296b808", "metadata": {}, "source": [ "### 5.3 `get_batch` — random windows\n", "\n", "`xb` = inputs, `yb` = targets (each row shifted +1 on the tape). Same shapes.\n", "\n", "**→ Training:** Called every training iteration. Returns the `(xb, yb)` pair that `loss.backward()` will score and correct.\n", "\n", "**In the chain:** `get_batch` is **step 0** of every training iteration — supplies `xb` (questions) and `yb` (answers) before `model()` can run.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "67c1ce6f", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(1337)\n", "batch_size = 4\n", "\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 + block_size + 1] for i in ix])\n", " return x.to(device), y.to(device)\n", "\n", "xb, yb = get_batch(\"train\")\n", "print(\"xb shape:\", xb.shape, \" yb shape:\", yb.shape)\n", "print(xb)\n" ] }, { "cell_type": "markdown", "id": "e72033cb", "metadata": {}, "source": [ "### 5.4 Decode a batch row\n", "\n", "Row `b`, column `t`: `yb[b,t]` is the next token after `xb[b,t]` on the tape.\n", "\n", "**→ Training:** Sanity check that `yb` is the answer key aligned with `xb` before trusting the loss number.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d8be50cf", "metadata": {}, "outputs": [], "source": [ "b = 0\n", "print(\"row\", b, \"input: \", decode(xb[b].tolist()))\n", "print(\"row\", b, \"target:\", decode(yb[b].tolist()))\n", "for t in range(block_size):\n", " print(f\" t={t} x={repr(itos[xb[b,t].item()])} y={repr(itos[yb[b,t].item()])}\")\n" ] }, { "cell_type": "markdown", "id": "3f69fc3b", "metadata": {}, "source": [ "**Your turn 5** — Call `get_batch(\"val\")` and decode row 2.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "cd1c51dc", "metadata": {}, "outputs": [], "source": [ "xb_val, yb_val = get_batch(\"val\")\n", "print(decode(xb_val[2].tolist()))\n" ] }, { "cell_type": "markdown", "id": "883c2705", "metadata": {}, "source": [ "### 5.5 `.to(device)`\n", "\n", "Moves tensors to GPU so the model (also on GPU) can process them in one place.\n", "\n", "**Section 5 checklist**\n", "- [ ] `(B, T)` = batch of `block_size`-long windows\n", "- [ ] `yb` is `xb` shifted +1 on the tape\n", "- [ ] `get_batch` randomises which part of the book each step\n", "\n", "**→ Training:** Moves batch tensors next to **weight matrices** on GPU so forward + backward run without device errors.\n" ] }, { "cell_type": "markdown", "id": "0e99e2c5", "metadata": {}, "source": [ "## 6. Logits, loss, and the Bigram model (overview)\n", "\n", "\n", "**You are here (section 6):** the **weights** + how we score wrongness.\n", "\n", "**Needs:** section 5 (`xb`, `yb` shapes); section 2 (`vocab_size`).\n", "**Unlocks:** `model(xb,yb)->loss` for `backward()` (section 8); `model(xb)->sample` for generation (section 7).\n", "The **Bigram** model: next character depends **only** on the current character.\n", "Before the full class, we unpack **logits**, **softmax**, and **cross-entropy** —\n", "the same three ideas GPT uses at scale.\n", "\n", "**→ Training:** Forward produces **logits**; loss measures wrongness; backward rewrites the **embedding matrix** — the only weights in this model.\n" ] }, { "cell_type": "markdown", "id": "f6bb711d", "metadata": {}, "source": [ "### 6.1 Logits — raw scores\n", "\n", "For each possible next token, the model outputs a **logit** (unnormalised score).\n", "Higher = more likely. Length = `vocab_size`.\n", "\n", "**→ Training:** Logits are not trained directly — **weights** produce logits. Gradients flow back through logits into those weights.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e7428034", "metadata": {}, "outputs": [], "source": [ "toy_logits = torch.tensor([2.0, 0.5, 0.1]) # pretend vocab_size=3\n", "print(\"logits:\", toy_logits.tolist())\n", "print(\"argmax (favourite class):\", toy_logits.argmax().item())\n" ] }, { "cell_type": "markdown", "id": "ff289741", "metadata": {}, "source": [ "### 6.2 Softmax — scores → probabilities\n", "\n", "Softmax squashes logits into positive numbers that **sum to 1** — a distribution.\n", "\n", "**→ Training:** Used heavily in **generation** (sampling). Training uses cross-entropy on logits directly (softmax is inside that).\n" ] }, { "cell_type": "code", "execution_count": null, "id": "71e447c8", "metadata": {}, "outputs": [], "source": [ "probs = F.softmax(toy_logits, dim=0)\n", "print(\"probs:\", [round(p, 3) for p in probs.tolist()])\n", "print(\"sum:\", probs.sum().item())\n" ] }, { "cell_type": "markdown", "id": "e60d5dd1", "metadata": {}, "source": [ "### 6.3 Cross-entropy loss — how wrong were we?\n", "\n", "Compare model distribution to the **true** next token id.\n", "- right answer gets high prob → low loss\n", "- wrong answer → high loss\n", "\n", "Random guessing over `vocab_size` tokens → loss ≈ `-ln(1/vocab_size)`.\n", "\n", "**→ Training:** The single number `loss` that `backward()` uses. Lower loss ⇒ nudge weights so true next tokens get higher score.\n", "\n", "**In the chain:** cross-entropy turns logits + `yb` into the **single number** `loss`. Section 8 cannot run `backward()` until this exists.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c9166861", "metadata": {}, "outputs": [], "source": [ "for target in [0, 1, 2]:\n", " loss = F.cross_entropy(toy_logits.unsqueeze(0), torch.tensor([target]))\n", " print(f\"true class {target} -> loss {loss.item():.3f}\")\n", "expected = -torch.log(torch.tensor(1.0 / vocab_size))\n", "print(f\"random baseline (vocab {vocab_size}): {expected.item():.3f}\")\n" ] }, { "cell_type": "markdown", "id": "7b45f3a7", "metadata": {}, "source": [ "**Your turn 6** — Make class 2 strongest; loss for target 2 should drop.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "dfe98faf", "metadata": {}, "outputs": [], "source": [ "strong = torch.tensor([0.1, 0.5, 2.5])\n", "for target in [0, 2]:\n", " loss = F.cross_entropy(strong.unsqueeze(0), torch.tensor([target]))\n", " print(f\"target {target} -> {loss.item():.3f}\")\n" ] }, { "cell_type": "markdown", "id": "7926ed10", "metadata": {}, "source": [ "### 6.4 Bigram = one lookup table\n", "\n", "`nn.Embedding(vocab_size, vocab_size)` stores a matrix: **row i** = logits for\n", "\"what follows token i?\" No memory of earlier tokens.\n", "\n", "**Analogy:** a cheat sheet — \"after `t` usually comes `h` or `e`\" — but only\n", "looking at the **last** letter.\n", "\n", "**→ Training:** The entire learned model is this **one matrix** (~vocab² numbers). Training adjusts every row: \"what follows char i?\"\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f705fa50", "metadata": {}, "outputs": [], "source": [ "table = nn.Embedding(vocab_size, vocab_size)\n", "ch = 't'\n", "tid = stoi[ch]\n", "row = table.weight[tid]\n", "print(\"row for\", repr(ch), \"has shape:\", row.shape)\n", "top_p, top_i = F.softmax(row, dim=0).topk(3)\n", "print(\"top-3 followers (random weights):\")\n", "for p, i in zip(top_p, top_i):\n", " print(f\" {repr(itos[i.item()])}: {p.item():.3f}\")\n" ] }, { "cell_type": "markdown", "id": "e5208051", "metadata": {}, "source": [ "### 6.5 The full `BigramLanguageModel`\n", "\n", "Forward: `idx` `(B,T)` → logits `(B,T,C)`. Flatten to `(B*T, C)` for one big\n", "cross-entropy over every position in the batch.\n", "\n", "**→ Training:** `forward(xb, yb)` = predict on all positions, compare to `yb`, return scalar loss. `generate()` uses same weights without `yb`.\n", "\n", "**In the chain:** `forward(xb, yb)` is **step 1** of training. Without `yb`, you get logits only (generation path) — no loss, no backward.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f7cc9f29", "metadata": {}, "outputs": [], "source": [ "class BigramLanguageModel(nn.Module):\n", " def __init__(self, vocab_size):\n", " super().__init__()\n", " self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n", "\n", " def forward(self, idx, targets=None):\n", " logits = self.token_embedding_table(idx)\n", " if targets is None:\n", " return logits, None\n", " B, T, C = logits.shape\n", " loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))\n", " return logits, loss\n", "\n", " def generate(self, idx, max_new_tokens):\n", " for _ in range(max_new_tokens):\n", " logits, _ = self(idx)\n", " logits = logits[:, -1, :]\n", " probs = F.softmax(logits, dim=-1)\n", " idx_next = torch.multinomial(probs, num_samples=1)\n", " idx = torch.cat((idx, idx_next), dim=1)\n", " return idx\n", "\n", "model = BigramLanguageModel(vocab_size).to(device)\n", "xb, yb = get_batch(\"train\")\n", "logits, loss = model(xb, yb)\n", "print(\"logits shape (B, T, C):\", logits.shape)\n", "print(f\"initial loss: {loss.item():.4f}\")\n" ] }, { "cell_type": "markdown", "id": "1614313d", "metadata": {}, "source": [ "**Your turn 6b** — Initial loss should be near `-ln(1/vocab_size)`. Compare.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "25b7140e", "metadata": {}, "outputs": [], "source": [ "baseline = -torch.log(torch.tensor(1.0 / vocab_size)).item()\n", "print(f\"loss {loss.item():.4f} vs random baseline {baseline:.4f}\")\n" ] }, { "cell_type": "markdown", "id": "23f37864", "metadata": {}, "source": [ "**Section 6 checklist**\n", "- [ ] Logits → softmax → probs (generation)\n", "- [ ] Cross-entropy compares preds to true `yb` (training)\n", "- [ ] Bigram row `i` = next-token scores after char `i` only\n" ] }, { "cell_type": "markdown", "id": "8f6036e2", "metadata": {}, "source": [ "## 7. Generation before training (overview)\n", "\n", "\n", "**You are here (section 7):** use weights **without** training (no `yb`, no `backward`).\n", "\n", "**Needs:** section 6 (`model`, random weights).\n", "**Unlocks:** baseline gibberish to compare after section 9 — proof that training changed weights.\n", "**Generation** = autoregressive sampling: predict one token, append it, repeat.\n", "With random weights, output is gibberish — our **before** snapshot.\n", "\n", "**→ Training:** With **random** weights, generation is nonsense. This is the baseline before any `optimizer.step()` changes the matrices.\n" ] }, { "cell_type": "markdown", "id": "f51c9655", "metadata": {}, "source": [ "### 7.1 One generation step by hand\n", "\n", "Context → forward → logits at **last** position → softmax → sample one id.\n", "\n", "**→ Training:** Same forward path as training, but no `yb` and no `backward()`. Shows what one row of trained weights **does** at inference.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f9cdda7f", "metadata": {}, "outputs": [], "source": [ "context = torch.zeros((1, 1), dtype=torch.long, device=device) # start: newline\n", "logits, _ = model(context)\n", "last_logits = logits[0, -1, :]\n", "probs = F.softmax(last_logits, dim=0)\n", "sampled = torch.multinomial(probs, num_samples=1)\n", "print(\"context:\", repr(itos[context[0,0].item()]))\n", "print(\"sampled next char:\", repr(itos[sampled.item()]))\n" ] }, { "cell_type": "markdown", "id": "f8461a58", "metadata": {}, "source": [ "### 7.2 Full generation (200 chars, untrained)\n", "\n", "Start from newline (id 0). Expect nonsense.\n", "\n", "**→ Training:** Save mentally (or in output) for contrast with section 9 after weights have been updated.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "4b5c3e45", "metadata": {}, "outputs": [], "source": [ "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n", "out = model.generate(context, max_new_tokens=200)\n", "print(decode(out[0].tolist()))\n" ] }, { "cell_type": "markdown", "id": "eb818ac2", "metadata": {}, "source": [ "## 8. Training (overview)\n", "\n", "**You are here (section 8):** run the full chain thousands of times.\n", "\n", "**Needs:** ALL above — especially `get_batch` (5), `model`+loss (6).\n", "**Unlocks:** updated weights — the **result** of training.\n", "\n", "One iteration — **do not reorder**:\n", "\n", "```\n", "prereq: xb,yb = get_batch() # step 0 — sections 4-5\n", "step 1: logits, loss = model(xb,yb) # forward — section 6\n", "step 2: (loss is ready) # scalar surprise\n", "step 3: loss.backward() # gradients per weight\n", "step 4: optimizer.step() # nudge weights\n", "```\n", "\n", "Plus `zero_grad` before step 1 each time — wipe last batch's gradients.\n", "\n", "**-> Training:** This section is where the ultimate goal happens: weights change.\n" ] }, { "cell_type": "markdown", "id": "4ac65ed4", "metadata": {}, "source": [ "### 8.0 What training is — fundamentals\n", "\n", "Strip away the jargon and training is four ideas:\n", "\n", "1. **Weights** — numbers in matrices (our Bigram: one `vocab_size × vocab_size` table).\n", "2. **Forward** — multiply/lookup through weights → predictions (logits).\n", "3. **Loss** — one number: how wrong predictions are vs true next tokens `yb`.\n", "4. **Update** — nudge every weight a tiny bit to lower loss (`backward` + `step`).\n", "\n", "**Analogy:** weights = knobs on a mixer; loss = how bad the song sounds; training turns knobs repeatedly until the mix fits the data.\n", "\n", "The model does **not** store the book inside weights. It stores **habits** like \"after `t` often comes `h`.\"\n", "\n", "**→ Training:** This is the goal every earlier section prepared: data → batches → loss → changed weights.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0881b2c7", "metadata": {}, "outputs": [], "source": [ "# THE trainable artifact: model parameters (weight matrices)\n", "n_params = sum(p.numel() for p in model.parameters())\n", "print(\"trainable parameters:\", n_params)\n", "print(\"embedding weight shape:\", model.token_embedding_table.weight.shape)\n", "print(\"dtype:\", model.token_embedding_table.weight.dtype)\n", "print()\n", "print(\"These numbers ARE what training rewrites.\")\n" ] }, { "cell_type": "markdown", "id": "362158ed", "metadata": {}, "source": [ "### 8.0b One step = forward, loss, backward, update\n", "\n", "| Step | What happens | Uses from earlier sections |\n", "|------|----------------|---------------------------|\n", "| forward | `model(xb, yb)` → logits, loss | `get_batch`, `xb`/`yb`, embedding |\n", "| zero_grad | clear old gradients | — |\n", "| backward | `loss.backward()` computes ∂loss/∂weight | cross-entropy on logits |\n", "| step | `optimizer.step()` adds small change to weights | AdamW |\n", "\n", "One step touches the **same** weight matrix thousands of times indirectly (via every position in the batch), then applies one averaged update.\n", "\n", "Run **8.1** next (creates `optimizer`), then the cell below watches one weight change.\n" ] }, { "cell_type": "markdown", "id": "eb72fa20", "metadata": {}, "source": [ "### 8.0d Backward (`loss.backward`) — why it exists in the chain\n", "\n", "**Goal of backward:** figure out *how each weight should move* to make loss smaller.\n", "\n", "You cannot run backward until the steps below exist:\n", "\n", "| Prerequisite | Built in | Provides |\n", "|--------------|----------|----------|\n", "| **Step 1** `xb`, `yb` | sections 4-5 | input + answer key |\n", "| **Step 2** `logits`, `loss` | section 6 forward | one scalar \"how wrong\" |\n", "| **Step 3** `loss.backward()` | **this step** | `.grad` on each weight |\n", "| **Step 4** `optimizer.step()` | section 8.1+ | weights actually change |\n", "\n", "**Analogy:** loss = \"how off-key is the song?\" Backward = \"for each knob,\n", "which direction fixes pitch?\" Step = \"turn the knobs slightly.\"\n", "\n", "Without step 2 (loss), step 3 has nothing to differentiate — PyTorch needs a\n", "single number that depends on the weights.\n", "\n", "**-> Training:** backward is the bridge between \"we were wrong\" (loss) and\n", "\"here is how to fix each weight\" (gradients).\n" ] }, { "cell_type": "code", "execution_count": null, "id": "019a14ab", "metadata": {}, "outputs": [], "source": [ "# trace the chain on one batch (steps 1 -> 2 -> 3)\n", "xb, yb = get_batch(\"train\")\n", "print(\"STEP 1 prerequisite: batch shapes\", xb.shape, yb.shape)\n", "\n", "model.zero_grad(set_to_none=True)\n", "logits, loss = model(xb, yb)\n", "print(\"STEP 2: forward produced loss =\", round(loss.item(), 4), \"(scalar)\")\n", "\n", "loss.backward() # STEP 3\n", "w = model.token_embedding_table.weight\n", "print(\"STEP 3: backward filled w.grad with shape\", w.grad.shape)\n", "print(\" example grad[t,h]:\", w.grad[stoi[\"t\"], stoi[\"h\"]].item())\n", "print(\"STEP 4: run optimizer.step() to add grad into weights (see 8.1)\")\n" ] }, { "cell_type": "markdown", "id": "843b881d", "metadata": {}, "source": [ "**Your turn 8d** — Run the cell above, then answer:\n", "1. What shape is `w.grad`? (Same as weights — one partial derivative per entry.)\n", "2. Why must `loss` be a single number, not a vector?\n", "3. What section would break if `yb` were missing?\n", "\n", "*(2: backward needs one objective to differentiate. 3: section 6 loss — no answer key.)*\n" ] }, { "cell_type": "markdown", "id": "80359cd3", "metadata": {}, "source": [ "### 8.1 One training step\n", "\n", "Run this once to see loss change after a single update.\n", "\n", "**→ Training:** Minimal version of the full loop. If loss changes after `step()`, weights moved — training is working.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8949b082", "metadata": {}, "outputs": [], "source": [ "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)\n", "xb, yb = get_batch(\"train\")\n", "_, loss_before = model(xb, yb)\n", "optimizer.zero_grad(set_to_none=True)\n", "_, loss = model(xb, yb)\n", "loss.backward()\n", "optimizer.step()\n", "_, loss_after = model(xb, yb)\n", "print(f\"before: {loss_before.item():.4f} after: {loss_after.item():.4f}\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8d2cc5da", "metadata": {}, "outputs": [], "source": [ "# watch ONE matrix entry change (run after 8.1 created optimizer)\n", "w = model.token_embedding_table.weight\n", "before = w[stoi['t'], stoi['h']].item()\n", "\n", "xb, yb = get_batch(\"train\")\n", "_, loss = model(xb, yb)\n", "optimizer.zero_grad(set_to_none=True)\n", "loss.backward()\n", "optimizer.step()\n", "\n", "after = w[stoi['t'], stoi['h']].item()\n", "print(f\"loss: {loss.item():.4f}\")\n", "print(f\"weight score for t->h: {before:.6f} -> {after:.6f}\")\n", "print(\"changed:\", before != after)\n" ] }, { "cell_type": "markdown", "id": "4b41e5fc", "metadata": {}, "source": [ "### 8.0c Where the result of training is stored\n", "\n", "| Location | What | Persists? |\n", "|----------|------|-----------|\n", "| `model.token_embedding_table.weight` | learned matrix in RAM/VRAM | until kernel restarts |\n", "| `model.state_dict()` | name → tensor map of all weights | same |\n", "| disk (`torch.save`) | checkpoint file | yes — **not done by default** |\n", "\n", "Re-run the notebook from the top → new random weights → must train again unless you load a saved checkpoint.\n", "\n", "**→ Training:** After section 8.2, `state_dict()` holds the **result**. Section 9 reads those same weights to generate text.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "66fd4109", "metadata": {}, "outputs": [], "source": [ "sd = model.state_dict()\n", "print(\"keys in state_dict:\", list(sd.keys()))\n", "print(\"weight tensor shape:\", sd[\"token_embedding_table.weight\"].shape)\n", "\n", "# optional: save / load (uncomment to use across sessions)\n", "# torch.save(sd, \"bigram_shakespeare.pt\")\n", "# model.load_state_dict(torch.load(\"bigram_shakespeare.pt\", map_location=device))\n", "print(\"trained weights live in memory inside model until you torch.save them.\")\n" ] }, { "cell_type": "markdown", "id": "603bd9b0", "metadata": {}, "source": [ "### 8.2 Full training loop + val loss\n", "\n", "Every 1000 steps we average loss on train **and** val batches. Val should track\n", "train; if train drops but val does not → overfitting.\n", "\n", "**→ Training:** 10,000 updates to the embedding matrix. Printed val loss checks that weights generalise, not only memorise train slices.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "b34ee95e", "metadata": {}, "outputs": [], "source": [ "batch_size = 32\n", "\n", "@torch.no_grad()\n", "def estimate_loss(eval_iters=100):\n", " out = {}\n", " model.eval()\n", " for split in [\"train\", \"val\"]:\n", " losses = torch.zeros(eval_iters)\n", " for k in range(eval_iters):\n", " X, Y = get_batch(split)\n", " _, L = model(X, Y)\n", " losses[k] = L.item()\n", " out[split] = losses.mean().item()\n", " model.train()\n", " return out\n", "\n", "for step in range(10000):\n", " xb, yb = get_batch(\"train\")\n", " _, loss = model(xb, yb)\n", " optimizer.zero_grad(set_to_none=True)\n", " loss.backward()\n", " optimizer.step()\n", " if step % 1000 == 0:\n", " L = estimate_loss()\n", " print(f\"step {step:5d} train {L['train']:.4f} val {L['val']:.4f}\")\n" ] }, { "cell_type": "markdown", "id": "69ff0dfc", "metadata": {}, "source": [ "**Your turn 8** — After training, is val loss below the random baseline (~4.17)?\n" ] }, { "cell_type": "code", "execution_count": null, "id": "23be596a", "metadata": {}, "outputs": [], "source": [ "L = estimate_loss(50)\n", "baseline = -torch.log(torch.tensor(1.0 / vocab_size)).item()\n", "print(f\"train {L['train']:.4f} val {L['val']:.4f} baseline {baseline:.4f}\")\n" ] }, { "cell_type": "markdown", "id": "1a49aef3", "metadata": {}, "source": [ "## 9. Generation after training + what's next\n", "\n", "\n", "**You are here (section 9):** see the **goal** achieved (partially — bigram is limited).\n", "\n", "**Needs:** section 8 (trained weights in `model`).\n", "**Unlocks:** next notebook — self-attention for longer context.\n", "Same `generate` as section 7 — but weights learned letter/space patterns.\n", "Still **not** real Shakespeare: bigram only sees **one** character back.\n", "\n", "Next notebook: **self-attention** so each token can use the full prefix.\n", "\n", "**Whole notebook checklist**\n", "- [ ] Text → tokens → 1-D tensor tape\n", "- [ ] `x`/`y` shift, prefix trick, `(B,T)` batching\n", "- [ ] Logits, softmax, cross-entropy\n", "- [ ] Train loop + val loss\n", "- [ ] Generate before vs after\n", "\n", "**→ Training:** Proof that weights changed. Better letter patterns = the matrix learned co-occurrence stats. Weights live in `model` until you `torch.save` them.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2ef869fb", "metadata": {}, "outputs": [], "source": [ "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n", "print(decode(model.generate(context, max_new_tokens=400)[0].tolist()))\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.12" } }, "nbformat": 4, "nbformat_minor": 5 }