nanoGPT

Part I: The mental model

raw text corpus (Shakespeare, OpenWebText)
        │
        ▼
data/<dataset>/prepare.py ── tokenize once, offline
        │
        ▼
train.bin / val.bin ── flat uint16 token streams on disk
        │
        ▼
get_batch() ── np.memmap + random offsets → (x, y) on GPU
        │
        ▼
GPT.forward() ── wte+wpe → 12 Blocks → ln_f → logits → cross_entropy
        │
        ▼
scaler.scale(loss).backward() ── AMP, grad accumulation, DDP sync
        │
        ▼
clip → AdamW step → zero_grad ── one iteration done
        │
        ▼
ckpt.pt ── periodically, when val loss says so

The one-sentence identity: nanoGPT is a faithful GPT-2 and an honest pretraining loop in two files you can read in one sitting. model.py defines the network with no configuration framework, no registry, no abstraction for future architectures: a causal self-attention module, an MLP, a block that stacks them, and a GPT class that ties it together and computes the loss. train.py contains, in linear order and without indirection, essentially every technique a real pretraining run uses: distributed data parallelism, gradient accumulation, mixed precision with a gradient scaler, gradient clipping, decoupled weight decay, warmup plus cosine decay, and torch.compile.

The rest of the repository is support: one prepare.py per dataset that turns a corpus into flat binary token files, sample.py for generation, bench.py for timing the forward-backward in isolation, configurator.py for the deliberately crude configuration scheme, and two notebooks about parameter counts and scaling. There is no dataloader class, no callback system, no plugin interface. Where a framework would hand you a lever, this repo hands you the line of code the lever would have pulled.

One honesty note up front: the README has carried a notice since November 2025 that nanoGPT is old and effectively frozen, with nanochat as its modern successor. For learning how transformers are trained, frozen is a feature; the code is stable, complete, and still the best first training codebase to read.

Part II: Using it, a practical tutorial

There is no package to install; you clone and run scripts from the repo root. On Linux and macOS alike:

git clone https://github.com/karpathy/nanoGPT
cd nanoGPT
pip install torch numpy transformers datasets tiktoken wandb tqdm

The fastest first result needs no training at all: sample.py can pull real GPT-2 weights through Hugging Face and generate with them, which doubles as a proof that the reimplementation in model.py is correct:

python sample.py --init_from=gpt2 --device=cpu \
    --start="The meaning of life is" --num_samples=2 --max_new_tokens=50

Expect a "loading weights from pretrained gpt: gpt2" line, then two plausible-but-drifty GPT-2 completions separated by dashes. The canonical first training run is a character-level model on the complete works of Shakespeare:

python data/shakespeare_char/prepare.py
python train.py config/train_shakespeare_char.py
python sample.py --out_dir=out-shakespeare-char

The prepare step prints the dataset stats (1,115,394 characters, vocab size 65, 1,003,854 train tokens) and writes train.bin, val.bin, and meta.pkl into the dataset directory. Training a 6-layer, 6-head, 384-dim model at context 256 takes about three minutes on an A100 and reaches a best validation loss around 1.47 per the README; the loop prints a loss line every 10 iterations and evaluates every 250. Sampling then produces Shakespeare-shaped pastiche, speaker names and all. On a machine with no GPU, the README's exact CPU flag set works in about the same three minutes at smaller scale:

python train.py config/train_shakespeare_char.py \
    --device=cpu --compile=False --eval_iters=20 --log_interval=1 \
    --block_size=64 --batch_size=12 --n_layer=4 --n_head=4 --n_embd=128 \
    --max_iters=2000 --lr_decay_iters=2000 --dropout=0.0

reaching a loss around 1.88; on Apple Silicon, --device=mps is markedly faster than CPU. The serious reproduction is the same two steps at scale: tokenize OpenWebText with python data/openwebtext/prepare.py (it downloads about 54 GB into the Hugging Face cache and writes a ~17 GB train.bin), then launch

torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py

which on 8 A100 40GB GPUs takes roughly four days to reach val loss ~2.85 on OWT, matching a GPT-2 checkpoint fine-tuned onto the same data. That the toy run and the reproduction are the same script with different configs is the whole point of the project. Now the beginner mistakes, each of which teaches a design decision.

Mistake one: running from the wrong directory

Every path in the repo is relative: train.py does exec(open('configurator.py').read()) and reads data from data/<dataset>/. Run it from anywhere but the repo root and you get FileNotFoundError: 'configurator.py':

# wrong
python nanoGPT/train.py nanoGPT/config/train_shakespeare_char.py

# right
cd nanoGPT && python train.py config/train_shakespeare_char.py

Mistake two: CPU without disabling compile

