13 — Pretraining a Small Model End-to-End (124M → 1B)¶

Why: every serious inference engineer must have pretrained something, once, from bytes to a checkpoint that generates coherent text. The lessons — tokenization economics, LR schedules, mixed precision, checkpointing, W&B hygiene, why your loss went NaN on step 4300 — do not transfer from reading. You must run it.

Target project: train a 124M–1B parameter GPT-style model on FineWeb-Edu for ~10–50B tokens, reach a validation loss competitive with Karpathy’s llm.c GPT-2 replication (~3.28 on the FineWeb val set is the community reference).

Realistic hardware: 1 H100 (rented, $2–3/hr) for 1B model over ~2–4 days; or 8×H100 for a fast run (see Phase 6.14 modded-nanogpt). Your 3090 can pretrain a 124M in ~1–2 days.

1. Choose your codebase¶

Codebase	Params	Framework	Best for
nanoGPT (Karpathy)	124M–350M	pure PyTorch, single-file	Start here. Read every line.
llm.c (Karpathy)	124M–1.5B (GPT-2)	C/CUDA + PyTorch	If you want to see the same run in C.
nanotron (HuggingFace)	1B–7B+	PyTorch + 5D parallelism	Once you’re comfortable, this is production-grade.
modded-nanogpt (K. Jordan)	124M	nanoGPT fork with Muon, FP8, speedrun optimizations	Phase 6.14 — do after your baseline.
picotron (HuggingFace)	124M	minimal 4D parallelism tutorial	Educational; read alongside Ultra-Scale Playbook.

Recommended path: nanoGPT baseline → nanotron reproduction → modded-nanogpt participation.

Github links:

2. The dataset: FineWeb-Edu¶

FineWeb-Edu (https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) is the current community standard: 1.3T tokens of educational-quality web text, filtered by a Llama-3-70B classifier. It punches well above older Common Crawl derivatives token-for-token.

Sizes:

sample-10BT — 10 billion tokens, ~20GB. This is your working set.
sample-100BT — 100 billion tokens.
Full — 1.3T tokens.

Preparation:

# Tokenize with GPT-2 or Llama tokenizer, save as .bin shards.
python prepare_fineweb_edu.py --sample 10BT --tokenizer gpt2

nanoGPT’s data/openwebtext/prepare.py is a template. For FineWeb-Edu:

from datasets import load_dataset
ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True)
# stream + tokenize + write to np.memmap in 100M-token shards

Discipline: keep a held-out val set the model never sees. 100M tokens is plenty. This is your only source of truth for progress.

3. Model config: the anatomy of a 124M¶

@dataclass
class GPTConfig:
    block_size: int = 1024      # sequence length
    vocab_size: int = 50304     # 50257 GPT-2 vocab padded to multiple of 64 for tensor cores
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0        # off for pretraining
    bias: bool = False          # LayerNorm+Linear biases off (Llama-style)

Parameter accounting (do this by hand):

Component	Formula	124M
Token embedding	vocab · n_embd	38.6M
Position embedding	block_size · n_embd	0.8M
Per-block attn (QKVO)	4 · n_embd²	2.36M
Per-block MLP (4× hidden)	8 · n_embd²	4.72M
Per-block layernorms	2 · n_embd	1.5K
Blocks total (12)	~ 7.08M · 12	85M
Final LN + head (tied w/ embed)	0 (tied)	0
Total		~124M

If your parameter count doesn’t match the checkpoint, your model is wrong. Print it before training.

4. The training loop, honest edition¶

# Precision
model = model.to(dtype=torch.bfloat16)   # bf16 throughout; no grad scaler needed
# (fp16 would require GradScaler; bf16 has enough exponent range)

# Optimizer
optim = torch.optim.AdamW(
    model.parameters(),
    lr=6e-4, betas=(0.9, 0.95), weight_decay=0.1,
    fused=True,   # fused kernel matters on H100
)

# LR schedule: warmup + cosine to 10%
def get_lr(step):
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    if step > max_steps:
        return min_lr
    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (max_lr - min_lr)

# Grad accumulation for effective batch = 0.5M tokens
accum_steps = 512 * 1024 // (batch_size * block_size)   # target 0.5M tokens

