Train a small LM, end to end.

Training a language model means showing a program billions of words of text and, every time it guesses the next word wrong, nudging its numbers a tiny bit so it guesses better next time. You have built the model itself, the arrangement of numbers that does the guessing, several times. What you have never done is the whole job, end to end: choose the text, chop it up, feed it in fast enough, keep the nudging going all night, save your progress, and then check whether what came out is any good.

That whole job matters more than it sounds. Almost every idea you will want to test later, a new twist on the model’s design, a cleaner pile of text, a small helper model that speeds up the server on your Spark, gets judged by comparing two training runs. A broken run does not crash. It just quietly scores worse, and a good idea looks like a bad one. So what you are really building here is a run you can trust, plus the habits that catch the silent mistakes in the first five minutes instead of the fifth hour.

This pulls together a lot of what came before: Tokenizers, fp32, bf16, fp16, fp8, Why training needs so much memory and the crash-safe save from Files and disks. And it becomes the yardstick for the rest of the track: Data a model can eat changes only the text, Train your own draft model reuses this loop, and Comparing transformer variants finally gives your own ideas a fair test.

Today

• Run the whole thing once, for real: raw web text, chopped into word pieces and saved to disk as numbers, fed in batches to a 124-million-number GPT for a night on the Spark, then saved, sampled from and scored.
• Know the handful of training settings that decide whether a run works (how big each step is and how that changes over the run, a gentle ramp-up at the start, which numbers get pulled toward zero, a cap on any single wild step, how much text goes into each step) and why each one is there.
• Read the chart of the model’s error over time like a doctor reads a chart: know what the error should be before training even starts, and tell a healthy run from one whose steps are too big, too small, missing the ramp-up, memorising instead of learning, or fed by a buggy data pipeline.
• Estimate how long a run will take before you start it, from just the model’s size and the amount of text, then measure what fraction of the chip’s top speed you actually used. Plus the rule of thumb for how much text a model of a given size deserves: about 20 word pieces per number in the model.
• Have a debugging routine that catches most broken runs in the first five minutes instead of the fifth hour.

You need to have written a GPT-2 forward pass before, which you have, several times. It helps, but is not required, to have done Tokenizers, fp32, bf16, fp16, fp8, Why training needs so much memory, and the crash-safe save from Files and disks. Plan on two evenings of building, one overnight run and one morning of evaluating.

The Problem

You have built GPT-2 from scratch more than once. Attention, the causal mask, RoPE in the Llama version: all fine. And yet every time you tried something of your own, a new attention variant, a different normalisation, a trick from a paper, the result was the same. The loss went down a bit, the samples were mush, and you could not tell whether your idea was bad or your training was.

That second possibility is the whole problem. A training run has maybe twenty settings that each, if wrong, quietly cost you a factor of two. A learning rate 10× too low does not crash; it just makes a model that looks dumb. Forgetting warmup does not crash; it costs you a few percent. Feeding the model the wrong targets does not crash; it gives you a beautiful loss curve. None of these throw an exception. All of them make your creative idea look like it failed.

People who get good results are not mostly more creative. They have a baseline that works, a run they trust, and they change one thing at a time against it. They also have a checklist that catches the silent bugs before they burn a night of compute. This lesson builds both: one trusted, end-to-end training run of a small language model on the Spark, plus the habits around it.

The Idea

The whole loop on one page

Here is everything that happens between “a pile of web pages” and “a model that writes”.

You have built the blue box in the middle many times. This lesson is about every other box, and about the arrows, because that is where runs die quietly.

Choosing the text

The model can only become as good as what it reads, so the first real decision is the dataset. Two good choices, for two different jobs:

• TinyStories (Eldan and Li, 2023). About two million short children’s stories, written by GPT-3.5 and GPT-4 with a deliberately tiny vocabulary of words. A few hundred million tokens in total (check the exact count on the dataset card). Its superpower is fast feedback: a model of only a few million parameters learns to write coherent little stories in minutes, so you can tell a working pipeline from a broken one before your coffee cools. Its weakness: it only teaches you about TinyStories. Nothing you learn about hyperparameters transfers cleanly to real text.
• FineWeb-Edu (Penedo et al., 2024, from Hugging Face). Web pages from Common Crawl, filtered by a classifier trained to spot “educational” content. It has ready-made samples, including sample-10BT, a random 10-billion-token slice. This is what llm.c and build-nanogpt trained on, which is the second superpower: there are published curves to compare against. When your run is at step 2,000, you can check whether a known-good run was at roughly the same loss.

