Anatomy of a Training Loop: nanochat Line by Line

By the end of this lesson you will be able to open a small pretraining codebase and find, without searching, the six places that matter: where the vocabulary is made, where the model’s size is decided, where text becomes batches, where the optimiser and its schedules live, where the evaluations are, and where generation happens. You will also be able to say what the numbers the training loop prints on every step mean, which is what makes the lab’s log readable rather than decorative.

The text we read is nanochat, nanochat main · verified 2026-09-08, which the course pins at main because the project does not cut releases. Everything quoted below was read from the repository on 2026-09-09; the file names are stable but the line contents move, so treat the quotes as a guide to where to look rather than as the current state of the code.

Clone it now and read alongside. The whole repository is small enough to hold in your head, which is the property that makes it worth reading at all.

RunnableAll tracks

get the code

git clone https://github.com/karpathy/nanochat.git ~/nanochat

The map

The README publishes its own file tree, and the annotations in it are the fastest orientation you will get. The part that matters is one package and one scripts directory.

Output — what you should see

nanochat
├── checkpoint_manager.py       # Save/Load model checkpoints
├── common.py                   # Misc small utilities, quality of life
├── core_eval.py                # Evaluates base model CORE score (DCLM paper)
├── dataloader.py               # Tokenizing Distributed Data Loader
├── dataset.py                  # Download/read utils for pretraining data
├── engine.py                   # Efficient model inference with KV Cache
├── gpt.py                      # The GPT nn.Module Transformer
├── loss_eval.py                # Evaluate bits per byte (instead of loss)
├── optim.py                    # AdamW + Muon optimizer, 1GPU and distributed
└── tokenizer.py                # BPE Tokenizer wrapper in style of GPT-4
scripts
├── base_eval.py                # Base model: CORE score, bits per byte, samples
├── base_train.py               # Base model: train
├── chat_cli.py                 # Chat model: talk to over CLI
├── tok_eval.py                 # Tokenizer: evaluate compression rate
└── tok_train.py                # Tokenizer: train it

Ten modules and five scripts is the entire pretraining path. The README’s own description of the design is worth taking seriously as an engineering position: “there are no giant configuration objects, model factories, or if-then-else monsters in the code base”.

What sits on what, reading upwards from the corpus

each layer only exists because the one below it produced a file

1. Corpus shardsParquet files with a text column, downloaded by nanochat/dataset.py. The last shard is held out as validation.on disk
2. Vocabularyscripts/tok_train.py trains byte-pair merges with rustbpe and saves a tiktoken encoding plus a table of bytes per token.frozen after this
3. Batchesnanochat/dataloader.py tokenises documents on the fly and packs them into rows that each begin with the beginning-of-sequence token.
4. Modelnanochat/gpt.py builds the transformer from one integer, the depth, with width, heads and horizon derived from it.
5. Optimiser and schedulesnanochat/optim.py runs Muon on matrices and AdamW on everything else; scripts/base_train.py schedules the learning rate, the momentum and the weight decay.
6. Evaluationsnanochat/loss_eval.py reports bits per byte on held-out text; nanochat/core_eval.py scores a benchmark bundle.
7. Engine and CLInanochat/engine.py generates with a KV cache; scripts/chat_cli.py drives it as a conversation once a chat model exists.what you talk to

Every arrow upwards is a file written to disk. That is why the pipeline can be run one stage at a time, and why a mistake in an early stage is invisible until a late one.

The tokeniser, and why it is first

scripts/tok_train.py is ninety lines and does one job. Its own summary is “Train a tokenizer using our own BPE Tokenizer library. In the style of GPT-4 tokenizer.” The defaults it exposes are the three decisions: --vocab-size at 32768, which its help text notes is 2^15; --max-chars at two billion, the amount of text the merges are learned from; and --doc-cap at ten thousand, the number of characters taken from any single document so that one enormous file cannot dominate the statistics.

nanochat/tokenizer.py explains the division of labour in its first line: “BPE Tokenizer in the style of GPT-4: train with rustbpe, inference with tiktoken.” Training a vocabulary and using one are different performance problems, and the file solves them with different libraries.

The regular expression that decides where text may be split before merging is worth reading closely, because it is a design decision most people never see:

Output — what you should see

# NOTE: this split pattern deviates from GPT-4 in that we use \p{N}{1,2} instead of \p{N}{1,3}
# I did this because I didn't want to "waste" too many tokens on numbers for smaller vocab sizes.
# I verified that 2 is the sweet spot for vocab size of 32K. 1 is a bit worse, 3 was worse still.