for step in range(max_steps):
    lr = get_lr(step)
    for pg in optim.param_groups: pg['lr'] = lr

    optim.zero_grad(set_to_none=True)
    for micro in range(accum_steps):
        x, y = next_batch()
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits, loss = model(x, y)
        loss = loss / accum_steps
        loss.backward()

    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optim.step()

    if step % 250 == 0:
        val_loss = evaluate(model, val_data)
        wandb.log({"train_loss": loss.item()*accum_steps, "val_loss": val_loss, "lr": lr, "step": step})

    if step % 1000 == 0:
        save_checkpoint(model, optim, step)

Non-negotiable hyperparameter choices (defensible defaults)¶

Knob	Value	Why
Precision	bf16 throughout	fp32 optimizer state, bf16 forward/back. No GradScaler needed.
Max LR	6e-4 for 124M	Empirical GPT-2/3 sweet spot. Scale inversely-ish with model size (Chinchilla-lineage).
Warmup	500–2000 steps	Prevents early divergence when init is far from stable regime.
Schedule	cosine → 10% of max	Kaplan/Chinchilla convention. Linear-decay is fine too.
Beta1, Beta2	0.9, 0.95	GPT-3 defaults; 0.999 for beta2 is too slow for LLMs.
Weight decay	0.1	Applied to matmul weights, NOT to bias/norm/embedding.
Grad clip	1.0 (global norm)	Non-optional for stability.
Effective batch	0.5M tokens	GPT-3 small-model convention. Grad-accumulate to reach it.
Sequence length	1024 (start), 2048 (later)	Longer is better for quality but expensive.
Tokens per param	20 (Chinchilla-optimal)	124M × 20 ≈ 2.5B tokens for compute-optimal, but train longer for a better model.

`torch.compile` — turn it on¶

model = torch.compile(model)   # ~30–50% speedup on H100

On first call it takes 60–120s to compile; from then on it’s free. If it crashes on dynamic shapes, set mode="reduce-overhead" or fall back to mode="default".

5. Weights & Biases (or the local equivalent)¶

pip install wandb
wandb login

Log these every 250 steps:

train_loss, val_loss (both bpc and natural loss)
lr
grad_norm (before clipping)
tokens/sec, MFU (6·N·B·S / (t · peak_flops))
per-parameter-group weight norm (diagnoses instabilities early)

Log these every checkpoint:

sample generations at fixed prompts — you’ll see quality improve; loss alone lies
gpu memory used
disk usage (checkpoints eat space fast)

W&B alternatives: TensorBoard (offline), MLflow (self-hostable), Aim (open source, similar UI to W&B).

6. Checkpointing that doesn’t lose your run¶

Rules learned the painful way:

Save every 1000–5000 steps. Include: model state_dict, optimizer state_dict, LR scheduler state, RNG state, step count, config hash. Miss any of these and you can’t resume.
Keep only last-3 + a checkpoint-every-10k. Full checkpoints for a 1B model are ~12GB. Disk explodes fast.

Use torch.save with a temp-file rename to survive mid-write crashes:

torch.save(state, tmp_path); os.replace(tmp_path, final_path)

For FSDP2, use torch.distributed.checkpoint (DCP) with async saving. Regular torch.save at scale will OOM your process.
Test your resume path on day 1. Kill the process, restart from checkpoint, verify loss picks up smoothly. If it doesn’t, your restore is broken and you’ll only discover it after a crash halfway through a 3-day run.

7. Common failure modes and their smells¶

Symptom	Likely cause	Fix
NaN loss at step 200	LR too high, no warmup, or fp16 without scaler	Add warmup, switch to bf16
Loss plateaus at ~10 forever	Wrong tokenizer / vocab / bos-eos	Print a decoded sample; verify tokenizer
Val loss diverges from train after 30% of run	Overfitting on small data, or wrong val set	Increase data; check val is not in train shard
Loss spike then recovery	Ordinary at scale (see DeepSeek-V3 report)	Ignore if it recovers; investigate if not
Loss spike no recovery	Corrupted batch, LR too high, bad init	Rewind to prior checkpoint, lower LR
Throughput drops 3x mid-run	Thermal throttle, disk contention, tokenizer bottleneck	Check GPU temps, `iostat`, tokenizer prefetch
GPU memory grows over time	Activation cache leak, W&B media buffer, dataloader worker leak	`torch.cuda.memory_snapshot()`, restart workers