Recommendation: FineWeb-Edu sample-10BT for the real run, TinyStories for the five-minute smoke tests. You are not trying to make the best 124M model in the world. You are trying to make a run you can trust, and the fastest way to trust a run is to have someone else’s run to hold it up against. TinyStories is your unit test; FineWeb-Edu is your integration test.

You will not use all 10B tokens. Below you will compute that about 2.5B is the sensible budget for a 124M model, which is a single night on the Spark.

Tokenize once, store as uint16

Tokenizing is slow, Python-heavy work (see Tokenizers for what BPE is actually doing). Training wants to read millions of tokens a second. So you never tokenize during training. You tokenize once, up front, and write the token IDs to disk as one flat array of integers. During training you just read slices of that array.

Three decisions make this file right:

1. One integer type, the smallest that fits. GPT-2’s tokenizer has 50,257 tokens, so every ID is between 0 and 50,256. An unsigned 16-bit integer holds 0 to 65,535. So each token takes 2 bytes as a uint16. 2.5B tokens is 5 GB on disk. As int64, the default for torch.tensor, it would be 20 GB. (If you ever switch to a tokenizer with more than 65,536 tokens, like Llama 3’s 128K vocabulary, you need uint32, and a silent cast to uint16 will wrap the big IDs around to garbage. Assert the max before you cast.)
2. A separator between documents. Put GPT-2’s special end-of-text tokenend-of-text tokenToken 50256, written <|endoftext|>. It marks the boundary between two documents so the model learns that what follows is unrelated to what came before. in front of every document. Then concatenate everything, ignoring document boundaries. A training window of 1,024 tokens will often span the end of one page and the start of the next; the separator tells the model “context resets here”. This is the simplest form of packingpackingConcatenating many documents into one long token stream and cutting it into fixed-length windows, so no compute is wasted on padding., and it is what GPT-2 and nanoGPT do.
3. Read it with a memory map. np.memmap(path, dtype=np.uint16, mode="r") gives you an array-shaped view of the file without reading it. Slicing it reads just those bytes, through the page cache, on demand. That is the mmap from Virtual memory and Files and disks: the OS keeps hot pages in RAM, and with 128 GB of unified memory on the Spark, your whole 5 GB dataset will end up cached after the first pass.

The loader: random windows, a held-out split

The loader’s whole job is: given a batch size B and a context length T, hand back two integer tensors x and y of shape (B, T), where y is x shifted by one. Three details make it correct.

Hold out validation data before anything else. Validation data has to be text the model never trains on, or the validation loss means nothing. The cleanest split is by file: the first shard is val, the rest are train. Never split by sampling random windows from one big file into “train” and “val”; windows overlap, so the same sentences end up in both.

Cover the data, don’t just sample it. The tempting loader picks B random start positions each step. That works, but it samples with replacement.

Keep the integers small until the last moment. Slice uint16 from the memmap, convert to int64 (PyTorch’s embedding lookup wants long), then move to the GPU. On the Spark the CPU and GPU share the same physical memory, but a .to("cuda") still goes through the CUDA allocator; at these sizes it is noise either way.

The model: GPT-2 small, with two small changes

Use the GPT-2 124M configuration: 12 layers, 12 heads, model width 768, context 1,024. You know this model. Two details worth getting exactly right:

• Tie the embeddings. The input embedding tableembedding tableA big matrix with one row per token. Looking up a token ID returns its row: the vector the model starts from. and the output layer (the “LM head”) share one matrix. That saves 38.6M parameters and is what GPT-2 did. With tying, the model has 124,475,904 parameters (I counted them with PyTorch on my Mac with the configuration below); 85M of those are in the transformer blocks, the rest are the embedding tables.
• Pad the vocabulary to 50,304. That is 50,257 rounded up to a multiple of 128. The extra 47 rows are never used as targets, so the model learns to give them near-zero probability. Why bother? GPU matrix-multiply hardware works in tiles, and dimensions that are multiples of 64 or 128 fill the tiles exactly. Karpathy measured a few percent speedup from this one change. It is your first taste of the theme of the speed section: the numbers the hardware likes are not the numbers the model needs.

