From 6b7927c16cbe89082552ccbaaa2f432a87c6702a Mon Sep 17 00:00:00 2001 From: jzzliu Date: Sat, 7 Nov 2026 20:19:10 -0700 Subject: [PATCH] Add microgpt: from-scratch GPT with RoPE, KV cache, sampling, tests and benchmark --- .gitignore | 9 ++ README.md | 93 ++++++++++++++++- pyproject.toml | 33 ++++++ scripts/benchmark.py | 44 ++++++++ src/microgpt/__init__.py | 8 ++ src/microgpt/data.py | 37 +++++++ src/microgpt/generate.py | 49 +++++++++ src/microgpt/model.py | 220 +++++++++++++++++++++++++++++++++++++++ src/microgpt/sampling.py | 49 +++++++++ src/microgpt/train.py | 101 ++++++++++++++++++ tests/test_model.py | 124 ++++++++++++++++++++++ tests/test_sampling.py | 60 +++++++++++ 12 files changed, 826 insertions(+), 1 deletion(-) create mode 100644 .gitignore create mode 100644 pyproject.toml create mode 100644 scripts/benchmark.py create mode 100644 src/microgpt/__init__.py create mode 100644 src/microgpt/data.py create mode 100644 src/microgpt/generate.py create mode 100644 src/microgpt/model.py create mode 100644 src/microgpt/sampling.py create mode 100644 src/microgpt/train.py create mode 100644 tests/test_model.py create mode 100644 tests/test_sampling.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3d33612 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +.venv/ +__pycache__/ +*.egg-info/ +.pytest_cache/ +.ruff_cache/ +data/ +checkpoints/ +*.pt +.DS_Store diff --git a/README.md b/README.md index a5daf24..2427c49 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,93 @@ -# python-project +# microgpt +A GPT-style language model written from scratch in PyTorch, with the pieces that matter in +production LLM systems: rotary position embeddings, RMSNorm, SwiGLU, an incremental-decoding +**KV cache**, nucleus sampling, and a test suite that proves each piece is correct. + +No `nn.Transformer`, no Hugging Face. Attention, masking, caching and sampling are all +written out explicitly (~250 lines of model code) so every line can be reasoned about. + +## What's inside + +| Piece | Where | Why it's here | +|---|---|---| +| Causal self-attention with explicit mask | `model.py` | Mask is offset by the cache length, which is the classic place incremental decoding goes wrong | +| RoPE | `model.py` | Relative position without a learned table; cached keys keep their absolute rotation | +| RMSNorm + SwiGLU, tied embeddings | `model.py` | The Llama-style recipe at toy scale | +| KV cache + sliding window | `model.py` | `O(T)` work per generated token instead of `O(T^2)` | +| Temperature / top-k / top-p | `sampling.py` | Written from first principles, tested against analytic distributions | +| Trainer | `train.py` | AdamW with decoupled weight decay (no decay on norms), warmup + cosine LR, grad clipping | +| Benchmark | `scripts/benchmark.py` | Measures the cache speedup instead of asserting it | + +## Quickstart + +```bash +python3 -m venv .venv && source .venv/bin/activate +pip install -e ".[dev]" +pytest # 21 tests, ~3s on CPU + +mkdir -p data && curl -sSL -o data/shakespeare.txt \ + https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt +python -m microgpt.train --data data/shakespeare.txt --steps 1200 --batch-size 24 +python -m microgpt.generate --prompt "ROMEO:" --tokens 200 --top-k 40 +python scripts/benchmark.py +``` + +## Correctness: what the tests actually prove + +* **No future leakage** — perturbing tokens at positions >= 8 leaves logits at positions < 8 unchanged. +* **KV cache == full forward** — prefill then one-token-at-a-time decoding reproduces the full + forward pass to `1e-5`; so does multi-token continuation (which exercises the offset causal mask). +* **Greedy generation is identical with and without the cache** (while within the context window). +* **RoPE** preserves vector norms, and `` depends only on `m - n`. +* **Sampling** — top-k and top-p keep exactly the intended tokens; empirical frequencies match + the softmax distribution; seeded runs are reproducible. +* **Learning** — the model can overfit a tiny batch (loss drops by more than 70%). + +## Results + +Measured on an Intel iMac (6 CPU threads, no GPU). Numbers are from this repo's own scripts. + +**Training** — 0.81M parameters, 4 layers, 128-dim, character-level Tiny Shakespeare, +1200 steps (~5.6 min on CPU): + +| step | train loss | val loss | +|---:|---:|---:| +| 0 | 4.242 | 4.235 | +| 400 | 1.657 | 1.830 | +| 800 | 1.444 | 1.655 | +| 1200 | 1.348 | 1.596 | + +(Initial loss 4.24 ≈ ln 65, i.e. uniform over the 65-character vocabulary, as expected.) + +**KV-cache decoding** — 2.68M-parameter model, greedy, 16-token prompt: + +| new tokens | no cache (tok/s) | KV cache (tok/s) | speedup | identical output | +|-----------:|-----------------:|-----------------:|--------:|:-----------------:| +| 32 | 158.3 | 194.2 | 1.2x | yes | +| 64 | 166.2 | 259.3 | 1.6x | yes | +| 128 | 98.2 | 174.4 | 1.8x | yes | +| 256 | 74.5 | 198.8 | 2.7x | yes | +| 480 | 34.7 | 167.4 | 4.8x | yes | + +Cached throughput stays roughly flat as the sequence grows while uncached throughput decays, +which is the `O(T)` vs `O(T^2)` difference showing up in practice. Timings on a shared CPU are +noisy, so treat the small-length rows as approximate. + +## Design notes and known limitations + +* **Sliding window past `block_size`.** The uncached path always sees the last `block_size` + tokens. The cached path drops its cache when full and rebuilds it from the latest half-window. + Outputs are therefore identical only while the sequence fits in the window; beyond that they + differ by design. (RoPE makes true cache eviction possible, but keys would need re-rotation.) +* **Attention is `O(T^2)` memory** per layer: no FlashAttention, no grouped-query attention. + Both are natural next steps and would also shrink the KV cache. +* **Character-level, small scale.** The point is correctness and clarity, not sample quality. + Samples are Shakespeare-flavoured but not coherent. + +## Roadmap + +- [ ] Grouped-query attention (smaller KV cache) with a parity test against full multi-head +- [ ] Batched generation with per-sequence cache lengths +- [ ] Speculative decoding with a draft model, verified to preserve the target distribution +- [ ] BPE tokenizer diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..bf3cd40 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,33 @@ +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[project] +name = "microgpt" +version = "0.1.0" +description = "A from-scratch GPT in PyTorch: RoPE, RMSNorm, SwiGLU, KV-cache inference, and the tests that prove it is correct." +readme = "README.md" +requires-python = ">=3.9" +dependencies = ["torch>=2.2", "numpy"] + +[project.optional-dependencies] +dev = ["pytest", "ruff"] + +[project.scripts] +microgpt-train = "microgpt.train:main" +microgpt-generate = "microgpt.generate:main" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] + +[tool.ruff] +line-length = 100 +target-version = "py39" + +[tool.ruff.lint] +select = ["E", "F", "I", "B", "UP"] +ignore = ["UP006", "UP007", "UP035", "UP045"] diff --git a/scripts/benchmark.py b/scripts/benchmark.py new file mode 100644 index 0000000..2c41068 --- /dev/null +++ b/scripts/benchmark.py @@ -0,0 +1,44 @@ +"""Measure decoding throughput with and without the KV cache. + + python scripts/benchmark.py # prints a markdown table +""" + +from __future__ import annotations + +import time + +import torch + +from microgpt.model import GPT, GPTConfig + + +def time_generate(model: GPT, prompt: torch.Tensor, n: int, use_cache: bool) -> float: + t0 = time.perf_counter() + model.generate(prompt, n, temperature=0, use_cache=use_cache) + return time.perf_counter() - t0 + + +def main() -> None: + torch.manual_seed(0) + torch.set_num_threads(max(1, torch.get_num_threads())) + cfg = GPTConfig(vocab_size=96, block_size=512, n_layer=6, n_head=6, n_embd=192) + model = GPT(cfg).eval() + prompt = torch.randint(0, cfg.vocab_size, (1, 16)) + print(f"model: {model.num_params() / 1e6:.2f}M params, {torch.get_num_threads()} CPU threads\n") + for use_cache in (False, True): # warm up allocator / thread pool before timing + time_generate(model, prompt, 32, use_cache) + print("| new tokens | no cache (tok/s) | KV cache (tok/s) | speedup | identical output |") + print("|-----------:|-----------------:|-----------------:|--------:|:-----------------:|") + for n in (32, 64, 128, 256, 480): + slow = time_generate(model, prompt, n, use_cache=False) + fast = time_generate(model, prompt, n, use_cache=True) + same = torch.equal( + model.generate(prompt, n, temperature=0, use_cache=False), + model.generate(prompt, n, temperature=0, use_cache=True), + ) + verdict = "yes" if same else "NO" + print(f"| {n} | {n / slow:.1f} | {n / fast:.1f} | {slow / fast:.1f}x | {verdict} |") + + +if __name__ == "__main__": + main() diff --git a/src/microgpt/__init__.py b/src/microgpt/__init__.py new file mode 100644 index 0000000..88d78be --- /dev/null +++ b/src/microgpt/__init__.py @@ -0,0 +1,8 @@ +"""microgpt: a from-scratch GPT with KV-cache inference.""" + +from microgpt.data import CharTokenizer +from microgpt.model import GPT, GPTConfig +from microgpt.sampling import sample_next + +__all__ = ["GPT", "GPTConfig", "CharTokenizer", "sample_next"] +__version__ = "0.1.0" diff --git a/src/microgpt/data.py b/src/microgpt/data.py new file mode 100644 index 0000000..27a1b85 --- /dev/null +++ b/src/microgpt/data.py @@ -0,0 +1,37 @@ +"""Character-level tokenizer and batching.""" + +from __future__ import annotations + +from typing import List, Sequence, Tuple + +import torch + + +class CharTokenizer: + def __init__(self, chars: Sequence[str]) -> None: + self.chars: List[str] = list(chars) + self._stoi = {c: i for i, c in enumerate(self.chars)} + + @classmethod + def from_text(cls, text: str) -> CharTokenizer: + return cls(sorted(set(text))) + + @property + def vocab_size(self) -> int: + return len(self.chars) + + def encode(self, text: str) -> List[int]: + return [self._stoi[c] for c in text] + + def decode(self, ids: Sequence[int]) -> str: + return "".join(self.chars[i] for i in ids) + + +def get_batch( + data: torch.Tensor, batch_size: int, block_size: int, generator: torch.Generator | None = None +) -> Tuple[torch.Tensor, torch.Tensor]: + """Random (x, y) windows where y is x shifted left by one token.""" + starts = torch.randint(len(data) - block_size - 1, (batch_size,), generator=generator) + x = torch.stack([data[s : s + block_size] for s in starts]) + y = torch.stack([data[s + 1 : s + 1 + block_size] for s in starts]) + return x, y diff --git a/src/microgpt/generate.py b/src/microgpt/generate.py new file mode 100644 index 0000000..c3a1441 --- /dev/null +++ b/src/microgpt/generate.py @@ -0,0 +1,49 @@ +"""Sample text from a trained checkpoint.""" + +from __future__ import annotations + +import argparse + +import torch + +from microgpt.data import CharTokenizer +from microgpt.model import GPT, GPTConfig + + +def load_checkpoint(path: str): + ckpt = torch.load(path, map_location="cpu") + model = GPT(GPTConfig(**ckpt["config"])) + model.load_state_dict(ckpt["model"]) + model.eval() + return model, CharTokenizer(ckpt["chars"]) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--ckpt", default="checkpoints/model.pt") + p.add_argument("--prompt", default="\n") + p.add_argument("--tokens", type=int, default=300) + p.add_argument("--temperature", type=float, default=0.8) + p.add_argument("--top-k", type=int, default=None) + p.add_argument("--top-p", type=float, default=None) + p.add_argument("--no-cache", action="store_true", help="disable the KV cache") + p.add_argument("--seed", type=int, default=0) + args = p.parse_args() + + model, tok = load_checkpoint(args.ckpt) + g = torch.Generator().manual_seed(args.seed) + idx = torch.tensor([tok.encode(args.prompt)], dtype=torch.long) + out = model.generate( + idx, + args.tokens, + temperature=args.temperature, + top_k=args.top_k, + top_p=args.top_p, + use_cache=not args.no_cache, + generator=g, + ) + print(tok.decode(out[0].tolist())) + + +if __name__ == "__main__": + main() diff --git a/src/microgpt/model.py b/src/microgpt/model.py new file mode 100644 index 0000000..4041ff8 --- /dev/null +++ b/src/microgpt/model.py @@ -0,0 +1,220 @@ +"""Decoder-only transformer written from scratch. + +Architecture (the Llama recipe, at toy scale): + * RMSNorm pre-normalisation + * Rotary position embeddings (RoPE) applied to queries and keys + * Causal multi-head self-attention, written out explicitly (no nn.Transformer) + * SwiGLU feed-forward network + * Tied input/output embeddings + * Per-layer key/value cache for incremental decoding +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from microgpt.sampling import sample_next + +KVCache = List[Tuple[torch.Tensor, torch.Tensor]] + + +@dataclass +class GPTConfig: + vocab_size: int = 256 + block_size: int = 256 + n_layer: int = 4 + n_head: int = 4 + n_embd: int = 128 + rope_base: float = 10000.0 + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight + + +def rope_tables(head_dim: int, max_len: int, base: float) -> Tuple[torch.Tensor, torch.Tensor]: + """cos/sin tables of shape (max_len, head_dim // 2).""" + inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) + angles = torch.outer(torch.arange(max_len).float(), inv_freq) + return angles.cos(), angles.sin() + + +def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Rotate consecutive channel pairs of x (B, H, T, hd) by position-dependent angles. + + Because rotations compose, depends only on m - n, which is + how relative position enters attention without any learned position table. + """ + x1, x2 = x[..., 0::2], x[..., 1::2] + out1 = x1 * cos - x2 * sin + out2 = x1 * sin + x2 * cos + return torch.stack((out1, out2), dim=-1).flatten(-2) + + +class CausalSelfAttention(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + assert cfg.n_embd % cfg.n_head == 0, "n_embd must be divisible by n_head" + self.n_head = cfg.n_head + self.head_dim = cfg.n_embd // cfg.n_head + assert self.head_dim % 2 == 0, "head_dim must be even for RoPE" + self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False) + self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False) + cos, sin = rope_tables(self.head_dim, cfg.block_size, cfg.rope_base) + self.register_buffer("cos", cos, persistent=False) + self.register_buffer("sin", sin, persistent=False) + + def forward( + self, x: torch.Tensor, cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + B, T, C = x.shape + past = 0 if cache is None else cache[0].size(2) + + q, k, v = self.qkv(x).split(C, dim=2) + q, k, v = (t.view(B, T, self.n_head, self.head_dim).transpose(1, 2) for t in (q, k, v)) + + # Positions are absolute, so cached keys keep the rotation they were stored with. + cos, sin = self.cos[past : past + T], self.sin[past : past + T] + q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) + + if cache is not None: + k = torch.cat((cache[0], k), dim=2) + v = torch.cat((cache[1], v), dim=2) + new_cache = (k, v) + + S = k.size(2) # total keys = past + T + scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) + # Query i sits at absolute position past + i and may see keys j <= past + i. + allowed = torch.ones(T, S, dtype=torch.bool, device=x.device).tril(diagonal=S - T) + scores = scores.masked_fill(~allowed, float("-inf")) + out = F.softmax(scores, dim=-1) @ v + out = out.transpose(1, 2).contiguous().view(B, T, C) + return self.proj(out), new_cache + + +class SwiGLU(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + hidden = int(8 * cfg.n_embd / 3) + hidden = 32 * ((hidden + 31) // 32) + self.w_gate = nn.Linear(cfg.n_embd, hidden, bias=False) + self.w_up = nn.Linear(cfg.n_embd, hidden, bias=False) + self.w_down = nn.Linear(hidden, cfg.n_embd, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) + + +class Block(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + self.norm1 = RMSNorm(cfg.n_embd) + self.attn = CausalSelfAttention(cfg) + self.norm2 = RMSNorm(cfg.n_embd) + self.mlp = SwiGLU(cfg) + + def forward(self, x, cache=None): + a, new_cache = self.attn(self.norm1(x), cache) + x = x + a + x = x + self.mlp(self.norm2(x)) + return x, new_cache + + +class GPT(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + self.cfg = cfg + self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd) + self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer)) + self.norm_f = RMSNorm(cfg.n_embd) + self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) + self.lm_head.weight = self.tok_emb.weight # weight tying + self.apply(self._init_weights) + # GPT-2 style: shrink residual-branch outputs so variance stays flat with depth. + for name, p in self.named_parameters(): + if name.endswith(("proj.weight", "w_down.weight")): + nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, (nn.Linear, nn.Embedding)): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + + def num_params(self) -> int: + return sum(p.numel() for p in self.parameters()) # tied weights counted once + + def forward( + self, + idx: torch.Tensor, + targets: Optional[torch.Tensor] = None, + cache: Optional[KVCache] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], KVCache]: + """Return (logits, loss, new_cache). Pass `cache` to decode incrementally.""" + _, T = idx.shape + past = 0 if cache is None else cache[0][0].size(2) + if past + T > self.cfg.block_size: + raise ValueError(f"sequence length {past + T} exceeds block_size {self.cfg.block_size}") + + x = self.tok_emb(idx) + new_cache: KVCache = [] + for i, block in enumerate(self.blocks): + x, layer_cache = block(x, None if cache is None else cache[i]) + new_cache.append(layer_cache) + logits = self.lm_head(self.norm_f(x)) + + loss = None + if targets is not None: + loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) + return logits, loss, new_cache + + @torch.no_grad() + def generate( + self, + idx: torch.Tensor, + max_new_tokens: int, + temperature: float = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + use_cache: bool = True, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + """Autoregressively extend idx (B, T). Beyond block_size the window slides. + + With the cache each step feeds one token (O(T) attention work per step); without it + every step re-runs the whole window (O(T^2)). While the sequence fits in block_size + both paths give identical tokens (verified in tests). Past that point they differ by + design: the uncached path always sees the last block_size tokens, whereas the cached + path drops its cache when full and rebuilds it from the most recent half-window, so + it pays one re-prefill instead of invalidating positions on every step. + """ + was_training = self.training + self.eval() + block = self.cfg.block_size + cache: Optional[KVCache] = None + for _ in range(max_new_tokens): + if use_cache: + if cache is not None and cache[0][0].size(2) >= block: + cache = None # window full: rebuild from the tail + if cache is None: + window = idx[:, -(block // 2) :] if idx.size(1) >= block else idx + logits, _, cache = self(window, cache=None) + else: + logits, _, cache = self(idx[:, -1:], cache=cache) + else: + logits, _, _ = self(idx[:, -block:]) + next_tok = sample_next(logits[:, -1, :], temperature, top_k, top_p, generator) + idx = torch.cat((idx, next_tok), dim=1) + self.train(was_training) + return idx diff --git a/src/microgpt/sampling.py b/src/microgpt/sampling.py new file mode 100644 index 0000000..6d9d159 --- /dev/null +++ b/src/microgpt/sampling.py @@ -0,0 +1,49 @@ +"""Next-token sampling: greedy, temperature, top-k and top-p (nucleus).""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F + + +def filter_top_k(logits: torch.Tensor, k: int) -> torch.Tensor: + """Keep the k largest logits per row; set the rest to -inf.""" + k = min(k, logits.size(-1)) + kth = torch.topk(logits, k, dim=-1).values[..., -1:] + return logits.masked_fill(logits < kth, float("-inf")) + + +def filter_top_p(logits: torch.Tensor, p: float) -> torch.Tensor: + """Keep the smallest set of tokens whose probability mass reaches p (always >= 1 token).""" + sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) + probs = F.softmax(sorted_logits, dim=-1) + # Mass strictly before each token: drop a token if the ones ahead of it already cover p. + mass_before = probs.cumsum(dim=-1) - probs + sorted_logits = sorted_logits.masked_fill(mass_before >= p, float("-inf")) + return torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits) + + +def sample_next( + logits: torch.Tensor, + temperature: float = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Sample one token id per row from logits (B, V). Returns (B, 1). + + temperature == 0 means greedy argmax. Filters are applied after temperature scaling. + """ + if temperature < 0: + raise ValueError("temperature must be >= 0") + if temperature == 0: + return logits.argmax(dim=-1, keepdim=True) + logits = logits / temperature + if top_k is not None: + logits = filter_top_k(logits, top_k) + if top_p is not None: + logits = filter_top_p(logits, top_p) + probs = F.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1, generator=generator) diff --git a/src/microgpt/train.py b/src/microgpt/train.py new file mode 100644 index 0000000..c667f51 --- /dev/null +++ b/src/microgpt/train.py @@ -0,0 +1,101 @@ +"""Train a character-level GPT on a text file.""" + +from __future__ import annotations + +import argparse +import math +import time +from dataclasses import asdict +from pathlib import Path + +import torch + +from microgpt.data import CharTokenizer, get_batch +from microgpt.model import GPT, GPTConfig + + +def lr_at(step: int, max_steps: int, max_lr: float, warmup: int, min_ratio: float = 0.1) -> float: + """Linear warmup then cosine decay to min_ratio * max_lr.""" + if step < warmup: + return max_lr * (step + 1) / warmup + progress = (step - warmup) / max(1, max_steps - warmup) + return max_lr * (min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * progress))) + + +@torch.no_grad() +def estimate_loss(model: GPT, data: torch.Tensor, batch_size: int, iters: int = 20) -> float: + model.eval() + losses = [] + for _ in range(iters): + x, y = get_batch(data, batch_size, model.cfg.block_size) + losses.append(model(x, y)[1].item()) + model.train() + return sum(losses) / len(losses) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--data", required=True, help="path to a UTF-8 text file") + p.add_argument("--out", default="checkpoints/model.pt") + p.add_argument("--steps", type=int, default=1000) + p.add_argument("--batch-size", type=int, default=32) + p.add_argument("--block-size", type=int, default=128) + p.add_argument("--n-layer", type=int, default=4) + p.add_argument("--n-head", type=int, default=4) + p.add_argument("--n-embd", type=int, default=128) + p.add_argument("--lr", type=float, default=3e-3) + p.add_argument("--warmup", type=int, default=50) + p.add_argument("--eval-every", type=int, default=100) + p.add_argument("--seed", type=int, default=1337) + args = p.parse_args() + + torch.manual_seed(args.seed) + text = Path(args.data).read_text(encoding="utf-8") + tok = CharTokenizer.from_text(text) + data = torch.tensor(tok.encode(text), dtype=torch.long) + split = int(0.9 * len(data)) + train_data, val_data = data[:split], data[split:] + + cfg = GPTConfig( + vocab_size=tok.vocab_size, + block_size=args.block_size, + n_layer=args.n_layer, + n_head=args.n_head, + n_embd=args.n_embd, + ) + model = GPT(cfg) + print(f"vocab={tok.vocab_size} params={model.num_params():,} train_chars={len(train_data):,}") + + decay = [p for p in model.parameters() if p.dim() >= 2] + no_decay = [p for p in model.parameters() if p.dim() < 2] + opt = torch.optim.AdamW( + [{"params": decay, "weight_decay": 0.1}, {"params": no_decay, "weight_decay": 0.0}], + lr=args.lr, + betas=(0.9, 0.95), + ) + + t0 = time.time() + for step in range(args.steps + 1): + for group in opt.param_groups: + group["lr"] = lr_at(step, args.steps, args.lr, args.warmup) + if step % args.eval_every == 0: + tr = estimate_loss(model, train_data, args.batch_size) + va = estimate_loss(model, val_data, args.batch_size) + print(f"step {step:5d} | train {tr:.3f} | val {va:.3f} | {time.time() - t0:6.1f}s") + if step == args.steps: + break + x, y = get_batch(train_data, args.batch_size, args.block_size) + _, loss, _ = model(x, y) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + torch.save({"config": asdict(cfg), "model": model.state_dict(), "chars": tok.chars}, out) + print(f"saved {out}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..e927da1 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,124 @@ +import pytest +import torch + +from microgpt.model import GPT, GPTConfig, apply_rope, rope_tables + +CFG = GPTConfig(vocab_size=50, block_size=32, n_layer=2, n_head=4, n_embd=32) + + +@pytest.fixture() +def model(): + torch.manual_seed(0) + return GPT(CFG).eval() + + +def test_output_shapes(model): + idx = torch.randint(0, CFG.vocab_size, (3, 10)) + logits, loss, cache = model(idx) + assert logits.shape == (3, 10, CFG.vocab_size) + assert loss is None + assert len(cache) == CFG.n_layer and cache[0][0].shape == (3, CFG.n_head, 10, 8) + + +def test_weights_are_tied(model): + assert model.lm_head.weight is model.tok_emb.weight + + +def test_no_future_leakage(model): + """Changing a later token must not change logits at earlier positions.""" + a = torch.randint(0, CFG.vocab_size, (1, 12)) + b = a.clone() + b[0, 8:] = (b[0, 8:] + 1) % CFG.vocab_size + with torch.no_grad(): + la, lb = model(a)[0], model(b)[0] + assert torch.allclose(la[:, :8], lb[:, :8], atol=1e-6) + assert not torch.allclose(la[:, 8:], lb[:, 8:], atol=1e-6) + + +def test_kv_cache_matches_full_forward(model): + """Prefill + one-token-at-a-time decoding == a single full forward pass.""" + idx = torch.randint(0, CFG.vocab_size, (2, 20)) + with torch.no_grad(): + full, _, _ = model(idx) + logits, _, cache = model(idx[:, :7]) + steps = [logits] + for t in range(7, 20): + step, _, cache = model(idx[:, t : t + 1], cache=cache) + steps.append(step) + assert torch.allclose(full, torch.cat(steps, dim=1), atol=1e-5) + + +def test_cache_chunked_prefill_matches_full(model): + """Multi-token continuation exercises the non-trivial causal-mask offset.""" + idx = torch.randint(0, CFG.vocab_size, (1, 18)) + with torch.no_grad(): + full, _, _ = model(idx) + l1, _, cache = model(idx[:, :6]) + l2, _, _ = model(idx[:, 6:], cache=cache) + assert torch.allclose(full, torch.cat([l1, l2], dim=1), atol=1e-5) + + +def test_greedy_generation_identical_with_and_without_cache(model): + prompt = torch.randint(0, CFG.vocab_size, (1, 5)) + a = model.generate(prompt, 20, temperature=0, use_cache=True) + b = model.generate(prompt, 20, temperature=0, use_cache=False) + assert torch.equal(a, b) + + +def test_generation_past_block_size(model): + prompt = torch.randint(0, CFG.vocab_size, (1, 5)) + out = model.generate(prompt, CFG.block_size * 2, temperature=1.0, top_k=10) + assert out.shape == (1, 5 + CFG.block_size * 2) + assert int(out.max()) < CFG.vocab_size + + +def test_sequence_longer_than_block_size_is_rejected(model): + with pytest.raises(ValueError): + model(torch.zeros(1, CFG.block_size + 1, dtype=torch.long)) + + +def test_seeded_sampling_is_reproducible(model): + prompt = torch.randint(0, CFG.vocab_size, (1, 4)) + runs = [ + model.generate(prompt, 15, top_p=0.9, generator=torch.Generator().manual_seed(7)) + for _ in range(2) + ] + assert torch.equal(runs[0], runs[1]) + + +def test_rope_is_norm_preserving(): + cos, sin = rope_tables(16, 40, 10000.0) + x = torch.randn(2, 3, 40, 16) + y = apply_rope(x, cos, sin) + assert torch.allclose(x.norm(dim=-1), y.norm(dim=-1), atol=1e-5) + + +def test_rope_scores_depend_only_on_relative_position(): + """ must equal .""" + cos, sin = rope_tables(16, 64, 10000.0) + q, k = torch.randn(1, 1, 1, 16), torch.randn(1, 1, 1, 16) + + def score(m: int, n: int) -> torch.Tensor: + qm = apply_rope(q, cos[m : m + 1], sin[m : m + 1]) + kn = apply_rope(k, cos[n : n + 1], sin[n : n + 1]) + return (qm * kn).sum() + + assert torch.allclose(score(5, 2), score(25, 22), atol=1e-4) + assert torch.allclose(score(0, 9), score(30, 39), atol=1e-4) + assert not torch.allclose(score(5, 2), score(5, 3), atol=1e-4) + + +def test_can_overfit_a_tiny_batch(): + torch.manual_seed(0) + m = GPT(CFG) + x = torch.randint(0, CFG.vocab_size, (4, 16)) + y = torch.roll(x, -1, dims=1) + opt = torch.optim.AdamW(m.parameters(), lr=3e-3) + first = None + for _ in range(80): + loss = m(x, y)[1] + first = first if first is not None else loss.item() + opt.zero_grad() + loss.backward() + opt.step() + assert loss.item() < 0.3 * first diff --git a/tests/test_sampling.py b/tests/test_sampling.py new file mode 100644 index 0000000..b62348f --- /dev/null +++ b/tests/test_sampling.py @@ -0,0 +1,60 @@ +import pytest +import torch + +from microgpt.data import CharTokenizer, get_batch +from microgpt.sampling import filter_top_k, filter_top_p, sample_next + + +def test_temperature_zero_is_argmax(): + logits = torch.tensor([[0.1, 2.0, -1.0], [3.0, 0.0, 0.0]]) + assert sample_next(logits, temperature=0).flatten().tolist() == [1, 0] + + +def test_negative_temperature_rejected(): + with pytest.raises(ValueError): + sample_next(torch.zeros(1, 3), temperature=-1.0) + + +def test_top_k_keeps_exactly_k(): + logits = torch.tensor([[1.0, 5.0, 3.0, 4.0, 2.0]]) + kept = torch.isfinite(filter_top_k(logits, 2)) + assert kept.tolist() == [[False, True, False, True, False]] + + +def test_top_p_keeps_smallest_nucleus(): + probs = torch.tensor([[0.5, 0.3, 0.15, 0.05]]) + out = filter_top_p(probs.log(), 0.7) + assert torch.isfinite(out).tolist() == [[True, True, False, False]] + + +def test_top_p_always_keeps_the_top_token(): + out = filter_top_p(torch.tensor([[10.0, 0.0, 0.0]]), 0.01) + assert torch.isfinite(out).tolist() == [[True, False, False]] + + +def test_samples_respect_top_k(): + g = torch.Generator().manual_seed(0) + logits = torch.tensor([[0.0, 1.0, 2.0, 3.0, 4.0]]).expand(500, -1) + draws = sample_next(logits, top_k=2, generator=g) + assert set(draws.flatten().tolist()) <= {3, 4} + + +def test_sampling_matches_softmax_distribution(): + g = torch.Generator().manual_seed(0) + logits = torch.log(torch.tensor([[0.7, 0.2, 0.1]])).expand(20000, -1) + freq = torch.bincount(sample_next(logits, generator=g).flatten(), minlength=3) / 20000 + assert torch.allclose(freq, torch.tensor([0.7, 0.2, 0.1]), atol=0.02) + + +def test_tokenizer_roundtrip(): + text = "hello, world\n" + tok = CharTokenizer.from_text(text) + assert tok.decode(tok.encode(text)) == text + assert tok.vocab_size == len(set(text)) + + +def test_get_batch_targets_are_shifted_inputs(): + data = torch.arange(100) + x, y = get_batch(data, batch_size=4, block_size=8) + assert x.shape == y.shape == (4, 8) + assert torch.equal(y, x + 1)