compile = True is the default, and torch.compile is slow to warm up and historically flaky off CUDA (the README's troubleshooting section says exactly this). On CPU or unsupported platforms, pass both flags:

# wrong: long compile stalls or backend errors
python train.py config/train_shakespeare_char.py --device=cpu

# right
python train.py config/train_shakespeare_char.py --device=cpu --compile=False

Mistake three: fighting the configurator's type check

configurator.py parses --key=value with literal_eval and then asserts the new value has the same type as the default. So --max_iters=5000.0 dies on the assert (float replacing int), --device=cpu works only because the eval fails and falls back to string, and a key that does not exist in train.py's globals raises ValueError: Unknown config key. The error messages are terse because the whole configurator is 47 lines; when an override misbehaves, read it, it is faster than guessing.

Mistake four: resuming into a fresh model

init_from defaults to 'scratch'. Restarting a crashed run without --init_from=resume reinitializes the weights and, because always_save_checkpoint defaults to true at GPT-2 scale, happily overwrites your checkpoint at the next eval:

# wrong: starts over and clobbers out/ckpt.pt at the next eval
python train.py config/train_gpt2.py

# right: restores model, optimizer state, and iter_num from ckpt.pt
python train.py config/train_gpt2.py --init_from=resume

Mistake five: expecting a dataset knob that does not exist

There is no epoch setting, no shuffle flag, no dataloader worker count. The data pipeline is random offsets into a memory-mapped token stream (Part IV, Stage 3), so those concepts genuinely do not exist here, and looking for them means you have not yet read get_batch, which is nine lines.

Part III: When it is the right tool

nanoGPT is the right tool for three jobs. Learning: it is the minimum complete example of real pretraining, and reading it is the fastest path to understanding what frameworks like Trainer abstract. Research on architecture and optimization ideas at small scale: the code is so short that a fork stays understandable, and an idea can be wired in within an hour, which is why so many papers' baselines trace back to this repo. And honest small-scale pretraining or GPT-2-class fine-tuning, where config/finetune_shakespeare.py shows the pattern: initialize from a GPT-2 checkpoint, lower the learning rate, train briefly.

The alternatives win elsewhere. For fine-tuning modern architectures with an ecosystem around them, Hugging Face Transformers and its Trainer are the practical choice; nanoGPT knows exactly one architecture, GPT-2 with learned position embeddings, and rotary embeddings or newer attention variants are things you would add yourself (the README's own todo list says as much). For serious multi-node pretraining, torchtitan and its like exist precisely where nanoGPT stops, at FSDP-style sharding and multi-dimensional parallelism. For the full modern pipeline through chat fine-tuning and inference, nanochat is the declared successor. For serving a trained model, sample.py is a demonstration, not an engine; that problem belongs to vLLM.

The architecture-shaped warning: nanoGPT scales by replication only, so it stops where one GPU's memory stops. DDP gives every GPU a complete copy of the model, its gradients, and both AdamW moment buffers; more GPUs buy you throughput, never a bigger model. Trying to train a multi-billion-parameter model by adding nodes to torchrun fails not gradually but immediately, at the first allocation. Sharded approaches split those tensors across devices, which is a different codebase, not a different flag. Relatedly, the README warns that multi-node DDP without a fast interconnect will run but crawl, because gradient all-reduce crosses the network every iteration.

nanoGPT (DDP): replicas             sharded (FSDP/ZeRO): partitions
GPU0 [model|grads|adam m,v]         GPU0 [1/N of each]
GPU1 [model|grads|adam m,v]         GPU1 [1/N of each]
  ...  all-reduce grads               ...  gather / reduce-scatter
max model = one GPU's memory        max model = cluster memory

Part IV: The full life of one training iteration

The canonical operation is one pass through the while True: loop at the bottom of train.py, with the GPT-2 defaults: micro-batch 12, block size 1024, gradient accumulation 40 (written 5 * 8 in the file, five steps per GPU across eight GPUs). Every stage below names lines you can find in the two files.

Stage 0: what exists before the loop

By the time the loop starts, train.py has already: executed configurator.py over its own globals to apply your config file and flags; detected DDP by checking the RANK environment variable that torchrun sets, initialized the NCCL process group, pinned each process to its GPU, and divided gradient_accumulation_steps by the world size so total batch stays fixed; seeded each rank as 1337 + seed_offset so data draws differ across ranks; enabled TF32 matmuls; chosen bfloat16 if the GPU supports it, else float16; built the model (from scratch, from ckpt.pt, or from GPT-2 weights); created a GradScaler that is only enabled for float16; built the AdamW optimizer via model.configure_optimizers; optionally compiled the model; and wrapped it in DDP. It also printed the number that makes the run legible: tokens_per_iter = 40 * 12 * 1024 = 491,520, about half a million tokens per iteration, times 600,000 iterations, roughly 300 billion tokens.

Stage 1: set the learning rate (train.py, get_lr)

Each iteration begins by computing lr = get_lr(iter_num) and writing it into every optimizer param group by hand. The schedule is three cases: linear warmup for the first 2,000 steps, cosine decay from 6e-4 down to 6e-5 between warmup and step 600,000, and the floor after that. No scheduler object; the deep dive draws the curve.

Stage 2: the eval gate (train.py, estimate_loss)

Every eval_interval iterations (2,000 by default), the master process switches the model to eval mode and averages loss over 200 random batches from each split, since a single batch's loss is far too noisy to checkpoint on. If val loss improved, or always_save_checkpoint is set, it writes ckpt.pt: model state dict, optimizer state (those AdamW moments are why resumption works properly), model_args, iter_num, best_val_loss, and the full config dict.

Stage 3: fetch a batch (train.py, get_batch)

The data loader is nine lines. It opens the split's .bin file as np.memmap with dtype uint16, meaning the OS pages in only the bytes touched and a 17 GB training file needs no RAM; the memmap is deliberately re-created on every call to dodge a known memory-leak pattern, per the comment's Stack Overflow citation. It draws batch_size random offsets, slices block_size tokens at each for x and the same window shifted one token right for y, upcasts uint16 to int64 because embedding lookups demand it, and moves both to GPU with pin_memory() and non_blocking=True so the copy overlaps compute. The targets are just the inputs shifted by one: that single line is the entire supervision signal of language model pretraining.

Stage 4: the micro-step loop (train.py + model.py forward)

Now the accumulation loop runs 40 micro-steps (5 per process under 8-GPU DDP). Each micro-step runs the forward pass inside the autocast context: model(X, Y) enters GPT.forward in model.py, which looks up token embeddings wte and learned position embeddings wpe, adds them, applies dropout, runs the twelve blocks (each: x = x + attn(ln_1(x)) then x = x + mlp(ln_2(x))), applies the final LayerNorm, projects through the tied lm_head to (12, 1024, 50304) logits, and computes F.cross_entropy against the shifted targets with ignore_index=-1. The loss is divided by gradient_accumulation_steps so that accumulating gradients over micro-steps averages rather than sums. Two subtleties sit right here in the loop body. First, the next batch is fetched immediately after the forward is launched, so the CPU prepares data while the GPU computes. Second, under DDP the flag model.require_backward_grad_sync is set true only on the last micro-step, so gradients all-reduce across GPUs once per iteration instead of once per micro-step; the comment explains this is what no_sync() does internally, minus the context-manager ceremony. Then scaler.scale(loss).backward() accumulates gradients, scaled up if float16 needs protecting from underflow.

Stage 5: clip, step, update (train.py)

After the micro-steps, four lines in strict order: scaler.unscale_(optimizer) brings gradients back to true scale, clip_grad_norm_(model.parameters(), 1.0) rescales the global gradient norm if it exceeds 1.0 (the order matters: clipping scaled gradients would clip against the wrong threshold), scaler.step(optimizer) runs AdamW unless the scaler found inf/nan gradients this step, in which case it skips and shrinks the scale, and scaler.update() adjusts the scale for next time. Under bfloat16 the scaler is disabled and all of this degenerates to plain clip-and-step. Finally optimizer.zero_grad(set_to_none=True) releases gradient memory rather than filling it with zeros.

Stage 6: timing and MFU (train.py + model.py, estimate_mfu)

The log line costs something: loss.item() forces a CPU-GPU synchronization, which the comment flags. The printed loss is multiplied back by the accumulation steps to approximate the true batch loss. After five settle-in iterations the script also prints model FLOPs utilization: estimate_mfu computes FLOPs per token as 6N + 12 * L * H * Q * T (the PaLM appendix formula: six FLOPs per parameter plus the attention term), multiplies out to FLOPs per iteration, divides by elapsed time, and expresses the result as a fraction of an A100's 312 bfloat16 TFLOPS, smoothed with an exponential moving average. Then iter_num += 1, and the loop breaks past max_iters.

And secondarily: one generation (sample.py + model.py, generate)

sample.py is the inference mirror. It loads ckpt.pt, strips the '_orig_mod.' prefix that torch.compile leaves on state dict keys (a real wart, acknowledged in a comment), rebuilds the GPT from the saved model_args, and picks its tokenizer by looking for the dataset's meta.pkl: character mappings if found, otherwise tiktoken's GPT-2 BPE. Generation is the naive loop in GPT.generate: crop the context to the last block_size tokens, forward the whole sequence, take the last position's logits, divide by temperature (default 0.8), zero out everything below the top-k cutoff (default 200), softmax, torch.multinomial, append, repeat. Note what is missing: there is no KV cache, so every step re-runs the full forward over the whole context. That honest inefficiency is exactly what makes the loop readable, and the entire engineering story of inference engines is what it costs to fix it, which is the vLLM chapter.

Part V: Internals deep dives

Deep dive 1: model.py as a curriculum

Read top to bottom, model.py teaches the transformer in five lessons. LayerNorm (the first class) exists only because PyTorch's has no bias=False switch; GPT-2 used biases, but the config lets you drop them, which the comment notes is slightly better and faster.

CausalSelfAttention is the file's core lesson, and the shape choreography is worth memorizing:

x: (B, T, C)                      B=batch, T=time, C=n_embd
c_attn(x): (B, T, 3C) ── one fused Linear computes q,k,v for all heads
split → q,k,v: (B, T, C)
view+transpose → (B, nh, T, hs)   nh=heads, hs=C/nh ── "heads" is a reshape
attention → (B, nh, T, hs)
transpose+view → (B, T, C) ── heads reassembled side by side
c_proj → (B, T, C)

When papers talk about heads, you can point at the reshape that creates them. The forward then branches: if torch.nn.functional.scaled_dot_product_attention exists, it dispatches to the fused flash kernel with is_causal=True; otherwise it runs the manual version, and having both in one screen is a gift, because the manual path is the semantics of the fused one: q @ k.T / sqrt(hs), mask the upper triangle to negative infinity with a precomputed tril buffer (registered, confusingly, under the name bias), softmax, dropout, multiply by v. My softmax page covers why the fused kernel can avoid materializing the (T, T) matrix at all. MLP is four lines of substance: expand to 4 * n_embd, GELU, project back, dropout. Block is the two-line pre-LN residual pattern, normalize-then-sublayer inside each residual branch, which is what keeps deep stacks trainable from initialization.

The GPT class holds the remaining lessons. GPTConfig defaults vocab size to 50,304, which is GPT-2's 50,257 padded up to a multiple of 64 purely for kernel efficiency, a classic example of hardware leaking into hyperparameters. Weight tying is one line, self.transformer.wte.weight = self.lm_head.weight, sharing the 38.6M-parameter embedding matrix with the output head. Initialization is normal with std 0.02 everywhere, except residual projections (c_proj.weight) get std 0.02 / sqrt(2 * n_layer) per the GPT-2 paper, so that the sum of 24 residual branch outputs does not grow with depth. The forward has an inference-only optimization worth noticing: with no targets, the lm_head is applied to just the final position, skipping 1023/1024 of the projection work. configure_optimizers encodes the standard decay rule by tensor shape, any parameter with p.dim() >= 2 (matmul weights, embeddings) gets weight decay 0.1 and everything else (biases, LayerNorm gains) gets none, and picks fused AdamW when available. from_pretrained maps Hugging Face GPT-2 checkpoints into this module, transposing the four weight matrices that OpenAI stored as Conv1D, and its shape asserts are a template for validating any reimplementation against a reference.

The famous trap in this file is parameter counting with tied weights. get_num_params reports the canonical "124M" figure (123.65M when loading GPT-2 weights) by subtracting only the position embeddings: the token embedding is genuinely also the output head, so it is counted once, not omitted and not double-counted. People who "fix" weight tying by cloning the matrix quietly add 38M parameters and break the GPT-2 correspondence.

Deep dive 2: train.py's real-world training practices

Each technique in train.py exists to kill a specific failure mode, and because nothing is hidden, the file works as an annotated checklist. Mixed precision: the forward runs under torch.amp.autocast in bfloat16 when the GPU supports it. Bfloat16 keeps float32's exponent range, so it needs no loss scaling and GradScaler(enabled=False) makes the scaler a no-op; float16 has a narrow exponent and small gradients underflow to zero, so the scaler multiplies the loss up before backward and unscales before the optimizer looks at gradients. The trap: the unscale must precede clipping, and train.py gets the order right in a way many hand-rolled loops do not. If gradients overflow, scaler.step skips that iteration entirely; occasional skipped steps in float16 logs are normal.

Gradient accumulation: 40 micro-batches of 12 sequences behave like one batch of 480 because gradients sum in place across backward() calls until zero_grad. The two supporting details are dividing the loss by the accumulation count (average, not sum) and, under DDP, syncing gradients only on the final micro-step. Clipping at global norm 1.0 caps the damage any single bad batch can do to the weights, cheap insurance against the loss spikes that plague long runs. The schedule is warmup plus cosine:

lr
6e-4 ┤        ╭─╮__
     │       ╱     ╲__
     │      ╱         ╲____
     │     ╱               ╲______
6e-5 ┤    ╱                       ╲________________
     └───┴────────────────────────┴────────────► iter
      0  2k (warmup)            600k (decay end, then floor)

Warmup protects the early steps when AdamW's second-moment estimates are garbage; cosine decay to a tenth of peak (both "per Chinchilla", say the comments) trades exploration for convergence. torch.compile is one opt-in line that the README credits with cutting iteration time from about 250 ms to 135 ms on the reference hardware, and TF32 is enabled explicitly for float32 matmuls. The PyTorch walkthrough covers what autocast, DDP buckets, and compile actually do beneath these calls.

The correction worth stating plainly: gradient accumulation is not a free-lunch batch size increase in wall-clock terms; it multiplies time per iteration by the number of micro-steps. It buys you the optimization behavior of a half-million-token batch on hardware that cannot hold one, and that is all it buys. Similarly, DDP is not a speedup knob for a model that already fits and saturates one GPU's compute: it scales tokens per second, not steps per second.

Deep dive 3: the data pipeline, bins and memmaps

The pipeline has exactly two moving parts: a prepare.py that runs once, and get_batch that runs forever. data/openwebtext/prepare.py downloads the dataset via Hugging Face datasets, carves off a 0.05% validation split, tokenizes every document with tiktoken's GPT-2 BPE (encode_ordinary, then an explicit end-of-text token, id 50256, appended per document), and concatenates all ids into flat binary files written through a memmap in 1,024 shards. The file's own comments record the outcome: train.bin is ~17 GB holding 9,035,582,198 tokens, val.bin ~8.5 MB holding 4,434,897. The dtype is uint16, which works only because GPT-2's max token id 50,256 fits under 2^16, halving disk and page-cache footprint versus int32. The character-level Shakespeare prepare is the same idea small enough to read in one breath, and it additionally pickles meta.pkl with the 65-entry stoi/itos tables, which is how train.py later discovers the vocab size and sample.py discovers the decoder.

documents ──tokenize──► ids + EOT │ ids + EOT │ ids + EOT │ ...
                                  └────── one flat uint16 stream ──────┘
train step: randint offsets ──► [i : i+1024] windows, anywhere at all

Training then treats the corpus as a featureless token stream. Sampling is with replacement at random offsets, so there are no epochs, no shuffling infrastructure, and no guarantee any given token is ever seen; "epoch" is simply not a concept this loop has, and for a 9B-token corpus consumed at 300B training tokens, expectation does the bookkeeping that dataloaders do elsewhere. Two consequences surprise people. Windows routinely straddle document boundaries, and the model simply learns that text after an EOT token starts fresh; there is no attention masking between documents, a simplification larger systems often revisit. And the random window start means the model trains at every position offset, which is what lets sample.py prompt at arbitrary lengths. The trap to correct: the bins are not compressed or structured, and that is the point; the memmap plus OS page cache is the entire I/O system, and it comfortably feeds half a million tokens per iteration because sequential-ish uint16 reads are nearly free next to the forward pass.

Part VI: Reading the repository

Stage 0, run before reading. Do the Part II tutorial: the pretrained sample, then the Shakespeare training run. Watch the printed config dump, the parameter count line (from GPT.__init__), and the loss trajectory. You should be able to answer: what files did prepare create? Why does the model print 10.65M parameters for the Shakespeare config when n_embd is 384 and n_layer is 6?

Stage 1, model.py, whole thing. The rare model file where nothing should be skipped: LayerNorm, CausalSelfAttention, MLP, Block, GPTConfig, then the GPT class including from_pretrained, configure_optimizers, estimate_mfu, and generate. You should be able to answer: what shape is the tensor between any two lines of the attention forward? Why is vocab_size 50304? Which parameters get weight decay and why? What would break if you removed weight tying?

Stage 2, train.py, whole thing. Read it as a checklist, treating each mechanism (DDP setup, autocast and the scaler, get_batch, the accumulation loop, clip and step, get_lr, estimate_loss, checkpointing) as an item to fully understand rather than boilerplate to skim. You should be able to answer: what is the exact order of scaler, clip, and step, and why? When do gradients cross the network in DDP? What is in ckpt.pt and why is optimizer state there?

Stage 3, the periphery. sample.py and data/shakespeare_char/prepare.py, then data/openwebtext/prepare.py, then configurator.py (all 47 lines, including the apology). You should be able to answer: how does sampling pick its tokenizer? Why is the memmap dtype uint16? What does the configurator's type assertion protect against?

Stage 4, quantitative extras. bench.py for isolated forward-backward timing, and the two notebooks, transformer_sizing.ipynb (parameter and FLOP accounting that reproduces estimate_mfu's formula) and scaling_laws.ipynb. You should be able to answer: where does 6N + 12·L·H·Q·T come from, and what MFU should you consider healthy on an A100?

Where not to start: do not start by forking and editing, the classic nanoGPT failure mode, because the repo is only trustworthy as a baseline if you understand what you changed; do not start with the DDP and AMP branches of train.py before reading the single-GPU float32 path through the same lines; and do not start in the notebooks, which assume you already know the code they measure.

Part VII: Hands-on labs

All labs run from the repo root; CPU is fine everywhere except lab 6.

Lab 1: prove the architecture against the reference (from_pretrained).

python sample.py --init_from=gpt2 --device=cpu \
    --start="What is the answer to life, the universe, and everything?" \
    --num_samples=3 --max_new_tokens=60

Observe the forced config line (vocab_size=50257, block_size=1024) and coherent GPT-2 text from weights loaded through the Conv1D transpose in model.py. If the reimplementation were wrong in any shape or wiring, the asserts in from_pretrained would have said so.

Lab 2: train end to end on CPU (the whole loop). Run the README's CPU flag set from Part II, watching three things: the tokens per iteration will be: 768 line (accumulation 1 × batch 12 × block 64, the formula from Stage 0 applied to this config), the eval lines every 250 iterations, and the final loss near 1.88 after 2,000 iterations. Then sample with python sample.py --out_dir=out-shakespeare-char --device=cpu and observe character-level pastiche. Loss values vary a little run to run; the shape of the curve should not.

Lab 3: open the bin files (data pipeline).

python - <<'EOF'
import numpy as np, pickle
m = np.memmap('data/shakespeare_char/train.bin', dtype=np.uint16, mode='r')
meta = pickle.load(open('data/shakespeare_char/meta.pkl', 'rb'))
print(len(m), "tokens")                      # 1003854
print(''.join(meta['itos'][int(i)] for i in m[:200]))
EOF

Observe the opening of the corpus ("First Citizen:") coming straight off disk. Then divide the corpus by the GPU config's 16,384 tokens per iteration (64 × 256): one corpus-worth every ~61 iterations, so 5,000 iterations pass over the data roughly 80 times in expectation, which is why this config overfits and its comments plan for it.

Lab 4: plot the schedule without running it (get_lr). Copy the three-case get_lr function out of train.py with the GPT-2 constants (warmup 2,000, decay end 600,000, lr 6e-4, min 6e-5) and print it at iterations 0, 1000, 2000, 150000, 301000, 600000, 700000. Observe roughly 3e-7 at step 0 (warmup starts near zero, not at peak), the peak at 2,000, about 3.3e-4 at the halfway point of decay, and the 6e-5 floor after 600,000. Sanity-checking a schedule by evaluation, before a run, is a habit that pays for itself the first time a warmup misconfiguration would have wasted a day.

Lab 5: gradient accumulation equivalence (the accumulation loop). Two short CPU runs on Shakespeare:

python train.py config/train_shakespeare_char.py --device=cpu --compile=False \
  --eval_iters=20 --block_size=64 --n_layer=4 --n_head=4 --n_embd=128 \
  --max_iters=300 --lr_decay_iters=300 --dropout=0.0 \
  --batch_size=32 --gradient_accumulation_steps=1

python train.py ...same flags... --batch_size=8 --gradient_accumulation_steps=4

Observe that both print the same tokens-per-iteration, that the loss curves track each other closely (not identically, the batches differ), and that the second run's iterations take roughly four times as long. That is the whole contract of accumulation, demonstrated.

Lab 6: what compile buys (torch.compile), GPU required. Run the Shakespeare GPU config twice, with --compile=False and then default, comparing the steady-state time ...ms per log line after the first few (the compile run pays a large one-time cost up front). Expect a meaningful drop, in the spirit of the README's 250 ms to 135 ms measurement for the GPT-2 config, with the exact ratio depending on GPU and PyTorch version. Watch the mfu percentage move in the same direction, and check it against estimate_mfu's formula.

Part VIII: Understanding checks

1. What is nanoGPT, in one sentence? A faithful GPT-2 model definition and an honest, complete pretraining loop in two roughly 330-line files, capable of reproducing GPT-2 (124M) on OpenWebText on one 8-GPU node in about four days.

2. Where do training targets come from? The input window shifted one token right: get_batch slices data[i : i+block_size] as x and data[i+1 : i+1+block_size] as y, and cross entropy over the vocabulary at every position is the entire objective.

3. Why is the vocab size 50,304 when GPT-2's tokenizer has 50,257 tokens? It is padded to the nearest multiple of 64 so the embedding and output projection dimensions are friendly to GPU kernels; the extra rows correspond to no real token and simply learn to be never predicted.

4. What does weight tying do here and what changes if you remove it? The token embedding matrix and the output head are the same tensor, saving about 38M parameters at this scale and matching GPT-2. Removing it unties input and output representations, inflates the parameter count, and breaks loading of GPT-2 checkpoints, which assume the tie.

5. Why do only some parameters get weight decay? configure_optimizers decays parameters with two or more dimensions (matmul weights and embeddings) and exempts biases and LayerNorm gains, because decaying those one-dimensional parameters regularizes nothing useful and measurably hurts; the shape-based rule is a compact encoding of that standard practice.

6. Walk through why the scaler-unscale-clip-step order matters in float16. The loss is scaled up before backward so small gradients survive float16's narrow range; clipping must act on true-scale gradients, so scaler.unscale_ comes first; then scaler.step checks for inf/nan and either applies AdamW or skips the step and reduces the scale, and scaler.update adapts the scale. Clipping before unscaling would compare a scaled norm against an unscaled threshold and effectively never clip, or always clip, depending on the scale.

7. Why does bfloat16 not need the scaler? It keeps float32's 8-bit exponent, trading mantissa precision instead, so gradient underflow is not the failure mode and the code constructs GradScaler(enabled=False), a no-op passthrough.

8. How does gradient accumulation interact with DDP in this code? The per-process accumulation count is the configured total divided by world size, keeping the global batch fixed, and require_backward_grad_sync is set true only on the last micro-step so the all-reduce happens once per iteration rather than once per micro-step, which is what no_sync() would do with more ceremony.

9. What is in ckpt.pt beyond weights, and why? Optimizer state (AdamW's per-parameter moment estimates), model_args, iter_num, best_val_loss, and the config. Without optimizer state, resumption would restart Adam's statistics from zero and the run would behave like a fresh warmup; without model_args, the checkpoint could not rebuild a matching module.

10. Why can training loss estimates not come from single batches? A single random 12×1024 window has high loss variance, so estimate_loss averages 200 batches per split in eval mode; checkpointing on a one-batch val loss would checkpoint on noise.

11. Why is there no KV cache in generate, and what does that cost? Simplicity: each generated token re-runs the full forward over the whole cropped context, so generation is quadratic-ish in output length instead of linear. It is fine for sampling a few hundred tokens and hopeless for serving, which is the problem inference engines solve with cached and paged attention state.

12. What does MFU measure, and what is the formula based on? Model FLOPs utilization: achieved FLOPs per second (estimated as 6N plus an attention term of 12·L·H·Q·T per token, from the PaLM paper's appendix) divided by the hardware's peak, 312 bfloat16 TFLOPS for an A100. It answers whether the loop is feeding the GPU efficiently, independent of model size.

13. Someone scales nanoGPT to a 7B model by renting more 8-GPU nodes. What happens and why? It fails at model construction or the first optimizer step with out-of-memory, because DDP replicates the full model, gradients, and both AdamW moments on every GPU; extra nodes add replicas, not capacity. They need sharded training (FSDP/ZeRO-style), which is a different codebase, not a bigger torchrun invocation.

14. When would you reach for nanoGPT versus Transformers' Trainer? nanoGPT when you need to see or modify the training mechanics themselves: research forks, teaching, small-scale pretraining where a 300-line loop you fully control beats a framework you partially understand. Trainer when you want fine-tuning with modern architectures, ecosystem integrations, and distributed execution handled for you, and the loop itself is not the object of study.

15. Why does the character-level Shakespeare run overfit, and how does the config acknowledge it? The corpus is about one million tokens and the model sees far more than that in expectation over 5,000 iterations, so validation loss bottoms out and climbs while train loss keeps falling. The config sets always_save_checkpoint = False so only val-improving checkpoints are kept, and adds dropout 0.2, which the pretraining configs set to zero.

16. What does the configurator's exec approach trade away, and for what? It gives up namespacing, static analysis, and safety (config files are executed code) in exchange for every default being an ordinary visible variable at the top of train.py and zero configuration machinery to learn. The comments call it a bad pattern to copy; it is the right trade only because this codebase is meant to be read and hacked.

Part IX: Design lessons

Make the pedagogical path and the production path the same code. The three-minute Shakespeare toy and the four-day GPT-2 reproduction are one script with different configs, so everything learned on the toy transfers. The same principle shows up in build systems and infrastructure-as-code: if the demo path and the real path diverge, the demo teaches the wrong system.

Precompute aggressively, then make runtime access dumb. Tokenization happens once in prepare.py; the hot loop reads flat bytes through the OS page cache. Moving work across the offline/online boundary is the same move as search indexing, compiled assets, and materialized views.

Degenerate gracefully to the simple case. DDP activates only if torchrun's environment variables exist, the scaler is a no-op under bfloat16, compile is one optional line, and the same file therefore serves one CPU and sixteen A100s. Feature-detection with a clean fallback beats parallel codepaths, in training loops as in systems software.

Encode policy in structure, not configuration. Weight decay applies to p.dim() >= 2; residual init scale is derived from n_layer; vocab is padded to hardware-friendly multiples. Rules stated as code over tensor properties cannot drift out of sync with the model the way a hand-maintained parameter list can.

Keep the honest version next to the fast version. The manual attention implementation sits a few lines from the flash dispatch, and naive generate ships beside the training loop. Readable reference implementations adjacent to optimized paths are how kernels, codecs, and crypto libraries stay debuggable, and it is why this repo can be validated against Hugging Face weights in a single method.

Part X: The memorization framework

The whole system in one sentence: tokenize a corpus once into a flat binary stream, sample random shifted windows from it forever, and push each batch through a plain GPT and a carefully ordered AMP-AdamW step until the cosine schedule runs out.

corpus → bin → batch → forward → loss → backward → clip → step → ckpt

The chain mapped to actual source:

corpus → bin     data/<dataset>/prepare.py   (tiktoken/char ids, uint16)
bin → batch      train.py get_batch           (np.memmap, randint, shift by 1)
batch → forward  model.py GPT.forward         (wte+wpe → Blocks → ln_f → lm_head)
forward → loss   model.py                     (F.cross_entropy, ignore_index=-1)
loss → backward  train.py micro-step loop     (autocast, scaler, DDP sync toggle)
backward → step  train.py                     (unscale → clip 1.0 → AdamW → update)
step → ckpt      train.py eval gate           (estimate_loss → ckpt.pt)

Memorize these blocks:

The GPT-2 recipe. 12 layers, 12 heads, 768 dim, block 1024, vocab 50304; batch 12 × accumulation 40 × block 1024 = 491,520 tokens per iteration; 600,000 iterations ≈ 300B tokens; AdamW lr 6e-4 → 6e-5, betas 0.9/0.95, weight decay 0.1, clip 1.0, warmup 2,000.

The AMP rules. bfloat16: autocast only, scaler disabled. float16: autocast plus GradScaler, and the order is scale, backward, unscale, clip, step, update.

The two one-liners that carry the theory. Targets are inputs shifted by one; weight tying is wte.weight = lm_head.weight.

The shape chant for attention. (B,T,C) → (B,T,3C) → three of (B,nh,T,hs) → (B,nh,T,hs) → (B,T,C).

The cost model. FLOPs per token ≈ 6N + 12·L·H·Q·T; MFU = achieved / 312 TFLOPS on A100 bfloat16; compile took the reference iteration from ~250 ms to ~135 ms.

Part XI: Papers and further reading

Nearly every line of these two files traces back to a small set of papers, and each one rewards a direct read. Where this site derives the same idea in depth, the companion link points there.

1. Vaswani et al., Attention Is All You Need, 2017. The transformer architecture that model.py implements in miniature, down to the multi-head reshape this chapter chants. The attention note on this site derives the math.
2. Radford et al., Language Models are Unsupervised Multitask Learners, 2019. The GPT-2 paper, source of the architecture, the residual init rule, and the 124M model the repo reproduces. The language models from scratch class builds the same model piece by piece.
3. Brown et al., Language Models are Few-Shot Learners, 2020. The GPT-3 paper whose training recipe the defaults follow, betas of 0.9 and 0.95, weight decay 0.1, warmup plus cosine decay.
4. Kaplan et al., Scaling Laws for Neural Language Models, 2020. The power laws relating loss to model size, data, and compute, which scaling_laws.ipynb reproduces at small scale.
5. Hoffmann et al., Training Compute-Optimal Large Language Models, 2022. The Chinchilla analysis the schedule comments cite for decaying to a tenth of the peak learning rate. The learning rate schedules note covers warmup plus cosine in general.
6. Loshchilov and Hutter, Decoupled Weight Decay Regularization, 2019. The AdamW optimizer that configure_optimizers builds, with decay applied by tensor shape. The Adam note works through the moment estimates.
7. Micikevicius et al., Mixed Precision Training, 2018. The loss-scaling recipe behind the scale-backward-unscale-clip-step order that Part IV walks through. The mixed precision note covers the number formats.
8. Li et al., PyTorch Distributed, Experiences on Accelerating Data Parallel Training, 2020. The DDP design whose bucketed all-reduce and no_sync behavior train.py toggles by hand, covered further in the PyTorch walkthrough.
9. Dao et al., FlashAttention, Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022. The fused kernel behind scaled_dot_product_attention, covered in the flash-attention walkthrough.
10. Sennrich et al., Neural Machine Translation of Rare Words with Subword Units, 2016. The byte pair encoding idea behind the GPT-2 tokenizer that prepare.py applies through tiktoken.
11. Chowdhery et al., PaLM, Scaling Language Modeling with Pathways, 2022. Its appendix supplies the FLOPs-per-token formula that estimate_mfu uses.