Initialise like GPT-2: linear and embedding weights from a normal distribution with standard deviation 0.02, biases at zero, and the two projections that write back into the residual stream (attention output, MLP output) scaled down by 1/√(2 × layers), so the residual stream does not grow with depth.

What the loss number means

The loss is cross-entropycross-entropyFor each position, minus the log of the probability the model gave to the correct next token. Averaged over all positions in the batch. Zero only if the model is certain and right., measured in natsnatsThe unit you get when you use the natural log (ln). One nat is about 1.44 bits.. If the model gives the correct next token probability 0.05, that position contributes −ln(0.05) ≈ 3.0.

This is the single cheapest sanity check in deep learning, and it catches broken inits, wrong vocab sizes, mismatched loss reductions, and labels from the wrong tokenizer.

To make loss numbers feel like something, convert them to perplexityperplexitye raised to the loss. Read it as: “the model is as uncertain as if it were choosing uniformly among this many tokens.”:

Two warnings. Losses are only comparable on the same tokenizer and the same validation text. A loss of 3.3 on FineWeb-Edu and 1.5 on TinyStories are not comparing models, they are comparing datasets (stories are far more predictable). And small differences matter: at this scale, a change that drops validation loss by 0.02 nats is a real improvement if it is consistent, and a change that costs 0.05 is a real regression.

The optimizer: AdamW, with weight decay on the right things

AdamWAdamWAdam with “decoupled” weight decay: the weights shrink a little each step, separately from the gradient update. Adam itself keeps a running average of each weight’s gradient (momentum) and of its square, and divides one by the square root of the other, so every weight gets its own step size. is the default for transformers, and the GPT-3 settings are a good place to stand: betas (0.9, 0.95), epsilon 1e-8, weight decay 0.1. The 0.95 (instead of Adam’s usual 0.999) makes the squared-gradient average forget faster, which helps when gradient sizes swing around, as they do in language models.

Weight decayWeight decayShrinking every weight slightly toward zero each step. A gentle pressure toward small weights that acts as regularisation. is where people get it subtly wrong. Apply it to the 2-D weight matrices (every linear layer and the embedding table), and not to biases or LayerNorm gains and offsets. Decaying a LayerNorm gain pulls it toward zero, which literally scales down that layer’s output; there is no reason to want that. The rule that implements this is one line: decay a parameter if p.dim() >= 2. For the 124M model that puts 124,354,560 parameters in the decay group and 121,344 (biases and norms) in the no-decay group, which I also counted. If your no-decay group is empty, or holds tens of millions of parameters, something is off.

The learning rate: warm up, then cosine down

The learning rate is the most important number in the run. Get it within a factor of two of good and most other settings stop mattering much.

The shape that works: linear warmup from near zero to the peak over a few hundred steps, then cosine decay down to about 10% of the peak by the last step.

Why each part is there:

• Warmup. At step 0 Adam’s running averages are empty, so its first few steps are oddly sized, and the gradients of a random network point in wild directions. A full-size step on top of that can throw the model somewhere it never fully recovers from. A few hundred steps of ramping is cheap insurance. GPT-3 Small warmed up over 375M tokens; for a 2.5B-token run, a few hundred steps of 0.5M tokens is in the same spirit.
• Peak. 6e-4 is GPT-3 Small’s peak and a safe starting point for 124M. Karpathy found 124M models tolerate about 3× that. You will find your own with a sweep.
• Decay. Early on, big steps make fast progress. Late in training the loss surface is a narrow valley, and big steps bounce off its walls. Lowering the LR lets the model settle into the valley floor. You will see this as an extra dip in the last fifth of the run.

Batch size is measured in tokens

For language models, think of batch size as tokens per optimizer step, not sequences. GPT-3 Small used about 0.5M tokens per step; nanoGPT uses 2¹⁹ = 524,288 (a power of two near 0.5M). With T = 1,024 that is 512 sequences per step.

