learn_nanogpt

the file's own shape

The model is five pieces: RMSNorm, CausalSelfAttention, MLP, Block and GPT, with a GPTConfig dataclass holding every default. Read top to bottom, it goes from the smallest normalization up to the loop that generates tokens.

RMSNorm instead of LayerNorm

Instead of subtracting a mean, it divides by a root-mean-square. The comment explains the reasoning: it keeps activations at a stable scale so gradients do not vanish or explode, and the bias is optional because modern models usually leave it off.

attention: RoPE on top of grouped queries

One c_attn linear produces queries plus two groups of keys and values, sized by n_kv_head, that is grouped-query attention, where several query heads share one key-value head. Positions are handled by rotary embeddings applied to the new tokens only, which is what lets a cached forward pass append a single token without recomputing the whole context.

Two details worth keeping. qk_rmsnorm optionally normalizes queries and keys on their last dimension before the score is taken. And the causal rule is handed to scaled dot-product attention as is_causal=True for a plain forward pass; an explicit additive mask is built only when a KV cache makes the key length longer than the query length.

That comment is the exercise in the matrix below, drawn out token by token.

MLP: a gated block

The hidden width is two thirds of the classic four-times rule, rounded up to a multiple of 64, then doubled by the gate split, so the parameter count stays near a GELU MLP's while using fewer matrix multiplies.

the defaults it ships with

block_size

1024 positions of context

vocab_size

50304, the GPT-2 vocabulary padded to a multiple of 64

n_layer / n_head

12 and 12

n_kv_head

3, so twelve query heads share three key-value heads

bias

False, unlike GPT-2

rope_base

10000.0, with a rope_scale_factor for longer inference contexts

qk_rmsnorm

False by default, switchable per run

train.py overrides these with much smaller numbers for a CPU run, four layers, four heads, 256 wide, 128 positions at a time. That is the configuration this worksheet's figures are sized for.

one file that runs on one GPU or on a cluster

The script's opening comment is a usage note: a single GPU in debug mode, or torchrun for distributed data parallel across one or many nodes. The distributed branch is ordinary nanoGPT machinery, ranks, world size, gradient accumulation divided across processes, and it is the part the author says is still unfamiliar, so this page describes the shape rather than teaching it.

everything is a variable at the top

Every setting is a plain assignment, and configurator.py is executed over them afterwards. That is why the same file accepts --batch_size=4 on the command line or a whole Python config file: both are just overrides to those globals.

Because tokens_per_iter = gradient_accumulation_steps * world_size * batch_size * block_size, that configuration moves 128 tokens per weight update. The script prints that number at startup.

batches come from a memory-mapped file

get_batch opens data/train.bin or data/val.bin as a uint16 memory map and slices random windows. The target is the input shifted one position forward, which is the whole language-modeling objective: given these tokens, predict the next one. The map is recreated per batch, on purpose, to avoid a known memory leak.

the rest of the loop, briefly

Estimate validation loss every eval_interval iterations, clip gradients at 1.0, use AdamW with weight decay only on matrices, and save the best checkpoint to saved_model/ckpt.pt. The model also offers an estimate_mfu helper that reports model-FLOPS utilization against an A100's peak, which is why training on a CPU prints numbers that look alarming.

from a checkpoint to text

The script loads saved_model/ckpt.pt, rebuilds the model from the arguments stored inside it, and loads the weights. If data/meta.pkl exists it builds the encoder from those character maps; otherwise it falls back to the GPT-2 encoding from tiktoken. Either way, encode turns text into ids and decode turns ids back into text.

The same prompt is repeated into a batch of four, so one forward pass produces four independent samples from identical starting text.

the sampling knobs

max_new_tokens

100 tokens per sample

temperature

0.8, applied before filtering; exactly 0 means greedy argmax

top_k

10 largest logits survive, the rest become -inf

top_p

0.95 nucleus cutoff, applied after top-k

min_new_tokens

10, during which the end token is banned by setting its logit to -inf

eos_id

the GPT-2 end-of-text id, so samples can stop early

num_samples

10 batches, each of four sequences

Note the order in generate: temperature divides the logits, top_k masks, then top_p masks again, and only then does F.softmax turn whatever is left into a probability distribution. A row of the attention matrix and a row of the sampling distribution obey exactly the same softmax rule, which is why both appear on this page.

how a finished sequence is trimmed

When the end token appears, the sequence is cut there and the batch is padded out to the longest surviving sample, with lengths returned alongside. That is what return_lengths is for, and it is why sample.py prints one sequence at a time rather than fixed-width rows.

Attention takes one row of the matrix, the scores between one query position and every key position, and softmaxes it into weights that sum to one. Causality says a position may only attend to itself and to what came before. Pick a query position to see its row.

prepare, train, inspect, sample

python data/prepareeos.py

python train.py --batch_size=4

python saved_model/inspect_ckpt.py

python sample.py
python chat.py

github.com/mnsky-tyan/learn_nanogpt holds the files, and Karpathy's nanoGPT is the repository this one started from.