Tiny GPT / mini-LLM end-to-end

Milestone centerpiece: your mini-LLM

This lesson is the capstone of the language-modeling module. You assemble tokenizer + embedding + stacked transformer blocks + LM head into a trainable GPT-style model, run a real training loop, and generate text from your own checkpoint.

You are not reproducing ChatGPT. You are building the same software shape at a scale that runs on a laptop or free Colab GPU: train.py, model.py, sample.py, checkpoints, samples, README. That shape is what teams ship — scaled up with more data, layers, and compute.

Callout — ship the pipeline, not just the math: Hiring managers and future you care that the repo runs. A tiny model with clean code beats a large notebook that only worked once.

GPT architecture (decoder-only transformer)

GPT is a decoder-only stack: causal self-attention blocks with no encoder cross-attention. Each block:

x = x + MultiHeadAttention(LayerNorm(x), causal_mask=True)
x = x + MLP(LayerNorm(x))   # two linear layers with GELU, expand 4x typical

Full model:

token_ids → TokenEmbedding + PositionEmbedding
           → N × TransformerBlock
           → LayerNorm
           → Linear (lm_head) → logits over vocab

Weight tying: Often share weights between token embedding and lm_head — reduces parameters and can improve small models. Optional for your tiny GPT.

Position embeddings

Attention is permutation-invariant without position info. GPT-2 uses learned absolute position embeddings added to token embeddings, one vector per position up to block_size.

Alternative in modern models: RoPE (rotary positional embeddings). For your first GPT, learned positions match Karpathy's nanoGPT and are simpler to implement.

Hyperparameters for a trainable tiny model

Start here if you have no GPU; scale slightly if you do:

Parameter count might land in 1–10M range — small enough for CPU overnight training on a modest corpus (1–10 MB text).

Corpus suggestions: tiny Shakespeare, TinyStories subset, or your own markdown collection. Domain-specific corpora (your notes, project READMEs) produce fun personalized samples.

Training loop: same harness, new model

Reuse the training patterns from the foundations module:

1. Load tokenized corpus as flat integer array or memmap.
2. Sample random contiguous chunks of length block_size + 1 (inputs = first T, targets = shifted by 1).
3. Forward → cross-entropy over vocab → backward.
4. Log train/val loss; checkpoint best val.

# One training step (conceptual)
xb, yb = get_batch()  # (B, T) each
logits = model(xb)    # (B, T, vocab)
loss = F.cross_entropy(logits.view(-1, vocab), yb.view(-1))

Val loss on language modeling is meaningful — compare runs with identical val holdout. Target: val loss clearly below untrained baseline (random ≈ ln(vocab_size)).

Hardware notes

• CPU: Set small model; expect hours. Use --max_iters smoke tests first.
• Colab free GPU: Move model and batches to CUDA; increase batch or model modestly.
• Apple MPS: device = "mps" often works; fall back to CPU on ops errors.

Document exactly what you used in README — reviewers want honest compute stories.

Generation after training

Load best checkpoint; run autoregressive sampling (lesson 3.6 goes deeper on decoding):

model.eval()
context = torch.tensor([[start_token_id]], device=device)
for _ in range(max_new_tokens):
    logits = model(context[:, -block_size:])
    next_logits = logits[0, -1, :] / temperature
    probs = F.softmax(next_logits, dim=-1)
    next_id = torch.multinomial(probs, 1)
    context = torch.cat([context, next_id], dim=1)

Decode with your BPE decoder. Save 3–5 samples of 300+ tokens each in samples/ with prompts noted.

Good samples on tiny models: local style matching, some valid words/phrases, thematic vocabulary from corpus. Bad signs: infinite repetition (fix with temperature/top-k), Unicode garbage (tokenizer bug), unchanged loss (broken mask or labels).

Debugging checklist

Run a 10-batch overfit test on 1 batch — loss should approach zero. If not, bug in model or data.

Repo layout

m3/mini_llm/
  train.py
  model.py
  sample.py
  config.yaml
  checkpoints/best.pt
  samples/
  README.md

README must include: architecture summary, hyperparameters, final val loss, hardware, commands to train and sample.

Engineering problem (staff framing)

Wire tokenizer→embeddings→blocks→LM head→loss→sample into one trainable GPT. Integration bugs dominate.

Diagram — Tiny GPT data path

flowchart TD
  Tok[Tokens] --> Emb[Tok + Pos emb]
  Emb --> Blk[N × Transformer blocks]
  Blk --> Head[LM head]
  Head --> CE[CE loss / sample]

Precise definitions & mental model

Positional encodings, residual stream, weight tying optional, context length limit.

Tradeoffs — when to use what

Tiny GPT teaches algorithms; scale teaches emergence — do not conflate.

Failure modes (interview + on-call)

Pos emb longer than train context at inference; untied vocab mismatch; eval in train mode.

Production & OSS practices

Config dataclass + seed + checkpoint schema; compare curves to known nanoGPT-style baselines.

Deep dive (FAANG / OSS bar)

Minimal config surface (treat as API)

Pin these in a dataclass and log them into every checkpoint:

n_layer, n_head, d_model, d_ff, vocab_size, context_len, dropout, lr, batch_size, seed

Gradient sanity checklist

1. Loss on random init ≈ ln(vocab_size).
2. Overfit 1–2 batches to ~0 loss (proves capacity + labeling).
3. Then generalize on held-out; if you cannot overfit, debug before scaling data.

Diagram — train vs generate modes

flowchart TD
  subgraph Train
    Batches --> Fwd_train[Forward parallel]
    Fwd_train --> CE[CE loss]
  end
  subgraph Generate
    Prompt --> Loop[Sample one token]
    Loop --> Append --> Loop
  end

Micro-project: Train on CPU/Colab; generate text

Assemble and ship the full mini-LLM:

1. Wire tokenizer vocab into model vocab_size.
2. Implement GPT with at least 4 layers; causal mask verified.
3. Train until val loss plateaus; save best checkpoint.
4. Generate ≥3 samples; commit to repo.
5. Document hyperparameters, iters, wall time, and device in README.

This folder is the milestone artifact — polish matters. Future modules reference this checkpoint for sampling ablations and honest limits essays.

Optional stretch: integrate Weights & Biases; export loss curve; share one sample in module README.

Checklist

• End-to-end train + sample scripts run from README commands
• Best checkpoint committed or reproducible within stated iters
• Val loss logged and reasonable vs random baseline
• Samples committed with prompts and decoding settings noted
• Hardware and hyperparameters documented