That will not fit in one forward pass on the Spark, and it does not need to. Gradient accumulationGradient accumulationRunning several small forward/backward passes and adding up their gradients before taking one optimizer step. Mathematically the same as one big batch, just slower per step and lighter on memory. splits the big batch into micro-batches: run, say, 32 micro-batches of 16 × 1,024 tokens each, summing gradients, then step once.

Why not just use a small batch and skip the accumulation? Because the batch size is a hyperparameter of the recipe. The LR of 6e-4 was tuned for 0.5M-token batches. Smaller batches give noisier gradients and want a smaller LR; if you change one, you have changed the other. Keeping the batch at 0.5M tokens means published numbers still apply to you.

Gradient clipping

Every so often a batch produces a gradient far larger than usual: a weird document, a rare token, an unlucky moment. One giant step can undo hours of progress. Gradient clippingGradient clippingBefore the optimizer step, measure the size (L2 norm) of the whole gradient, all parameters together. If it is bigger than a threshold, scale it down to the threshold. Direction kept, size capped. at a norm of 1.0 caps the damage. torch.nn.utils.clip_grad_norm_ returns the norm before clipping; log it every step. A healthy run’s gradient norm settles somewhere below 1 after the early phase and drifts slowly. Spikes in the norm come before spikes in the loss, so it is your early warning.

Making it fast: four switches

None of these change what the model learns (to within noise). All of them change how many tokens per second you get.

1. bf16 autocast. Wrap the forward pass and loss in torch.autocast(device_type="cuda", dtype=torch.bfloat16). Matrix multiplies run in bf16bf16A 16-bit float with the same exponent range as fp32 but fewer digits of precision. Big range means no overflow/underflow trouble, so no loss scaling is needed (unlike fp16). on the tensor cores; the weights, the optimizer state, and sensitive operations like softmax and the loss stay in fp32. You do not need a gradient scaler with bf16. That is fp16’s problem, and why fp32, bf16, fp16, fp8 says everyone switched.
2. TF32 for whatever is left in fp32. torch.set_float32_matmul_precision("high") lets fp32 matmuls use tensor cores at reduced precision. Free speed, harmless for training.
3. Flash attention. Use F.scaled_dot_product_attention(q, k, v, is_causal=True) instead of writing out softmax(QKᵀ)V by hand. PyTorch picks a fused kernel that never writes the T × T attention matrix to memory, which is the whole story of FlashAttention.
4. torch.compile(model). Traces your model and fuses chains of small operations (add, LayerNorm, GELU, dropout) into single kernels. On the Spark this matters more than on an H100: its memory bandwidth (273 GB/s, spec sheet) is low relative to its compute, so every unfused elementwise op that reads and writes a whole activation tensor is relatively expensive. The first step takes a minute or more while it compiles; the steps after are faster. How much faster is one of the things you will measure.

Also pass fused=True to AdamW if your PyTorch version supports it on your device: the optimizer update becomes one kernel instead of a loop over parameter tensors.

Reading loss curves

Now the part that would have saved your earlier experiments. Every training run produces a loss curve, and the shape tells you far more than the final number. Below are the six curves you will meet most often. They are shaped from the real behaviour of runs like this one, but they are illustrations, not measurements. Click through them all and read each annotation.

The patterns to burn in:

• The first number is ln(V). If step 0 is not near 10.8, stop.
• Too good to be true is a bug. A 124M model on web text does not reach loss 1. If you see it, you are leaking the answer. I checked this on my Mac with a tiny GPT fed random tokens: a correct pipeline cannot do better than ln(V) on random data (there is nothing to learn), yet with y = x instead of the shifted target, the loss fell from 10.5 to 5.8 in 100 steps. Random-token training is a two-minute test for leaks.
• Spikes that recover are a warning; spikes that don’t are a verdict. Lower the LR, check clipping, look at the gradient norm.
• Still falling steeply at the end means you stopped too early or the LR is too low: compute left on the table.
• Val pulling away from train is overfitting. In a single-epoch pretraining run it should not happen, so if it does, suspect the data (too small, or replayed).

Two more habits. Evaluate validation loss on a fixed set of batches (say 20 batches, the same ones every time), so the val number moves only when the model changes. And always plot new runs on top of your baseline, same axes. Most “did my idea work?” questions are answered by whether the new curve is below the old one at the same step.

