Pretraining From Scratch

Domain: building a transformer/GPT and a BPE tokenizer from first principles — the from-first-principles training-layer competency. Does NOT cover applications-layer fine-tuning, RLHF, or inference optimization; those belong to sibling skills.

Canonical teachers: Karpathy "Neural Networks: Zero to Hero" (micrograd → makemore → "Let's build GPT" → "Let's build the GPT Tokenizer" → "Let's reproduce GPT-2"), Karpathy nanochat (full-stack from-scratch successor to nanoGPT), Raschka "Build a Large Language Model From Scratch", nanoGPT, minbpe, "Attention Is All You Need".

GPT-2 is the pedagogical spine here — the right thing to build first. The modern from-scratch baseline then swaps four components onto that spine (RoPE, RMSNorm, SwiGLU, GQA) and runs attention through FlashAttention/SDPA; see Modern Architecture Deltas.

When to Use This Skill

Activate when the user asks about:

• Implementing autograd / backprop from scratch (micrograd-style)
• Building makemore (bigram, MLP, WaveNet-style character LMs)
• Implementing self-attention, multi-head attention, causal masking
• Building the transformer block (pre-norm vs post-norm, residual, FFN)
• Stacking blocks into a GPT with an LM head and weight tying
• Writing the pretraining loop: cross-entropy, bf16 mixed precision, gradient accumulation, gradient checkpointing, cosine LR schedule with warmup, model checkpointing
• Building a BPE tokenizer from scratch: byte-level, merge algorithm, vocab construction, encode/decode (minbpe-style)
• Reproducing GPT-2 (124M) from scratch end-to-end (nanoGPT path)
• Implementing temperature scaling and top-k sampling for text generation

Scope Boundaries (Use These Skills for Depth)

• Fine-tuning and adaptation (SFT/PEFT/distillation) -> ai-llm; model/provider selection -> ai-architecture-advisor; deployment -> ai-mlops
• Multi-GPU training: DDP, FSDP, tensor/pipeline parallelism -> ai-distributed-training
• Token/param budget, Chinchilla scaling, compute-optimal runs -> ai-scaling-laws
• Dataset curation, deduplication, quality filtering for pretraining -> ai-data-curation-pretraining
• Evaluation harnesses, benchmark design, evals post-pretraining -> ai-evals
• Mixture-of-Experts (MoE): swaps the dense FFN for a router + expert FFNs (DeepSeek-V3/V4, Qwen3-MoE, Kimi-K2, Mixtral). A frontier architectural variant, not a from-scratch fundamental. Build-time mechanics (top-k routing, load-balancing loss, expert granularity, failure modes) are taught in §5 of Architecture Limitations and Workarounds; distributed training stays with ai-distributed-training and serving/inference with ai-llm-inference.
• Classification fine-tuning, instruction/SFT fine-tuning, LoRA/PEFT: post-pretraining applications. Raschka's book covers these; this skill stops at pretraining. -> ai-llm

Default Workflow

1. Autograd first: implement Value class with backward(), build MLP, verify gradients against PyTorch.
2. Character LM ladder: bigram table -> MLP (makemore) -> verify loss convergence and sampling.
3. Attention module: single-head self-attention with causal mask; verify attention weights sum to 1 per row.
4. Multi-head attention: split heads, concatenate, project; compare with nn.MultiheadAttention using the same weights, mask, dropout setting and a numerical tolerance.
5. Transformer block: add FFN (4x, GELU), pre-LayerNorm, residuals; match nanoGPT block.
6. GPT assembly: stack N blocks, add LM head, tie weights with embedding; verify forward pass shape.
7. Pretraining loop: DataLoader, cross-entropy, CUDA bf16 via torch.autocast('cuda', dtype=torch.bfloat16) when supported, gradient accumulation, cosine or WSD LR, checkpoint (weights and data position). Verify on your model with python3 scripts/check_loop.py --model yourmodule:build (a zero-argument factory returning a fresh model): it checks step-0 loss ≈ ln(V), that accumulated micro-batch gradients equal the large-batch gradient, and that changing future tokens never changes earlier logits. Without --model it only runs a built-in demo. Then take the throughput wins: torch.set_float32_matmul_precision('high') for TF32, pad vocab_size 50257 → 50304, and track MFU rather than tokens/sec. See Pretraining Loop; parameter counts, FLOPs per token and run times for GPT-2 124M are in its reference card.
8. BPE tokenizer: byte-level text encoding, count bigram frequencies, greedy merge loop, build vocab, encode/decode round-trip.
9. GPT-2 reproduction: load OpenAI weights via HuggingFace, verify logits match, then train from scratch on FineWeb-Edu. 9a. Sampling: implement temperature scaling and top-k sampling for generation; optionally add a KV-cache for inference speed (see Quick Reference).
10. Modernize: swap to the modern baseline — RoPE for wpe, RMSNorm for LayerNorm, SwiGLU for the GELU-MLP, GQA, and F.scaled_dot_product_attention; optionally train with Muon. See Modern Architecture Deltas.

