Add microgpt: from-scratch GPT with RoPE, KV cache, sampling, tests and benchmark

.gitignore

@@ -0,0 +1,9 @@
+.venv/
+__pycache__/
+*.egg-info/
+.pytest_cache/
+.ruff_cache/
+data/
+checkpoints/
+*.pt
+.DS_Store

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 `<rope(q, m), rope(k, n)>` 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

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"]

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()

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"

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

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()

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, <rope(q, m), rope(k, n)> 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

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)

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()

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():
+    """<rope(q, m), rope(k, n)> must equal <rope(q, m+s), rope(k, n+s)>."""
+    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

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)