How big, how long: scaling laws and 6·N·D

How many tokens should a 124M model see? In 2022, DeepMind’s Chinchilla paper (Hoffmann et al.) trained over 400 models of different sizes on different amounts of data, all for different compute budgets, and asked: for a fixed compute budget, what split between model size and data gives the lowest loss? The answer, roughly: grow them together, with about 20 tokens per parameter. Earlier work (Kaplan et al., 2020) had said to favour bigger models; Chinchilla’s correction was that most big models of the time were badly undertrained.

So for 124M parameters: about 2.5B tokens. That is the “compute-optimal” budget: the cheapest way to train to a given loss. (Plenty of real models go far past 20 tokens per parameter on purpose: Llama 3 8B saw over 15 trillion tokens, because a small model trained long is cheaper to serve. That trade is about inference cost, not training.)

Now the compute. Training a transformer costs about 6 FLOPs per parameter per token:

• The forward pass does one multiply and one add per weight per token: 2·N FLOPs per token.
• The backward pass computes two things for every layer, the gradient with respect to its inputs (to pass back) and with respect to its weights (to update). Each costs about as much as the forward: 4·N.

Total training compute ≈ 6·N·D, where D is the number of training tokens. This ignores the attention scores themselves (QKᵀ and the weighted sum of V), which add about 12 × layers × width × T FLOPs per token. For our model at T = 1,024 that is 113M on top of 6N = 747M, about 15% more, so a fairer count is ~860M FLOPs per token. Both numbers below; neither is exact.

For 124M parameters and 2.5B tokens: 6 × 1.245e8 × 2.5e9 ≈ 1.9e18 FLOPs (2.2e18 with attention).

What the Spark can do: about 125 TFLOP/s of dense bf16 at most. That is not a published spec; it is the estimate derived in Bandwidth bound from NVIDIA’s marketing number, and you will measure the real one. No training run gets the peak. The fraction it does get is the MFUMFUModel FLOPs utilization: the FLOPs your model’s math actually needs per second (tokens/s × FLOPs per token), divided by the hardware’s peak. Counts only useful model math, not recomputation or overhead. Introduced in Google’s PaLM paper.. Well-tuned large-scale runs reach about 40-60%; llm.c’s GPT-2 run on 8 A100s works out to roughly 60% by this formula (10B tokens in 90 minutes is about 1.85M tokens/s). For a single small model on the Spark, 30-45% would be a good result. I do not know your number; that is why you will measure it.

At 40% MFU, 50 TFLOP/s achieved: 1.9e18 / 5e13 ≈ 37,000 s ≈ 10 hours (about 12 with attention). At 25% MFU, more like 17-19 hours. That is an overnight run, and it tells you something important before you write any code: you get about one full-size run per day. So everything you want to learn about hyperparameters has to be learned on smaller, faster runs first.

Play with the budget. Then try the reverse panel with the tokens/second you measure later.

The debugging recipe

Karpathy’s “A Recipe for Training Neural Networks” (2019) boils down to: neural nets fail silently, so be paranoid, and build up from things you can verify. For a language model, in this order, before any long run:

1. Look at the data. Decode a few training windows back to text and read them. Is it text? Are there EOT tokens between documents? Print one (x, y) pair side by side and check y is x moved left by one.
2. Check the init loss. One forward pass, no training. ≈ ln(V)? (Above.)
3. Overfit one batch. Take a single small batch and train on it, over and over, for a few hundred steps. A working model and optimizer must drive the loss to near zero, because memorising 256 tokens is trivial. I ran this on my Mac with a 4-layer GPT on one batch of 4 × 64 tokens at LR 1e-3: loss went from 10.87 to 0.08 in 50 steps and 0.005 by step 300. If yours plateaus at 5, something between the loss and the weights is broken: wrong targets, detached tensors, an optimizer that was given the wrong parameters, an LR of zero.
4. Train on random tokens for a minute. The loss must stay at or above ln(V). If it drops, the model can see its targets.
5. Print the LR every step for the first few. Is warmup happening? Is it the number you think? (Off-by-one-order-of-magnitude LRs are the most common silent bug.)
6. A short real run on TinyStories or a small FineWeb slice. Does the curve look like the “good run” in Figure 005? Do samples get better?
7. Only then, the long run. And keep notes.