8. Evaluation beyond loss¶

Loss is a proxy. At checkpoint time, run:

Perplexity on WikiText-2 / WikiText-103 — apples-to-apples with the literature.
HellaSwag, ARC-Easy, PIQA (via lm-evaluation-harness) — sanity checks; small models score just above random on these, but movement is meaningful.
Generation from fixed prompts — log a few paragraphs at each checkpoint. Grep them for fluency.
KL divergence from a reference model (e.g., GPT-2-124M) at matched steps — the local-community-favored sensitive metric.

lm-evaluation-harness: https://github.com/EleutherAI/lm-evaluation-harness

9. From single-GPU to multi-GPU¶

Once your single-GPU 124M run is clean, ladder up:

DDP on 2 GPUs — confirm you get ~1.9x throughput. If not, dataloader is your bottleneck.
FSDP2 on 4–8 GPUs — use fully_shard(model, mesh=...). Confirm memory drops linearly.
FSDP2 + TP (2D mesh) on 8 GPUs for a 1B model — mesh (dp=4, tp=2).
nanotron for a real 5D-parallel run of a 1.3B on 8 H100s. Read the config carefully.

See Phase 6.04 (FSDP/ZeRO) and 6.11 (Ultra-Scale Playbook) for the mechanics.

10. Cost budget¶

Config	Model	Tokens	Runtime	Cost @ $2.5/hr H100
1×3090	124M	10B	3–5 days	Owned
1×H100	124M	10B	~12 hrs	~$30
1×H100	350M	10B	~30 hrs	~$75
1×H100	1B	10B	~5 days	~$300
8×H100	1B	100B	~2 days	~$960
8×H100 (speedrun)	124M	~2B	~3 min	~$1.20

A 124M pretrain on a rented H100 for $30 is one of the cheapest inflection points in your career. Do it.

11. Exercises¶

Parameter-count derivation. Before running anything, write out the parameter count of your chosen config by hand and match it to sum(p.numel() for p in model.parameters()). Report the two numbers.
Predict-then-measure MFU. Predict your MFU: 6·N·B·S / (t·peak). Predict tokens/sec on H100 for 124M at batch=64, seq=1024. Measure. Report gap.
Run the 124M for 10B tokens. Log everything. Publish the W&B run as a shareable link.
Ablation: LR sweep. Rerun with LR ∈ {3e-4, 6e-4, 1.2e-3}. Overlay curves. Which was best and did it match the Chinchilla scaling prediction?
Ablation: batch size. Rerun with effective batch ∈ {128K, 512K, 2M} tokens. Overlay. Discuss the loss-vs-tokens tradeoff.
Resume test. Kill the process at step 3000; resume. Show the loss curve is continuous across the restart.

12. Deliverable¶

A public repo + write-up with:

Config + code that reproduces the run.
W&B report link showing the training curve, sample generations at each checkpoint, and a final eval table.
One-page “what surprised me” retrospective.

This is one of the highest-signal artifacts in your portfolio. Recruiters have to squint to distinguish two candidates who have “read the papers.” They don’t have to squint at someone who has pretrained a model.

13. References¶

Karpathy — nanoGPT: github.com/karpathy/nanoGPT
Karpathy — Let’s build GPT-2 (llm.c): github.com/karpathy/llm.c/discussions/677
HF FineWeb-Edu dataset card: huggingface.co/datasets/HuggingFaceFW/fineweb-edu
HF FineWeb blog: huggingface.co/spaces/HuggingFaceFW/blogpost-fineweb-v1
Chinchilla paper (Hoffmann et al., DeepMind, 2022) — arxiv:<phone_number_or_numberic_id_or_random_id_151>
GPT-3 paper (Brown et al., 2020) — arxiv:<phone_number_or_numberic_id_or_random_id_152>
lm-evaluation-harness: github.com/EleutherAI/lm-evaluation-harness
W&B: wandb.ai

Next: 14_modded_nanogpt.md — the speedrun community as an accelerator for modern training tricks.