480 lines
20 KiB
Plaintext
480 lines
20 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "3d5936ca",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Let's build GPT — self-attention to a full transformer\n",
|
|
"\n",
|
|
"Give every token the power to **look back** — **concept -> code -> Your turn** each step.\n",
|
|
"\n",
|
|
"This is notebook **06**, the payoff. Notebook `01` built a bigram that only sees one character\n",
|
|
"back. Here we add **self-attention** so each position can gather information from all earlier\n",
|
|
"positions, then stack it into a real (small) **GPT** and train it on tiny Shakespeare.\n",
|
|
"Follows Karpathy's *\"Let's build GPT from scratch\"*."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "8b6c4fb3",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Prologue — in plain English\n",
|
|
"\n",
|
|
"The bigram's weakness: predicting the next letter from only the previous one. WaveNet (`05`)\n",
|
|
"widened the view but fused neighbours in a fixed pattern. **Attention** is smarter: every\n",
|
|
"token looks at all earlier tokens and decides, on its own, *which* ones matter right now.\n",
|
|
"\n",
|
|
"Real-life picture: you are reading a sentence and hit the word \"it.\" To know what \"it\" refers\n",
|
|
"to, your eyes jump back to the relevant earlier noun — not to a fixed position, but to whichever\n",
|
|
"word is relevant. Self-attention lets each token make that kind of targeted look-back.\n",
|
|
"\n",
|
|
"We build it up in small pieces: the averaging trick, then queries/keys/values, then multiple\n",
|
|
"heads, then the full transformer block (attention + feed-forward + residual + LayerNorm)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "028dce18",
|
|
"metadata": {},
|
|
"source": [
|
|
"### 0.1 Setup and data (standalone)\n",
|
|
"\n",
|
|
"Self-contained: load `input.txt`, build the character vocabulary, and make `get_batch`. Uses\n",
|
|
"GPU if available, else CPU. The model is intentionally **small** so it trains on a laptop CPU;\n",
|
|
"the comments show how to scale it up if you have a GPU."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "36f4f6e6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\")\n",
|
|
"\n",
|
|
"import torch\n",
|
|
"import torch.nn as nn\n",
|
|
"from torch.nn import functional as F\n",
|
|
"\n",
|
|
"torch.manual_seed(1337)\n",
|
|
"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
|
|
"print(\"device:\", device, \"| torch\", torch.__version__)\n",
|
|
"\n",
|
|
"# --- hyperparameters (small for CPU; bump these up on a GPU) ---\n",
|
|
"batch_size = 16\n",
|
|
"block_size = 32 # max context length\n",
|
|
"n_embd = 64\n",
|
|
"n_head = 4\n",
|
|
"n_layer = 3\n",
|
|
"dropout = 0.0\n",
|
|
"learning_rate = 1e-3\n",
|
|
"max_iters = 3000\n",
|
|
"eval_interval = 1000\n",
|
|
"eval_iters = 50\n",
|
|
"\n",
|
|
"# --- data ---\n",
|
|
"text = open(\"input.txt\", \"r\", encoding=\"utf-8\").read()\n",
|
|
"chars = sorted(list(set(text)))\n",
|
|
"vocab_size = len(chars)\n",
|
|
"stoi = {ch: i for i, ch in enumerate(chars)}\n",
|
|
"itos = {i: ch for i, ch in enumerate(chars)}\n",
|
|
"encode = lambda s: [stoi[c] for c in s]\n",
|
|
"decode = lambda l: \"\".join(itos[i] for i in l)\n",
|
|
"\n",
|
|
"data = torch.tensor(encode(text), dtype=torch.long)\n",
|
|
"n = int(0.9 * len(data))\n",
|
|
"train_data, val_data = data[:n], data[n:]\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",
|
|
"print(\"vocab_size:\", vocab_size, \"| chars in dataset:\", len(text))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "d2b12520",
|
|
"metadata": {},
|
|
"source": [
|
|
"### 1.1 The math trick: a token can average its past\n",
|
|
"\n",
|
|
"Before real attention, here is the key move in its simplest form: let each position become the\n",
|
|
"**average of itself and all positions before it** (never the future — that would be cheating).\n",
|
|
"We show three ways that give the *same* result, from slow to elegant."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "54c69615",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"torch.manual_seed(1337)\n",
|
|
"B, T, C = 4, 8, 2\n",
|
|
"x = torch.randn(B, T, C)\n",
|
|
"\n",
|
|
"# v1: explicit loop — average of all previous tokens\n",
|
|
"xbow = torch.zeros((B, T, C))\n",
|
|
"for b in range(B):\n",
|
|
" for t in range(T):\n",
|
|
" xbow[b, t] = x[b, :t + 1].mean(0)\n",
|
|
"\n",
|
|
"# v2: same thing as a matrix multiply with a lower-triangular average matrix\n",
|
|
"wei = torch.tril(torch.ones(T, T))\n",
|
|
"wei = wei / wei.sum(1, keepdim=True)\n",
|
|
"xbow2 = wei @ x\n",
|
|
"\n",
|
|
"# v3: same again, but the weights come from softmax over a masked matrix\n",
|
|
"tril = torch.tril(torch.ones(T, T))\n",
|
|
"wei = torch.zeros((T, T)).masked_fill(tril == 0, float(\"-inf\"))\n",
|
|
"wei = F.softmax(wei, dim=-1)\n",
|
|
"xbow3 = wei @ x\n",
|
|
"\n",
|
|
"print(\"v1 == v2 :\", torch.allclose(xbow, xbow2, atol=1e-6), \" (max diff %.1e - just float rounding)\" % (xbow - xbow2).abs().max())\n",
|
|
"print(\"v1 == v3 :\", torch.allclose(xbow, xbow3, atol=1e-6), \" (max diff %.1e - just float rounding)\" % (xbow - xbow3).abs().max())\n",
|
|
"print(\"\\nthe softmax weight matrix (row t can only attend to columns <= t):\")\n",
|
|
"print(wei)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "aa46672a",
|
|
"metadata": {},
|
|
"source": [
|
|
"Why the `-inf` + softmax version matters: those weights do **not** have to be a flat\n",
|
|
"average. If the numbers in `wei` are *learned*, each token can decide how much to listen to each\n",
|
|
"earlier token. That is exactly what attention does next.\n",
|
|
"\n",
|
|
"Analogy: a meeting where you may only hear people who already spoke (the mask), and you choose\n",
|
|
"how much weight to give each speaker (the learned weights)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "b522ef98",
|
|
"metadata": {},
|
|
"source": [
|
|
"### 2.1 A single self-attention head — query, key, value\n",
|
|
"\n",
|
|
"Each token emits three vectors:\n",
|
|
"\n",
|
|
"- **query**: \"what am I looking for?\"\n",
|
|
"- **key**: \"what do I offer?\"\n",
|
|
"- **value**: \"what will I actually pass on if chosen?\"\n",
|
|
"\n",
|
|
"A token's attention to another = how well its **query** matches the other's **key**\n",
|
|
"(dot product). We scale by `1/sqrt(head_size)` to keep the numbers tame, mask the future, softmax\n",
|
|
"into weights, then take the weighted sum of **values**."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9adfb9ea",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"torch.manual_seed(1337)\n",
|
|
"B, T, C = 4, 8, 32\n",
|
|
"x = torch.randn(B, T, C)\n",
|
|
"head_size = 16\n",
|
|
"\n",
|
|
"key = nn.Linear(C, head_size, bias=False)\n",
|
|
"query = nn.Linear(C, head_size, bias=False)\n",
|
|
"value = nn.Linear(C, head_size, bias=False)\n",
|
|
"\n",
|
|
"k = key(x) # (B, T, head_size)\n",
|
|
"q = query(x) # (B, T, head_size)\n",
|
|
"wei = q @ k.transpose(-2, -1) * head_size ** -0.5 # (B, T, T) affinities\n",
|
|
"\n",
|
|
"tril = torch.tril(torch.ones(T, T))\n",
|
|
"wei = wei.masked_fill(tril == 0, float(\"-inf\"))\n",
|
|
"wei = F.softmax(wei, dim=-1)\n",
|
|
"\n",
|
|
"v = value(x)\n",
|
|
"out = wei @ v # (B, T, head_size)\n",
|
|
"print(\"attention output shape:\", tuple(out.shape))\n",
|
|
"print(\"\\nrow 0 of the learned attention weights (note they are NOT uniform):\")\n",
|
|
"print(wei[0, -1])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ff4eea4c",
|
|
"metadata": {},
|
|
"source": [
|
|
"### 3.1 The full GPT, assembled from clean modules\n",
|
|
"\n",
|
|
"Now we package it. The pieces:\n",
|
|
"\n",
|
|
"- **Head** — one self-attention head (above), with dropout and a causal mask buffer.\n",
|
|
"- **MultiHeadAttention** — several heads in parallel, concatenated, then projected. Analogy:\n",
|
|
" several readers each look back for a different kind of clue, then pool notes.\n",
|
|
"- **FeedForward** — a small per-token MLP: \"time to think about what I gathered.\"\n",
|
|
"- **Block** — attention + feed-forward, each wrapped with a **residual** connection and\n",
|
|
" **LayerNorm**. Residual = keep the original and add the new insight (so deep stacks still\n",
|
|
" train). LayerNorm = re-level each token's vector (the `03` lesson, per-token).\n",
|
|
"- **GPTLanguageModel** — token + position embeddings -> a stack of Blocks -> final LayerNorm ->\n",
|
|
" a linear head to vocab logits."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b0a4b08d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"class Head(nn.Module):\n",
|
|
" def __init__(self, head_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",
|
|
" self.dropout = nn.Dropout(dropout)\n",
|
|
"\n",
|
|
" def forward(self, x):\n",
|
|
" B, T, C = x.shape\n",
|
|
" k = self.key(x)\n",
|
|
" q = self.query(x)\n",
|
|
" wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5\n",
|
|
" wei = wei.masked_fill(self.tril[:T, :T] == 0, float(\"-inf\"))\n",
|
|
" wei = F.softmax(wei, dim=-1)\n",
|
|
" wei = self.dropout(wei)\n",
|
|
" v = self.value(x)\n",
|
|
" return wei @ v\n",
|
|
"\n",
|
|
"\n",
|
|
"class MultiHeadAttention(nn.Module):\n",
|
|
" def __init__(self, num_heads, head_size):\n",
|
|
" super().__init__()\n",
|
|
" self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
|
|
" self.proj = nn.Linear(head_size * num_heads, n_embd)\n",
|
|
" self.dropout = nn.Dropout(dropout)\n",
|
|
"\n",
|
|
" def forward(self, x):\n",
|
|
" out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
|
|
" return self.dropout(self.proj(out))\n",
|
|
"\n",
|
|
"\n",
|
|
"class FeedForward(nn.Module):\n",
|
|
" def __init__(self, n_embd):\n",
|
|
" super().__init__()\n",
|
|
" self.net = nn.Sequential(\n",
|
|
" nn.Linear(n_embd, 4 * n_embd),\n",
|
|
" nn.ReLU(),\n",
|
|
" nn.Linear(4 * n_embd, n_embd),\n",
|
|
" nn.Dropout(dropout),\n",
|
|
" )\n",
|
|
"\n",
|
|
" def forward(self, x):\n",
|
|
" return self.net(x)\n",
|
|
"\n",
|
|
"\n",
|
|
"class Block(nn.Module):\n",
|
|
" def __init__(self, n_embd, n_head):\n",
|
|
" super().__init__()\n",
|
|
" head_size = n_embd // n_head\n",
|
|
" self.sa = MultiHeadAttention(n_head, head_size)\n",
|
|
" self.ffwd = FeedForward(n_embd)\n",
|
|
" self.ln1 = nn.LayerNorm(n_embd)\n",
|
|
" self.ln2 = nn.LayerNorm(n_embd)\n",
|
|
"\n",
|
|
" def forward(self, x):\n",
|
|
" x = x + self.sa(self.ln1(x)) # residual around attention\n",
|
|
" x = x + self.ffwd(self.ln2(x)) # residual around feed-forward\n",
|
|
" return x\n",
|
|
"\n",
|
|
"\n",
|
|
"class GPTLanguageModel(nn.Module):\n",
|
|
" def __init__(self):\n",
|
|
" super().__init__()\n",
|
|
" self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
|
|
" self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
|
|
" self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])\n",
|
|
" self.ln_f = nn.LayerNorm(n_embd)\n",
|
|
" self.lm_head = nn.Linear(n_embd, vocab_size)\n",
|
|
"\n",
|
|
" def forward(self, idx, targets=None):\n",
|
|
" B, T = idx.shape\n",
|
|
" tok_emb = self.token_embedding_table(idx) # (B,T,C)\n",
|
|
" pos_emb = self.position_embedding_table(torch.arange(T, device=device))\n",
|
|
" x = tok_emb + pos_emb\n",
|
|
" x = self.blocks(x)\n",
|
|
" x = self.ln_f(x)\n",
|
|
" logits = self.lm_head(x) # (B,T,vocab)\n",
|
|
" if targets is None:\n",
|
|
" return logits, None\n",
|
|
" B, T, Cc = logits.shape\n",
|
|
" loss = F.cross_entropy(logits.view(B * T, Cc), 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",
|
|
" idx_cond = idx[:, -block_size:] # never exceed the context window\n",
|
|
" logits, _ = self(idx_cond)\n",
|
|
" logits = logits[:, -1, :] # last position only\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",
|
|
"\n",
|
|
"model = GPTLanguageModel().to(device)\n",
|
|
"print(\"parameters (millions):\", sum(p.numel() for p in model.parameters()) / 1e6)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "dbb781c5",
|
|
"metadata": {},
|
|
"source": [
|
|
"### 3.2 Positional embeddings — why we need them\n",
|
|
"\n",
|
|
"Attention by itself is **order-blind**: it treats the context as a bag of tokens. But \"dog\n",
|
|
"bites man\" and \"man bites dog\" must differ. So besides each token's meaning\n",
|
|
"(`token_embedding_table`), we add a vector for its **position**\n",
|
|
"(`position_embedding_table`). Now a token knows both *what* it is and *where* it is.\n",
|
|
"\n",
|
|
"(That addition already happens inside `forward`: `x = tok_emb + pos_emb`.)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0385a0ef",
|
|
"metadata": {},
|
|
"source": [
|
|
"### 4.1 Train the GPT\n",
|
|
"\n",
|
|
"Same training shape as always, with the AdamW optimizer. Watch train and val loss fall well\n",
|
|
"below the bigram's ~2.5. (Small model + CPU, so this is a few minutes; scale up on a GPU for\n",
|
|
"real Shakespeare.)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "395e6922",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"@torch.no_grad()\n",
|
|
"def estimate_loss():\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",
|
|
" _, loss = model(X, Y)\n",
|
|
" losses[k] = loss.item()\n",
|
|
" out[split] = losses.mean().item()\n",
|
|
" model.train()\n",
|
|
" return out\n",
|
|
"\n",
|
|
"optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
|
|
"\n",
|
|
"for it in range(max_iters):\n",
|
|
" if it % eval_interval == 0 or it == max_iters - 1:\n",
|
|
" losses = estimate_loss()\n",
|
|
" print(f\"step {it:4d}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\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",
|
|
"print(\"done training\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0ae5a608",
|
|
"metadata": {},
|
|
"source": [
|
|
"### 5.1 Generate Shakespeare-ish text\n",
|
|
"\n",
|
|
"Start from a single newline and let the model write. It will not be perfect Shakespeare (tiny\n",
|
|
"model, short training), but compared to the bigram it now produces real-looking words, spacing,\n",
|
|
"and line breaks — because each character can attend to the ones before it."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "61db808c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
|
|
"print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "6b43261f",
|
|
"metadata": {},
|
|
"source": [
|
|
"**Your turn 6** — If you have a GPU, scale up (e.g. `n_embd=384`, `n_head=6`,\n",
|
|
"`n_layer=6`, `block_size=256`, `dropout=0.2`, `max_iters=5000`) and retrain — the samples get\n",
|
|
"dramatically better. On CPU, try just raising `n_layer` to 4 or `max_iters` to 5000 and compare\n",
|
|
"the val loss."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "1e1142ab",
|
|
"metadata": {},
|
|
"source": [
|
|
"## What's next — and the connection to ChatGPT\n",
|
|
"\n",
|
|
"You built a real transformer: token + position embeddings, multi-head causal self-attention,\n",
|
|
"feed-forward layers, residuals, and LayerNorm — the exact architecture behind GPT.\n",
|
|
"\n",
|
|
"What you trained here is a **pretrained** model in miniature: it only learns to *continue text*\n",
|
|
"(predict the next token) from a big pile of text. ChatGPT starts the same way, then adds steps\n",
|
|
"this series does not cover:\n",
|
|
"\n",
|
|
"- **Supervised finetuning** on example conversations (teach it to answer, not just continue).\n",
|
|
"- **Reinforcement learning from human feedback (RLHF)** to align responses with what people\n",
|
|
" prefer.\n",
|
|
"\n",
|
|
"The remaining notebook:\n",
|
|
"\n",
|
|
"- `07_gpt_tokenizer.ipynb` — real models do not use one-character-per-token; they use **Byte\n",
|
|
" Pair Encoding**. That is the final piece of the pipeline.\n",
|
|
"\n",
|
|
"**Checklist**\n",
|
|
"- [ ] The averaging trick (loop = tril matmul = masked softmax)\n",
|
|
"- [ ] Query / key / value self-attention head\n",
|
|
"- [ ] Multi-head attention + projection\n",
|
|
"- [ ] Feed-forward, residuals, LayerNorm -> Block\n",
|
|
"- [ ] Token + position embeddings -> full GPT\n",
|
|
"- [ ] Trained below the bigram baseline; generated text"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv (3.12.10)",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"name": "python",
|
|
"version": "3.12.10"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
} |