Checkpoints

A 10-hour run will be interrupted at some point: a reboot, an OOM, a Ctrl-C you regret. A checkpoint is everything needed to continue as if nothing happened:

• the model weights (from the uncompiled model; torch.compile wraps it and prefixes every key with _orig_mod.),
• the optimizer state (Adam’s two running averages, twice the size of the model),
• the step number (so the LR schedule resumes in the right place),
• the data position (epoch, shuffle seed, index into the window list), so you don’t retrain on the same data,
• the config, and ideally the git commit of your code.

Write it with the crash-safe recipe from Files and disks: write to a temp file in the same directory, fsync, os.replace over the old one, fsync the directory. A checkpoint that is half-written when the machine dies is worse than none, because it overwrote the good one. Keep the last two or three, plus a few milestones, so a bad stretch of training doesn’t leave you with only bad checkpoints.

How big? Weights in fp32: 124.5M × 4 bytes ≈ 0.5 GB. Adam state: 1.0 GB. About 1.5 GB per checkpoint. (See Why training needs so much memory for where those multiples come from.)

Sampling: turning probabilities into text

The model outputs a probability for every token. How you pick from them changes the text a lot, and none of it changes the model.

• Greedy: always take the most likely token. Deterministic, and for small models falls into loops quickly (“the the the”, or one sentence forever).
• Temperature: divide the logits by τ before the softmax. τ < 1 sharpens the distribution (safer, duller); τ > 1 flattens it (more surprising, more nonsense). τ = 1 is the model’s own distribution.
• Top-k: keep only the k most likely tokens, renormalise, sample. Cuts off the long tail of very unlikely tokens, each individually rare but collectively likely to be picked once every few words.
• Top-p (nucleus sampling, Holtzman et al., 2019): keep the smallest set of tokens whose probabilities add to at least p, say 0.9. Adapts to context: when the model is sure, the set is two tokens; when it is unsure, it is hundreds.

For a 124M model, temperature 0.8 to 1.0 with top-p 0.9 or top-k 50 is a reasonable start. Fix a prompt and a seed, so samples from different checkpoints are comparable.

Small evals

Validation loss is your main number: cheap, smooth, and directly what you optimised. But loss does not tell you whether the model knows anything, so add one small task eval.

HellaSwag (Zellers et al., 2019) gives a short context and four possible endings; one is right, three are machine-generated traps. A language model can do it without any instruction following: compute the model’s average loss on each ending’s tokens given the context, and pick the ending with the lowest loss. Random guessing scores 25%. Karpathy measured OpenAI’s GPT-2 124M at about 29.5% on the 10,042-example validation set this way; a 124M model is barely above chance, which is exactly why it is a good small eval: it moves as you train. The standard error on 10,042 examples is about ±0.45 points, so differences under about one point between two runs are noise (more on that in Statistics for evals).

Build It

Everything runs on the Spark except where noted. Keep one folder, ~/tinylm/, under git. Commit before every run, so every result in your log points at the exact code that made it.

pip install tiktoken datasets numpy matplotlib
python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name())"

# prepare.py
import os, numpy as np, tiktoken, multiprocessing as mp
from datasets import load_dataset

OUT = os.path.expanduser("~/data/fineweb_edu")
SHARD = 100_000_000            # tokens per shard = 200 MB of uint16
BUDGET = 2_700_000_000         # val shard + ~2.6B train tokens
enc = tiktoken.get_encoding("gpt2")
EOT = enc.eot_token            # 50256

def tok(doc):
    ids = np.array([EOT] + enc.encode_ordinary(doc["text"]), dtype=np.int64)
    assert ids.max() < 2**16, "token id does not fit in uint16"
    return ids.astype(np.uint16)