Modern Baseline

Build GPT-2 first to understand the mechanics, then apply the deltas — the pre-norm residual skeleton is unchanged; you swap sublayers, not the architecture.

GPT-2 (2019)	Modern baseline	Why
Learned absolute pos embed (wpe)	RoPE (rotary, in attention)	Relative position without a learned position table; test quality beyond the trained context
LayerNorm	RMSNorm	Cheaper, no centering/bias, stable at depth
GELU-MLP (4×)	SwiGLU (~8/3×)	Gated FFN improves quality per param
MHA (KV heads = query heads)	GQA (fewer KV heads)	Shrinks KV cache for inference
Hand-rolled softmax attention	F.scaled_dot_product_attention	PyTorch selects an eligible fused kernel on supported CUDA inputs; check the selected backend before claiming its memory or speed gains
AdamW for all params	Muon (2D matrices) + AdamW (embed/head/norms)	Newton-Schulz orthogonalized updates; per-step speedup that shrinks with scale (~2× reported at Moonlight scale, ~1.1× at 1.2B in a controlled benchmark)
No q/k normalization	QK-Norm (RMSNorm on q/k before attention)	Bounds attention-logit growth — a stability default in new dense/MoE recipes, not just a speedrun trick (cf. Kimi K2's MuonClip QK-Clip)

Frontier reference: the modded-nanoGPT speedrun stacks Muon, QK-Norm, ReLU², logit softcap, and embedding-skip connections to reach GPT-2-grade FineWeb validation loss far faster than the original nanoGPT on 8×H100. For the full from-scratch pipeline (tokenizer → pretrain → SFT → RL → serve), Karpathy's nanochat is the successor to nanoGPT. Lookup step: both records move; read the repo README for the current target loss, hardware and time before quoting one, and quote it with the date you read it.

Hyperparameters, Schedule, Stability

• Transfer hyperparameters instead of searching at target scale. Default: parameterize with µP and tune LR/init on a small proxy, then transfer zero-shot to the wide target. µTransfer (Yang et al. 2022, arXiv 2203.03466) tuned on a 40M proxy and beat the published GPT-3 6.7B numbers, with a tuning cost of about 7% of the total pretraining cost. Never grid-search LR at target scale. If you are not using µP, fit η_opt(C) and B_opt(C) on 3+ small proxy budgets and read the target off the fitted law (ai-scaling-laws, workflow step 6). Not needed for a single-scale 124M reproduction, where one LR sweep suffices.
• Schedule and annealing. Cosine when the token budget is fixed; WSD (warmup–stable–decay) when it may change. The decay phase doubles as a cheap data probe (same short anneal with and without a candidate dataset) and as a mid-training stage for the highest-quality data. State warmup in tokens or steps, not as a percentage of the run. See Pretraining Loop.
• Loss spikes (working protocol, not a validated standard). Log max attention logit, max output logit and the update/weight ratio every step; prevent with QK-norm and a small z-loss; on a spike that does not recover, skip the batch window and roll back to the last good checkpoint; if spikes repeat, lower peak LR. See Pretraining Loop.

Quick Reference

Component	Key Detail	Common Mistake
Autograd	Value.backward() accumulates += into .grad, not =	Forgetting to zero grads before .backward()
Embedding	nn.Embedding(vocab_size, n_embd) — random init, learned	Confusing token embed with positional embed shape
Causal mask	torch.tril(torch.ones(T,T)) before softmax; fill -inf or torch.finfo(dtype).min, not 0	Using 0 fill — attention leaks future tokens
Attention math	softmax(QK^T / sqrt(d_k)) * V	Forgetting /sqrt(d_k) — variance explodes
LayerNorm placement	Pre-norm (before attention/FFN) in GPT-2; original paper was post-norm	Post-norm makes deep stacks hard to train
FFN expansion	4x hidden dim, GELU activation	Using ReLU — slight quality difference, matters at scale
Weight tying	LM head shares the embedding matrix (head.weight = wte.weight)	Forgetting tying adds a second V×d matrix: +38.6M params (+31%) at GPT-2 124M, not 2×
Init scaling	std=0.02 for most; residual projections: std=0.02/sqrt(2*n_layer)	Flat 0.02 everywhere — residual stream variance grows
Gradient accumulation	accumulate N micro-batches, divide loss by N, step once	Forgetting to divide loss — effective LR N× too large
bf16 autocast	torch.autocast('cuda', dtype=torch.bfloat16)	Using fp16 without loss scaling — NaN on older GPUs
BPE merges	greedy highest-frequency pair; merge in-place, repeat	Not updating pair counts after each merge — wrong vocab
LR schedule	cosine: linear warmup (GPT-3 used ≈375M tokens; nanoGPT's 715 steps × 524,288 tokens matches it), then cosine decay to ~10% of peak. WSD (trapezoidal) when the token budget is not fixed up front	Skipping warmup — loss spike at start; stating warmup as a % of the run, which changes meaning with run length
Temperature	logits / temperature before softmax; T<1 sharpens (more deterministic), T>1 flattens (more random)	Applying temperature after softmax — has no effect on the distribution
Top-k sampling	zero out all logits except the top-k before softmax; draw from the remaining distribution	Top-k=1 is greedy decoding; top-k=vocab_size is pure sampling
KV-cache	at inference, cache K and V tensors for all past positions; on each new token only compute Q/K/V for the single new position and append to cache	Re-computing all K/V at each generation step; the cache removes the redundant projection work (O(T²) → O(T) for K/V), not the attention itself — scoring is still O(T) per step, so total generation stays O(T²)

Scale-Up Gate

Prove the tokenizer, data loader, masking, loss, optimizer order, checkpoint restore, and sample generation on a tiny run before reserving large compute. Then run a fixed-budget pilot that records effective tokens, loss by source slice, gradient and activation health, throughput, utilization, and restart equivalence. Scale only when the loss curve and downstream probes improve as expected, the input pipeline is not the bottleneck, and a costed stop rule is written. Successful allocation or falling training loss alone does not justify the next scale.

Known Traps

• Zero-grad placement: call optimizer.zero_grad() before the forward pass (or set_to_none=True for speed), not after .step().

• Post-norm vs pre-norm: original "Attention Is All You Need" uses post-norm; GPT-2 and nanoGPT use pre-norm. Pre-norm trains more stably at depth.

• Causal mask fill value: use float('-inf') or torch.finfo(dtype).min, not a hard-coded -1e9. In fp32, -1e9 does not leak (future weights are exactly 0), but in fp16 masked_fill(-1e9) raises an overflow error. Guard fully masked rows (padding): with -inf they turn into NaN, with finfo.min into a uniform row.

• Gradient accumulation scaling: divide the loss by the accumulation steps inside the micro-batch loop, not outside.

• Weight tying in state_dict: when saving checkpoints, the LM head weight is the same tensor as the embedding weight — loading requires care to avoid double-counting params.

• BPE encode-decode round-trip: bytes, not characters — always encode text as UTF-8 bytes first before running BPE.

• DataLoader seeding: fix random seeds for reproducibility across runs; DataLoader worker seeds need explicit worker_init_fn.

• torch.compile interaction: torch.compile + gradient checkpointing can conflict in some PyTorch versions — test before enabling both. A compiled model also prefixes state_dict keys with _orig_mod., which breaks checkpoint loading into an uncompiled model.

• DDP gradient sync in the micro-loop: require_backward_grad_sync is reset to True by DDP on every forward, so it must be re-assigned per micro-step (or use model.no_sync()). Setting it once outside the loop either all-reduces every micro-step or never syncs at all — both silent.

• Resuming without the data position: restoring weights and optimizer but restarting the loader re-trains on seen shards with no error.

• Verify PyTorch API details (autocast dtype names, torch.compile flags, DataLoader args) against current PyTorch docs before recommending.

• Verify current nanoGPT and minbpe repo states (file structure, hyperparameters) against the GitHub repos — they are actively maintained.

Common Anti-Patterns

• Implementing attention without verifying attn_weights.sum(dim=-1) is all-ones (no causal leak check).
• Skipping the PyTorch parity check: always compare custom layer output to torch.nn. equivalent before stacking.
• Starting with the full GPT before the single-head attention works — build bottom-up.
• Training without a baseline loss: for character-level with vocab V, random model should give ln(V) loss; check this at step 0.
• Using Adam with default betas=(0.9, 0.999) — GPT-3 used betas=(0.9, 0.95) for stability at scale, and nanoGPT follows it (the GPT-2 paper does not report betas).
• Tokenizing the entire dataset in memory — stream and chunk for large corpora.
• Shipping the GPT-2 architecture as the final product — it is the teaching spine, not the modern baseline. Apply the modern deltas (RoPE/RMSNorm/SwiGLU/GQA/SDPA) once the GPT-2 build verifies.

Core Principles

1. Build then read: implement first, then verify against PyTorch source or the paper. Reading first encourages copy-paste, not understanding.
2. No black boxes: every component must be verified with a unit check before it's stacked.
3. One component at a time: single-head attention -> multi-head -> block -> GPT. Never jump layers.
4. PyTorch parity check: compare custom attention with nn.MultiheadAttention using matched weights, masks and dropout, within a numerical tolerance.
5. Fail loud on training metrics: if step-0 loss deviates from ln(vocab_size) by >10%, stop and debug — don't train through bad initialization.

Navigation: Core References

• Transformer From Scratch — attention math, block assembly, weight init, GPT architecture notes
• BPE Tokenizer — byte-level BPE algorithm, merge loop, vocab construction, encode/decode; plus the tokenizer landscape (SentencePiece, Unigram, SuperBPE, vocab sizing, fertility/compression evaluation)
• Pretraining Loop — GPT-2 124M reference card (params, FLOPs/token, run times, warmup), training loop anatomy, mixed precision, gradient accumulation, cosine/WSD and annealing, MFU, loss-spike protocol, checkpointing, per-slice eval during pretraining
• Research Recipes — µP/muP hyperparameter transfer and the coordinate check, WSD cooldown shape and length, loss-spike classes and their fixes, mid-training and continued-pretraining mixes (replay, re-warming, anneal probes)
• Modern Architecture Deltas — GPT-2 → modern baseline: RoPE, RMSNorm, SwiGLU, GQA, FlashAttention/SDPA, Muon and the speedrun frontier
• Architecture Limitations and Workarounds — failure-mode companion: each component's limitation → workaround → tradeoff (softmax pathologies/attention sinks, MHA→MQA→GQA→MLA + decoupled RoPE, positional design space + YaRN/NTK, MoE routing pitfalls, norm/residual/depth stability, fp8/fp4 precision, long-context, encoder/decoder/encoder-decoder contrast)
• Adaptive Depth and Conditional Compute — depth-axis conditional computation: Mixture-of-Depths, early exit (LayerSkip/TIDE), looped and recursive transformers (Mixture-of-Recursions, AdaPonderLM), cross-layer weight sharing (ALBERT, tied experts), modular-NN framing; lever-selection table
• Structured and Low-Rank Parameterization — replacing dense weights: low-rank, block-diagonal, butterfly, Monarch, Kronecker, BLAST; SVD-LLM/ASVD post-hoc factorization; MLA as low-rank KV; MonarchLinear and BlockLowRankLinear sketches; when it beats or composes with pruning/quantization

Scripts

• scripts/check_loop.py --model module:factory — runs three checks on your model: step-0 loss ≈ ln(V), accumulated micro-batch gradient ≡ large-batch gradient, and no future leak (earlier logits unchanged when later tokens change). The factory takes no arguments and returns a fresh nn.Module whose forward returns logits or (logits, ...). Without --model it runs a built-in demo that says nothing about your code. Requires PyTorch (CPU is fine). Exit 0 pass, 1 a check failed, 2 usage/import error. Tests: python3 scripts/test_check_loop.py.

Learnings Loop

When prior decisions or pitfalls are relevant, consult learnings.consolidated.md if present; use learnings.md only for needed history or as the available fallback. Otherwise skip both.

After applying it, if you encountered a pattern worth remembering, a mistake worth preventing, or a domain fact that surprised you, append one dated bullet to learnings.md via agents-skills-feedback-loop/scripts/append_learning.py. Do not modify SKILL.md itself.