nanoGPT++ — a modern decoder-only Transformer

Goal

Build, train, and evaluate a ~50-150M parameter decoder-only Transformer with modern conventions: pre-norm RMSNorm, SwiGLU FFN, RoPE, GQA, weight tying, FlashAttention via SDPA, bf16, AdamW, cosine LR with warmup. Train on a public corpus. Compare to a reference (HuggingFace gpt2).

The single most respected portfolio project in the field. Karpathy's nanoGPT showed that a usable LLM is ~300 lines. Yours will be ~600-900 lines including infra and eval.

Time: 2-4 weeks part-time.

Prerequisites

• 00_foundations/, 01_neural_foundations/, 02_transformer/ (all 5 lessons), 03_pretraining/ (1, 3, 5).
• A GPU. A single 4090, A100, or even Colab Pro will do at this scale.

Tech stack

• PyTorch 2.x
• tiktoken (or your tokenizer from Project 1)
• wandb for monitoring
• datasets for streaming
• safetensors for checkpointing

Architecture target

flowchart TB
    TOKENS[input ids B,T] --> EMB[Token embedding d=512-1024]
    EMB --> B1[Block 1]
    B1 --> B2[Block 2]
    B2 --> B3[...]
    B3 --> BL[Block L=8-12]
    BL --> NORM[Final RMSNorm]
    NORM --> HEAD[LM head, weight-tied]
    HEAD --> LOGITS[logits B,T,V]

    subgraph Block
       direction LR
       X[x] --> NA[RMSNorm] --> A[GQA + RoPE + SDPA causal] --> R1((+))
       X --> R1
       R1 --> NM[RMSNorm] --> M[SwiGLU FFN] --> R2((+))
       R1 --> R2
       R2 --> Y[y]
    end

Step-by-step

1. Tokenize

Use cl100k_base (or o200k_base) and pre-tokenize a public corpus into uint16/uint32 .bin shards (Lesson 3.2 pattern).

Recommended starter corpora:

• 1B-token slice of HuggingFaceFW/fineweb-edu (sample-1BT config).
• TinyStories (~1B tokens, very forgiving).
• enwik8 / enwik9 for character-level baselines.

Save train and val splits.

2. Implement components

Per Lesson 2.4: RMSNorm, SwiGLU, RoPE helpers, GQA, Block, GPT. Use F.scaled_dot_product_attention for FlashAttention.

Match standard configs:

Print param count; verify against the formula L * (4 + 3m) * d^2 + V*d.

3. Training loop

• Cosine LR with 200-step warmup.
• AdamW (β=(0.9,0.95), wd=0.1).
• Gradient clipping max-norm 1.0.
• bf16 autocast.
• Gradient accumulation to simulate batch ≥ 256k tokens.
• WandB logging (loss, lr, grad_norm, tokens/s).

4. Eval during training

• Val loss every 200 steps on a held-out shard.
• Sample 200 tokens from a fixed prompt every 1000 steps; eyeball.
• Save checkpoint when val loss improves.

5. Evaluation

After training:

• Compute final perplexity on wikitext-2.
• Compute bits-per-byte (loss * tokens_per_byte / ln(2)).
• Run a small MMLU subset (10 random subjects, 100 questions) zero-shot. (Use lm-eval-harness.)
• Compare to gpt2-small (124M) baseline.

6. Inference improvements

• Implement KV cache for model.generate.
• Add top-p / top-k sampling.
• Measure tokens/sec at batch=1.
• Bonus: load via vLLM for serving.

Acceptance criteria

• Modern stack: RMSNorm, SwiGLU, RoPE, GQA, weight tying.
• Trains end-to-end with no NaNs on 1B-token corpus.
• Achieves perplexity within 1.5× of gpt2-small on wikitext-2 (or strictly better at the same param count).
• WandB dashboard with loss / lr / grad_norm.
• KV-cached generate() is at least 5× faster than no-cache.
• README with config, training curves, sample outputs, MMLU mini-eval.
• All hyperparameters and seed in a config.yaml.

Stretch goals

• FSDP for multi-GPU on a 410M model.
• torch.compile for ~20% speedup.
• FP8 training if you have an H100.
• Mixture-of-experts variant: Mixtral-style 4×30M experts → similar quality, fewer active params per token.
• Mamba block drop-in to compare.
• Pretrain a coding mini-model (FIM data) and benchmark on HumanEval-tiny.

Common pitfalls

• Forgetting is_causal=True in SDPA → loss looks ok but model "memorises."
• LR too high for the chosen d_model → divergence at step ~500.
• Forgetting EOS between docs in tokenization → model never learns to stop.
• Using fp16 without a scaler → silent NaNs.
• Buffer not on device (cos/sin / causal mask) → mysterious slowdowns.

Story / portfolio

• Title: "Pretraining a 125M Transformer the Llama-3 way."
• Tables: param count, train hours, val ppl over time, mini-MMLU.
• Plot: loss curve with annotated phases (warmup, cosine).
• Sample output: 3 fully-generated paragraphs, one for each temperature 0.5/0.8/1.2.
• Code: single train.py <500 lines, plus modules.

This is the project that says "I can build the thing, not just call the API."