if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT",
                      split="train", streaming=True)
    buf = np.empty(SHARD, dtype=np.uint16)
    n = shard = total = 0
    with mp.Pool(max(1, os.cpu_count() - 2)) as pool:
        for ids in pool.imap(tok, ds, chunksize=64):
            while len(ids):
                take = min(SHARD - n, len(ids))
                buf[n:n + take] = ids[:take]
                n += take
                ids = ids[take:]
                if n == SHARD:
                    split = "val" if shard == 0 else "train"
                    buf.tofile(f"{OUT}/{split}_{shard:04d}.bin")
                    total += n
                    shard += 1
                    n = 0
                    print(f"wrote shard {shard}, {total / 1e9:.2f}B tokens", flush=True)
            if total >= BUDGET:
                break

class Loader:
    def __init__(self, data_dir: str, split: str, B: int, T: int, seed: int = 0):
        """Memmap every {split}_*.bin shard (np.memmap, dtype=np.uint16, mode='r').
        Build the list of (shard, start) window positions, stride T, each window
        needing T+1 tokens. Shuffle it with a seeded generator."""

    def next_batch(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Return x, y: int64 tensors of shape (B, T) on the GPU. y is x shifted by one.
        Walk the shuffled window list; when it runs out, bump the epoch and reshuffle."""

    def state_dict(self) -> dict: ...      # epoch, position, seed: enough to resume exactly
    def load_state_dict(self, d: dict): ...

# train.py (shape, not a solution)
@dataclass
class TrainConfig:
    data_dir: str = "~/data/fineweb_edu"
    total_batch: int = 524_288        # tokens per optimizer step
    B: int = 16; T: int = 1024        # micro-batch
    max_lr: float = 6e-4; min_lr_frac: float = 0.1
    warmup_steps: int = 300; max_steps: int = 4_800   # ~2.5B tokens
    weight_decay: float = 0.1; grad_clip: float = 1.0
    eval_every: int = 250; ckpt_every: int = 500; compile: bool = True

def get_lr(step: int, cfg: TrainConfig) -> float:
    """Linear warmup to max_lr over warmup_steps, then cosine to max_lr*min_lr_frac at max_steps."""

def configure_optimizer(model, cfg) -> torch.optim.Optimizer:
    """Two param groups: dim >= 2 gets weight_decay, the rest gets 0.
    Print both group sizes. AdamW, betas=(0.9, 0.95), eps=1e-8, fused if available."""

def save_checkpoint(path, raw_model, opt, step, loader, cfg): ...   # atomic: tmp, fsync, os.replace, fsync dir

# main loop, per step:
#   set lr on every param group from get_lr(step)
#   for micro in range(grad_accum):  (grad_accum = total_batch // (B*T); assert it divides)
#       x, y = loader.next_batch()
#       with torch.autocast("cuda", dtype=torch.bfloat16): _, loss = model(x, y)
#       (loss / grad_accum).backward()
#   norm = clip_grad_norm_(model.parameters(), cfg.grad_clip)
#   opt.step(); opt.zero_grad(set_to_none=True)
#   torch.cuda.synchronize(); time the step; log

step 1200 | loss 3.8123 | lr 5.81e-04 | norm 0.42 | dt 9012ms | tok/s 58,170 | mfu 39.9%

tmux new -s train
cd ~/tinylm && git commit -am "main run: 124M, 2.5B tok, lr 6e-4" && \
python train.py 2>&1 | tee runs/main/stdout.log
# detach: Ctrl-b d      reattach: tmux attach -t train

## 2026-10-18 main-124m
commit: a1b2c3d   config: 124M, 2.5B tok, B=16x1024x32, lr 6e-4, warmup 300
hypothesis: matches build-nanogpt's curve at equal tokens, within ~0.05
started: 22:10   expected: ~10-12 h at 40% MFU
result: (fill in) final val, hellaswag, tok/s, MFU, anything weird and at which step

Measure

Fill this in as you go. The “expect” column is my best guess from the arithmetic above and from published reproductions; where I can’t know (anything Spark-specific), it says so.

What to look for:

• The MFU row is the honest number. If it is under 20%, the GPU is waiting on something: data loading (watch htop for one pegged CPU core), Python overhead (micro-batch too small), or unfused elementwise ops (compile off). Each has a different fix. If it is over 60%, check your FLOP count and your timer; you probably forgot a synchronize().
• Your val loss vs. the published curve at the same number of tokens. If you are within a few hundredths, you have a trusted baseline. If you are 0.2 worse, something in the recipe differs, and finding it will teach you more than any new idea would have.
• Samples vs. loss. Read the step-500, step-2,000 and final samples next to their val losses. This is how you build the instinct for “what does 4.0 look like?” that the perplexity table above only describes.

Break It

Each of these is a short run (a few hundred steps of the 30M model is enough) plotted on the same axes as a good run. The point is to see the failure once, so you recognise it the next time it happens by accident.

1. Unshifted targets. Set y = x in the loader. The loss dives toward zero and both train and val look superb. Then sample from it: garbage. Lesson: the loss measures what you told it to measure; nothing more. Leaking validation data into training is the same bug, slower.
2. No warmup, 5× the LR. Set warmup_steps=0 and multiply max_lr by 5. Watch the first 100 steps on a log axis and watch the gradient-norm column: the norm spikes before the loss does. Then turn clipping off as well and see whether it diverges completely. Lesson: warmup and clipping exist for the first and worst moments of the run.
3. Forget to divide by grad_accum. Use loss.backward() instead of (loss / grad_accum).backward().

4. Replay one shard. Make the loader reshuffle and restart after the first shard (100M tokens) instead of moving to the next one, and run the 30M model long enough for several passes. You will see the “tiny data” curve from Figure 005: train keeps dropping with a cliff at each pass, val turns up. Lesson: in pretraining, rising val loss almost always means a data-pipeline bug, not a model problem.

Explain It Back

Answer out loud, in plain words, as if in an interview. Then check.

Go Deeper

• Andrej Karpathy, nanoGPT and build-nanogpt, with the video “Let’s reproduce GPT-2 (124M)” (2024). The reference implementation for everything in this lesson. The video is four hours and every minute earns its place, especially the parts about vocab padding, the weight-decay groups and the HellaSwag eval.
• Andrej Karpathy, llm.c, and its discussion post “Reproducing GPT-2 (124M) in llm.c in 90 minutes for $20”. The same training in C/CUDA, with the loss and HellaSwag curves you should compare yours against.
• Andrej Karpathy, “A Recipe for Training Neural Networks” (2019). The debugging philosophy: neural nets fail silently, so build up from things you can verify. Short, and it will change how you start every project.
• Hoffmann et al., “Training Compute-Optimal Large Language Models” (DeepMind, 2022, arXiv 2203.15556). Chinchilla. Read the figures and the first two sections for where “20 tokens per parameter” comes from, and how they fit it three different ways.
• Kaplan et al., “Scaling Laws for Neural Language Models” (OpenAI, 2020, arXiv 2001.08361). The original power laws, and the clearest derivation of the 6N FLOPs-per-token rule (their Section 2.1).
• Penedo et al., “The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale” (Hugging Face, 2024, arXiv 2406.17557), and Eldan and Li, “TinyStories: How Small Can Language Models Be and Still Speak Coherent English?” (2023, arXiv 2305.07759). Where your two datasets come from, and good reading for Data a model can eat.
• Holtzman et al., “The Curious Case of Neural Text Degeneration” (2019, arXiv 1904.09751). Why greedy and pure sampling both fail, and the paper that introduced top-p.
• Chowdhery et al., “PaLM: Scaling Language Modeling with Pathways” (Google, 2022, arXiv 2204.02311). Appendix B defines MFU and gives the FLOPs-per-token formula with the attention term, the one used in this lesson.

Connections

This lesson cashes in several earlier ones. The uint16 memmap is Virtual memory and Files and disks in practice, and the checkpoint uses the latter’s crash-safe save. bf16 autocast is the payoff of fp32, bf16, fp16, fp8; the checkpoint and optimizer sizes come from Why training needs so much memory; the attention kernel is FlashAttention; the peak FLOP/s and the “why compile matters more on the Spark” argument come from Bandwidth bound; the tokenizer is the one you built in Tokenizers.

Looking ahead, this run is the baseline for everything else in the track. Data a model can eat changes only the data and measures the effect against this curve. Train your own draft model reuses this loop with a distillation loss. Comparing transformer variants is where your creative ideas finally get a fair test: one change at a time, plotted against a run you trust. And Distributed training is what you reach for when one night per run stops being enough.