L4 · 09 · Capstone: Train & Serve a Tiny GPT

The Level 4 capstone closes the loop on everything this level covered: implement a tiny GPT (page 02), train it (03), quantise and serve it (04), benchmark it (04/05), and document the whole thing like a real system (06–07). You end with a model you built and a design doc you'd actually show a team.

The arc

Step 1: implement and train

Use the attention you wrote in lab 2, or adapt nanoGPT. Keep it small enough to train on CPU/a laptop in minutes–hours: d_model ~128, 4 layers, 4 heads, character-level tokeniser. Pick a corpus you know well (a book, your docs, code).

import torch, torch.nn as nn, torch.nn.functional as F

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.attn = nn.MultiheadAttention(d, h, batch_first=True)
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.ff = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))
    def forward(self, x, mask):
        a, _ = self.attn(x, x, x, attn_mask=mask, need_weights=False)
        x = self.ln1(x + a)
        return self.ln2(x + self.ff(x))

class TinyGPT(nn.Module):
    def __init__(self, V, d=128, h=4, L=4, T=128):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(T, d)
        self.blocks = nn.ModuleList(Block(d, h) for _ in range(L))
        self.head, self.T = nn.Linear(d, V), T
        mask = torch.triu(torch.ones(T, T), 1).bool()   # causal
        self.register_buffer("mask", mask)
    def forward(self, idx):
        t = idx.size(1)
        x = self.tok(idx) + self.pos(torch.arange(t))
        for b in self.blocks: x = b(x, self.mask[:t, :t])
        return self.head(x)

Train it (page 02's loop): get batch → logits → cross-entropy vs. shifted targets → backward → step. Watch the loss drop; generate samples along the way to see it learn.

Step 2: quantise and serve

Quantisation for your own tiny model is dramatic but simple — cast weights to int8/fp16 and measure the size/speed change. For serving, wrap generation behind a small endpoint:

import json
import time


class TinyModel:
    """Stand-in for a real GPT.generate — demonstrates the serving pattern."""
    def generate(self, prompt_tokens, max_new):
        t0 = time.perf_counter()
        out = [7, 3, 3, 7][:max_new]          # pretend to decode
        time.sleep(0.005 * max_new)          # pretend decode latency
        return out, (time.perf_counter() - t0)


def handle_request(body: dict, model) -> dict:
    """OpenAI-ish request -> response with usage + latency."""
    prompt = body.get("prompt_tokens", [])
    max_new = min(body.get("max_tokens", 16), 128)
    t0 = time.perf_counter()
    tokens, decode_s = model.generate(prompt, max_new)
    return {
        "choices": [{"tokens": tokens}],
        "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(tokens)},
        "latency_ms": round((time.perf_counter() - t0) * 1000, 1),
        "finish_reason": "length" if len(tokens) == max_new else "stop",
    }


if __name__ == "__main__":
    m = TinyModel()
    r = handle_request({"prompt_tokens": [1, 2, 3], "max_tokens": 3}, m)
    print(json.dumps(r))
    assert r["usage"]["completion_tokens"] == 3

Point: the serving layer tracks usage and latency per call — the same fields you logged in L3·07. A real version uses your trained model and a real HTTP server (or vLLM/llama.cpp once converted).

Step 3: benchmark

• TTFT and tokens/s for short vs. long prompts.
• Quality on 5–10 prompts from your corpus's domain (does it capture the style?).
• Before/after quantisation: size, speed, and any quality change you can spot.

Step 4: the design document

One to two pages, model-card + design-doc style:

• Goal & scope — what it does, what it doesn't.
• Architecture — layers, heads, d_model, context length, parameter count (compute it!).
• Data — corpus, size, cleaning, known gaps.
• Training — objective, optimizer, steps, final loss, hardware.
• Evaluation — benchmark numbers, quality examples, failures.
• Serving — quantisation, latency/throughput, cost notes.
• Limitations & risks — what it can't do, misuse, licence/provenance of the corpus.

Acceptance criteria

• You trained the model and can point to the loss curve dropping.
• It generates output recognisably in the corpus's style.
• A serving endpoint returns usage + latency; you benchmarked it.
• You computed the parameter count and the KV-cache size at max context.
• The design doc would let a teammate understand and reproduce it.