Numbers are allowed to merge into runs of at most two digits rather than three. At a vocabulary of thirty-two thousand, three-digit number tokens would consume entries that other text needs more. That is the whole vocabulary-size trade-off in one line: every entry you spend on one kind of text is an entry another kind does not get.

The last thing the script does is not obvious and matters for the rest of the run. It writes token_bytes.pt, a table of how many bytes each token id represents, with special tokens set to zero. The comment says why: it lets the run “report a loss that is invariant to the vocab size of the tokenizer”. We come back to that under evaluation.

Nine special tokens are declared in nanochat/tokenizer.py, and only the first of them is used during pretraining: <|bos|>, which “delimits documents”. The other eight mark user and assistant turns and Python tool calls, and the file says outright that they “are only used during finetuning to render Conversations into token ids”. They exist in the vocabulary of your base model and mean nothing to it.

The model, and its one dial

nanochat/gpt.py opens with a list of what is in the architecture, and reading it is a compact tour of what has become standard since GPT-2: “rotary embeddings (and no positional embeddings)”, “QK norm”, “untied weights for token embedding and lm_head”, “relu^2 activation in MLP”, “norm after token embedding”, “no learnable params in rmsnorm”, “no bias in linear layers”, “Group-Query Attention (GQA) support for more efficient inference”.

GPTConfig is seven fields: sequence length, vocabulary size, number of layers, query heads, key/value heads, embedding width, and a sliding-window pattern string whose characters mean full or quarter context per layer.

You never set most of those. scripts/base_train.py derives them from --depth:

Fragment — not complete on its own

