Pretraining¶
Stage 1 teaches the model English from scratch with plain next-token prediction on TinyStories. There is no instruction or reward here — just "given these tokens, predict the next one," millions of times.
Watch the pretraining walkthrough
Stage summary
Script: pretrain.py ·
Produces: ckpt_base.pt ·
Run: python pretrain.py [minutes] (default 10)
The objective¶
Cross-entropy on the next token. At every position t, the target is the token
at t+1 — the input window shifted right by one. Loss starts near
ln(vocab) ≈ 9.2 (uniform guessing over 10,000 tokens) and falls toward ~2 as
the model learns the statistics of English.
Data¶
ensure_data() downloads four files on first run (and is idempotent, so later
stages reuse them):
| File | Role |
|---|---|
pretrain_train.txt |
TinyStories training text (capped at ~80 MB) |
pretrain_valid.txt |
TinyStories validation text |
instruct_train.txt |
TinyStories-Instruct (used for the tokenizer + later stages) |
instruct_valid.txt |
TinyStories-Instruct validation |
Large training files are streamed and truncated at the last <|endoftext|> so
they always end on a clean story boundary.
Building the tokenizer¶
Token frequencies are counted across both training corpora; the top 10,000
(those appearing at least MIN_FREQ = 2 times) become the vocabulary, with the
three specials reserved first. Ties are broken alphabetically so the vocab is
deterministic given the same data.
The canonical vocabulary is whatever gets saved into ckpt_base.pt — every
later stage loads it from there, so token ids keep their meaning across the whole
pipeline.
Encoding to disk¶
Each corpus is encoded to a uint16 .bin file and read back with np.memmap,
so the dataset stays on disk and each batch only touches
batch_size × block_size tokens of it. Encoding is skipped when a .bin is
already newer than both its source text and the vocab.
The training loop¶
The classic loop, wall-clock bounded:
while time.time() < deadline:
x, y = get_train_batch() # random windows from the .bin
opt.zero_grad(set_to_none=True)
_, loss = model(x, y) # mean next-token cross-entropy
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
scheduler.step()
Hyperparameters¶
| Knob | Value | Why |
|---|---|---|
| Batch size | 48 | windows of 256 tokens each |
| Learning rate | 6e-4 | peak, AdamW |
| Betas | (0.9, 0.95) | β₂ below the 0.999 default — standard for noisy LM gradients |
| Weight decay | 0.1 | matrices only (dim ≥ 2), not biases/LayerNorm |
| Grad clip | 1.0 | a rare bad batch can't blow up the weights |
| Schedule | warmup 100 → cosine | decays to lr/10 over 8000 iters |
Why warmup + cosine? Adam's per-parameter step sizes are unreliable for the first ~100 steps (its moment estimates are still mostly zero), so the LR ramps up from near-zero. Then cosine decay lets the late steps make small, careful updates.
Why decay only the matrices? Decaying biases and LayerNorm gains toward zero
fights normalization instead of regularizing, so those parameter groups get
weight_decay = 0.
Watching it learn¶
Every 500 iterations the script prints train/val loss and samples a short story, then checkpoints. Early samples are word salad; within a few thousand iterations you get recognizable (if surreal) TinyStories prose — fluency emerging in real time.
📄 Source: pretrain.py