base_dim = depth * args.aspect_ratio
model_dim = ((base_dim + args.head_dim - 1) // args.head_dim) * args.head_dim
num_heads = model_dim // args.head_dim

With the defaults, aspect ratio 64 and head dimension 128, asking for depth 20 gives a width of 1,280 and ten heads. The README states the design intent plainly: the depth “automatically determines all other hyperparameters (the width of the transformer, number of heads, learning rate adjustments, training horizons, weight decays, …) so that the trained model comes out compute optimal”, and “the user doesn’t have to think about or set any of this”.

✦Tip

One dial is a strong choice for a teaching codebase and a limitation for research. It is why the lab’s “go further” step is to change the depth and nothing else: every other knob you turn moves you off the curve the defaults were tuned on, and you will not know which change caused what.

One more class in gpt.py is worth understanding, because it explains why the run does not use PyTorch’s automatic mixed precision. Linear subclasses nn.Linear and casts on the way in: “Replaces autocast: master weights stay fp32 for optimizer precision, but matmuls run in the activation dtype”. The README’s precision table says the compute dtype defaults to bfloat16 on CUDA devices of compute capability 80 and above, and to float32 on CPU and MPS, with an environment variable to override it. On a Mac the note is specific: “On recent macOS, MPS also runs NANOCHAT_DTYPE=bfloat16 fine (~25% less memory, similar speed).”

The data loader, and the tokens it throws away

nanochat/dataset.py is download and read utilities. Shards are parquet files fetched on demand from a Hugging Face repository, and the split rule is one line: “split can be ‘train’ or ‘val’. the last parquet file will be val.” Your validation set is a whole shard that training never touches, which is exactly the discipline Part 1 taught with MNIST, at a different scale.

nanochat/dataloader.py is where a decision most tutorials skip is made explicit. Its docstring describes the packing strategy:

Output — what you should see

BOS-aligned bestfit:
   - Every row starts with BOS token
   - Documents packed using best-fit algorithm to minimize cropping
   - When no document fits remaining space, crops a document to fill exactly
   - 100% utilization (no padding), ~35% tokens cropped at T=2048

Read the last two lines together. Nothing is padded, so no compute is wasted on filler; but roughly a third of the tokens are discarded to make documents line up with row boundaries. The file states the reason for paying that price: it “ensures that there are fewer ‘confusing’ tokens in the train/val batches as every token can now attend back to the BOS token and sees the full context of the document”.

That is a real trade, and it is invisible in the training log. If you plan a run around a token budget, the budget you consume from the corpus is larger than the budget the model trains on.

The optimiser, and the schedules around it

nanochat/optim.py describes itself as “A nice and efficient mixed AdamW/Muon Combined Optimizer. Usually the embeddings and scalars go into AdamW, and the matrix parameters go into Muon”, and credits its origin: “Adapted from: https://github.com/KellerJordan/modded-nanogpt”. Two optimisers, split by what kind of tensor a parameter is, is not what most training scripts do, and it is the single biggest reason this codebase trains faster than a textbook one.

scripts/base_train.py wraps three schedules around it, all of which need to know the total number of steps before the first one:

• Learning rate: a linear warm-up over --warmup-steps (40 by default), a constant middle, and a linear warm-down over the last --warmdown-ratio of the run (0.65 by default) to a final fraction of the initial rate.
• Muon momentum: warms from 0.85 to 0.97 over the first four hundred steps, then decays to 0.90 through the warm-down.
• Weight decay: a cosine decay to zero over the course of training.

Because all three depend on the horizon, the horizon is not something you can stop early without changing what the run was. A run interrupted at eighty per cent is not the same as a run planned for eighty per cent of the length.

The horizon itself is chosen three ways, in a stated order of precedence: an explicit --num-iterations, a --target-flops, or, by default, a --target-param-data-ratio whose help text reads “calculate num_iterations to maintain data:param ratio (Chinchilla=20, -1 = disable)”. The default value in the argument parser is 12, and the reference speedrun script overrides it to 8 with the comment “slightly undertrained to beat GPT-2 => decrease data:params ratio from compute optimal 10.5 (default) to 8”. Notice that the comment and the parser disagree about what the default is, which is the ordinary state of a codebase under active tuning and a reminder to read the code rather than the comment when a number matters. Hold on to those figures; the next two lessons are about them.

Two derived quantities follow from the horizon, and both cite papers in the code. The total batch size is computed from the token horizon following the Power Lines paper, whose abstract states that optimal and critical batch sizes “scale as power laws in D, independent of model size, N”; the code comment gives the exponent it uses as approximately 0.383 and rounds the result to a power of two. The learning rates are then scaled by the square root of the batch-size ratio, and the weight decay is scaled to keep a quantity the code calls T_epoch constant. The comments are honest about the uncertainty: “these papers study AdamW, not Muon. We are blindly following AdamW theory for scaling hoping it ~works for Muon too.”

∴Research note

That admission is the most valuable line in the file. A recipe that works is not the same as a recipe that is understood, and a codebase that says which is which is teaching you something a polished framework would hide. When your own run behaves oddly at a batch size far from the reference, this is the first comment to reread.

The loop itself is the loop from Part 1. Accumulate gradients over micro-steps, step the optimiser, zero the gradients, log. Each step prints a line whose fields are worth learning to read: loss is an exponentially smoothed and debiased training loss, lrm the current learning-rate multiplier, dt the milliseconds the step took, tok/sec the throughput, bf16_mfu the fraction of the device’s peak that the step achieved, epoch how many times the loader has been round the corpus, and eta the projected time remaining, computed from the average step time after the first ten steps. The lab uses eta to fit the run into your afternoon and epoch to check that you are not training on the same text twice.

The two evaluations

nanochat/loss_eval.py implements the one that runs during training, and its docstring is the clearest explanation of the metric you will find:

Output — what you should see

Instead of the naive 'mean loss', this function returns the bits per byte (bpb),
which is a tokenization vocab size-independent metric, meaning you are still comparing
apples:apples if you change the vocab size.

The mechanism is to sum the loss in nats and independently sum the bytes the target tokens represent, then divide, converting to base two. A model with a bigger vocabulary predicts fewer, longer tokens and so gets a lower mean loss for free; bits per byte removes that advantage. Special tokens are excluded because their byte count was set to zero by the tokeniser trainer, which is why that file wrote the table.

This matters to you the moment you train two models with two different tokenisers, which is exactly what the project in this part asks you to do. Compare their mean losses and you are comparing vocabularies. Compare their bits per byte and you are comparing models.

The second evaluation is the CORE metric in nanochat/core_eval.py, scored against a bundle of in-context learning tasks downloaded on first use. It comes from DataComp-LM, whose abstract describes “a broad suite of 53 downstream evaluations”. It is expensive relative to a loss evaluation, which is why --core-metric-every exists and why the small-machine recipes disable it with -1 and run it once at the end through scripts/base_eval.py instead.

scripts/base_eval.py runs three things, selected by a comma-separated --eval: core, bpb and sample. Its own usage comment gives the single-GPU form, which is the one the lab uses: python -m scripts.base_eval --model-tag d24 --device-batch-size=16 --max-per-task=100 --split-tokens=524288.

Sampling during training is hard-coded to seven prompts, and they are chosen well: “The capital of France is”, “The chemical symbol of gold is”, “If yesterday was Friday, then tomorrow will be”, “The opposite of hot is”, “The planets of the solar system are:”, “My favorite color is”, “If 5*x + 3 = 13, then x is”. Each probes a different thing: a fact, a lookup, a small piece of reasoning, an antonym, a list, an open continuation, and arithmetic. Watching which of them stops being nonsense first tells you more about the run than the loss curve does.

Inference, and the wrapper you talk to

nanochat/engine.py is described in the tree as “Efficient model inference with KV Cache”, which is the same mechanism Part 6 measured in llama.cpp and Part 9 measured under load in vLLM, here in about as little code as it can be written in.

scripts/chat_cli.py is the conversation layer, and reading it shows exactly how thin that layer is. It fetches the special token ids for the user and assistant markers, keeps a list of tokens for the conversation, appends the user’s message between its markers, and generates until the assistant end token appears. That is the whole “chat” abstraction. Its default --source is sft, because the base model you produce in the lab has never been trained to respect those markers; talking to a base model means prompting it with a prefix and reading the continuation, which is what the lab’s sampling script does instead.

Where the neighbours fit

nanochat did not appear from nowhere, and its README credits the line it descends from. Three repositories are worth knowing about, and each answers a different question.

nanoGPT is the ancestor, described by its README as “The simplest, fastest repository for training/finetuning medium-sized GPTs”, with “train.py a ~300-line boilerplate training loop and model.py a ~300-line GPT model definition”. It covers pretraining only. Its README reports reproducing GPT-2 at 124M parameters on OpenWebText “on a single 8XA100 40GB node in about 4 days of training”. Its character-level Shakespeare example is still the shortest path to watching a transformer learn anything: the README reports about three minutes on one A100, and gives an explicit reduced configuration for a laptop, noting that on Apple silicon you should “add --device=mps”. If nanochat feels like too much at once, start there.

llm.c answers “what is actually being computed”: “LLMs in simple, pure C/CUDA with no need for 245MB of PyTorch or 107MB of cPython”, including “a simple reference CPU fp32 implementation in ~1,000 lines of clean code in one file”. Nothing clarifies what a framework is doing for you like reading the version that has no framework.

modded-nanoGPT answers “how fast can this get”. Its README describes a competition to “search for the fastest algorithm to use 8 NVIDIA H100 GPUs to train a language model that attains 3.28 cross-entropy loss on the FineWeb validation set”, and reports reaching that target “under 75 seconds on 8xH100” against forty-five minutes for the llm.c reproduction it took the target from. The Muon optimiser in nanochat came from there, as did the rotary embeddings, QK-norm and squared ReLU in the architecture list above. It is where techniques are proven before they become defaults elsewhere, and it is also a good demonstration of why a leaderboard needs a fixed target: without one, “faster” is meaningless.

Trace one update and one unit of progress

Write the loop in this order: obtain a batch, compute logits and loss, scale the loss for accumulation if needed, backpropagate, update at the accumulation boundary, advance the scheduler according to its contract and clear gradients. Evaluation and checkpointing occur at explicitly counted steps. Confusing microsteps with optimiser updates changes the meaning of a learning-rate schedule.

For example, accumulating four microbatches before an update means four forward/backward passes and one optimiser step. An epoch counter tracks dataset passes; a token counter tracks processed text; neither is interchangeable with update count when packing or batch size changes.

Check the first and last batch of an epoch. A partial final accumulation window may need different scaling or deliberate dropping. Padding should not alter the loss denominator unexpectedly. Compare a tiny controlled run with a hand-counted number of examples and tokens, then verify the logged counters. Correct accounting prevents a misleading comparison where one training recipe appears better because it quietly processed more data or took more updates.

Recap

nanochat is ten modules and five scripts, and each stage writes a file the next one reads. The tokeniser is trained first with rustbpe and used with tiktoken, its split pattern trades number tokens for everything else, and it writes a bytes-per-token table so that loss can be reported in vocabulary-independent units. The model is built from one integer, the depth, with width, heads and horizon derived from it, and precision is managed by an explicit cast in a custom linear layer rather than by autocast. The data loader packs documents so that every row begins at a document boundary, at the cost of cropping roughly a third of the tokens. Muon runs on the matrices and AdamW on everything else, wrapped in three schedules that all depend on the horizon being fixed in advance. Bits per byte is the metric to compare models across tokenisers; the CORE score is the benchmark, run rarely because it is expensive. nanoGPT, llm.c and modded-nanoGPT sit around it as the simpler version, the frameworkless version and the fast version.

Check your understanding

Question 1. You train two models on the same corpus, one with a 16k vocabulary and one with a 64k vocabulary, and compare their final mean cross-entropy losses. What is wrong with the comparison?

• Nothing, mean loss is the standard metric
• The larger vocabulary predicts fewer, longer tokens, so its mean loss is lower for reasons that have nothing to do with model quality; bits per byte normalises this away
• The losses are not comparable because the models have different depths
• Cross-entropy cannot be compared between runs at all

Show the answer and why

Answer: The larger vocabulary predicts fewer, longer tokens, so its mean loss is lower for reasons that have nothing to do with model quality; bits per byte normalises this away

This is exactly why nanochat/loss_eval.py exists and why the tokeniser trainer writes a table of bytes per token. Bits per byte divides the summed loss by the summed bytes of the target tokens, so a model that packs more bytes into each token gets no free advantage.

Question 2. The data loader's docstring says roughly 35% of tokens are cropped at a sequence length of 2048. You plan a run on a 100 million token budget. What does that tell you?

• That the model will train on about 135 million tokens
• That you need to read more of the corpus than the number of tokens the model trains on, because cropping happens when packing documents into rows
• That 35% of the corpus is corrupt
• That you should raise the sequence length to 4096 to avoid it

Show the answer and why

Answer: That you need to read more of the corpus than the number of tokens the model trains on, because cropping happens when packing documents into rows

Packing every row to start at a document boundary means the tail of a document that does not fit is discarded. The compute is fully utilised, but the corpus is consumed faster than the token counter suggests. A longer sequence length reduces the fraction cropped but costs more memory in attention, which is a separate trade.

Question 3. You interrupt a planned 5,000-step run at step 4,000 and keep the checkpoint. Why is that not the same model as a run planned for 4,000 steps?

• It is the same; only the number of steps matters
• The learning rate, Muon momentum and weight decay all follow schedules computed from the total number of steps, so an interrupted run stops partway through a warm-down that was aimed at a different endpoint
• The data loader would have shuffled differently
• The checkpoint would be corrupt

Show the answer and why

Answer: The learning rate, Muon momentum and weight decay all follow schedules computed from the total number of steps, so an interrupted run stops partway through a warm-down that was aimed at a different endpoint

The warm-down is a fraction of the planned horizon and the weight decay is a cosine over it. Stopping early leaves the model at a learning rate that has not finished decaying, which usually costs a visible amount of loss. Choose the horizon before you start, which is what the lab's calibration step is for.

Question 4. Which of these does the base model produced by pretraining actually respect?

• The user and assistant special tokens in the vocabulary
• A system prompt
• Only the beginning-of-sequence token, which delimits documents during pretraining; the conversation tokens exist in the vocabulary but nothing has trained the model to use them
• The chat template stored alongside the tokeniser

Show the answer and why

Answer: Only the beginning-of-sequence token, which delimits documents during pretraining; the conversation tokens exist in the vocabulary but nothing has trained the model to use them

nanochat/tokenizer.py says the conversation tokens "are only used during finetuning to render Conversations into token ids". They occupy vocabulary entries from the start, which is why supervised fine-tuning can teach them cheaply later, but the base model has never seen them in a meaningful position.

Question 5. The training log prints "epoch: 2" partway through your run. What has happened, and does it matter?

• Nothing; epochs are just a counter
• The data loader has been round the whole corpus and started again, so the model is seeing text it has already trained on; whether that matters depends on your goals, but it means your effective token budget is smaller than your step count suggests
• The learning rate schedule has restarted
• The validation shard has been included in training

Show the answer and why

Answer: The data loader has been round the whole corpus and started again, so the model is seeing text it has already trained on; whether that matters depends on your goals, but it means your effective token budget is smaller than your step count suggests

The loader iterates infinitely and counts cycles. Repeated data is learned more strongly, which is the deduplication argument from Part 3 arriving in your own run. If you see an epoch above one and did not intend it, download more shards rather than shortening the run.

Check my answersClear