DL S053 · Training a 100M-Param Model End-to-End

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 53 / 80 · Module M09 · ~2 hours

The story we're starting with

This is the session everything so far has been building toward. Fifty-two sessions of tensors, attention math, tokenizers, data cleaning, scaling laws, GQA, RoPE, DDP, and GPU rentals — all so we can point a Python file at a rented A100 and watch it become a language model.

Here's what "training a 100M-param foundation model" actually looks like from the outside: a python train.py command, a wandb dashboard that ticks a loss number down for ~3 hours, and a .pt file at the end that produces something-like-English when you sample from it. From the inside it's a training loop, a scheduler, checkpointing, gradient accumulation, careful init, a few dozen hyperparameters, and the constant discipline of not touching things while they're working.

We're going to build the loop end-to-end, launch it, watch it, and interpret every early-warning signal. By the end you'll have a functioning train.py you can adapt to any size model, and you'll have watched (in your head or in real dollars) a real training run finish.

Prerequisites

• S050 (architecture) — you have LlamaLike model class.
• S052 (renting GPU) — you have an A100 rented and code synced.
• S046+S049 (data) — you have .bin shards ready to memmap.

1 · The full config

Everything in one place. This is the file you commit and never touch mid-run:

# configs/100m.py
from dataclasses import dataclass
 
@dataclass
class TrainConfig:
    # Model
    vocab_size: int = 32_000
    d_model:    int = 768
    n_layers:   int = 12
    n_heads:    int = 12
    n_kv_heads: int = 4
    hidden_dim: int = 2048
    max_seq_len:int = 1024
    rope_theta: float = 10_000.0
    norm_eps:   float = 1e-6
    tie_embed:  bool = True
 
    # Training
    total_tokens:  int = 2_000_000_000       # Chinchilla-optimal for 100M
    micro_batch:   int = 32
    accum_steps:   int = 8                    # → effective batch 256
    seq_len:       int = 1024                 # 256 × 1024 = 262144 tok/step
    lr:            float = 6e-4
    min_lr:        float = 6e-5               # 10% of peak, cosine floor
    warmup_steps:  int = 500
    weight_decay:  float = 0.1
    grad_clip:     float = 1.0
    beta1:         float = 0.9
    beta2:         float = 0.95
 
    # Data
    train_shards:  str = "/workspace/data/train-*.bin"
    val_shards:    str = "/workspace/data/val-*.bin"
    val_every:     int = 500
    val_batches:   int = 100
 
    # Checkpointing
    ckpt_dir:      str = "/workspace/checkpoints"
    ckpt_every:    int = 500
    ckpt_keep:     int = 3                    # rolling window
 
    # Logging
    wandb_project: str = "dl-llms-100m"
    log_every:     int = 10
 
    # Precision
    dtype: str = "bfloat16"
    compile: bool = True                       # torch.compile — free 20-30% speedup

Total training steps: 2e9 / (256 × 1024) ≈ 7630 steps.

Notice the LR of 6e-4. That's high — GPT-3 175B used 6e-5. But small models tolerate much higher LR; the Chinchilla paper suggests LR ≈ (1e-3) × sqrt(1e6 / N), giving 3-4e-4 for 100M. We push to 6e-4 with warmup for faster convergence.

2 · The training loop, annotated

Here's the whole thing, split into digestible chunks.

2.1 Setup — model, optimizer, scheduler

import torch, torch.nn as nn, os, math, time, wandb
from configs.hundred_m import TrainConfig
from model import LlamaLike
 
cfg = TrainConfig()
torch.manual_seed(42)
device = "cuda"
dtype = getattr(torch, cfg.dtype)
 
model = LlamaLike(cfg).to(device)
if cfg.compile:
    model = torch.compile(model)
print(f"Params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M")
 
# Adam with weight decay applied only to 2D+ tensors (matmul weights),
# never to norms or biases. Standard trick.
decay_params = [p for p in model.parameters() if p.dim() >= 2]
nodecay_params = [p for p in model.parameters() if p.dim() < 2]
optimizer = torch.optim.AdamW([
    {"params": decay_params,   "weight_decay": cfg.weight_decay},
    {"params": nodecay_params, "weight_decay": 0.0},
], lr=cfg.lr, betas=(cfg.beta1, cfg.beta2), fused=True)

The fused=True on AdamW is a free 10% speedup on A100+ — it fuses the parameter update into a single kernel launch.

2.2 Learning rate schedule — warmup + cosine

def get_lr(step: int, cfg: TrainConfig, total_steps: int) -> float:
    if step < cfg.warmup_steps:
        return cfg.lr * (step + 1) / cfg.warmup_steps
    if step > total_steps:
        return cfg.min_lr
    progress = (step - cfg.warmup_steps) / (total_steps - cfg.warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
    return cfg.min_lr + coeff * (cfg.lr - cfg.min_lr)

Warmup for 500 steps prevents the "loss spike on step 3" bug that hits every initialization if LR starts hot. Cosine decay from peak to 10% peak follows Chinchilla's exact schedule.

2.3 Data loading — memmap the shards

import glob, numpy as np
 
class ShardedDataset:
    def __init__(self, glob_pattern: str, seq_len: int):
        self.shards = sorted(glob.glob(glob_pattern))
        self.seq_len = seq_len
        self.rng = np.random.default_rng(int(os.environ.get("RANK", "0")))
 
    def get_batch(self, batch_size: int) -> tuple[torch.Tensor, torch.Tensor]:
        shard = self.rng.choice(self.shards)
        arr = np.memmap(shard, dtype=np.uint16, mode="r")
        idx = self.rng.integers(0, len(arr) - self.seq_len - 1, size=batch_size)
        x = np.stack([arr[i:i+self.seq_len].astype(np.int64) for i in idx])
        y = np.stack([arr[i+1:i+self.seq_len+1].astype(np.int64) for i in idx])
        return (torch.from_numpy(x).pin_memory().to(device, non_blocking=True),
                torch.from_numpy(y).pin_memory().to(device, non_blocking=True))
 
train_ds = ShardedDataset(cfg.train_shards, cfg.seq_len)
val_ds   = ShardedDataset(cfg.val_shards, cfg.seq_len)

Deliberately minimal — no PyTorch DataLoader overhead. Direct memmap + numpy sampling is 2× faster than DataLoader for this shape of data.

2.4 The main loop

total_steps = cfg.total_tokens // (cfg.micro_batch * cfg.accum_steps * cfg.seq_len)
wandb.init(project=cfg.wandb_project, config=vars(cfg))
 
start_step = maybe_load_checkpoint(model, optimizer)
 
model.train()
t0 = time.time()
for step in range(start_step, total_steps):
    # Set LR
    lr = get_lr(step, cfg, total_steps)
    for g in optimizer.param_groups:
        g["lr"] = lr
 
    # Gradient accumulation
    optimizer.zero_grad(set_to_none=True)
    total_loss = 0
    for micro in range(cfg.accum_steps):
        x, y = train_ds.get_batch(cfg.micro_batch)
        with torch.autocast(device_type="cuda", dtype=dtype):
            logits = model(x)
            loss = nn.functional.cross_entropy(
                logits.view(-1, cfg.vocab_size), y.view(-1)
            ) / cfg.accum_steps
        loss.backward()
        total_loss += loss.item()
 
    # Clip + step
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
    optimizer.step()
 
    # Log
    if step % cfg.log_every == 0:
        dt = time.time() - t0
        tok_per_sec = cfg.micro_batch * cfg.accum_steps * cfg.seq_len / dt * cfg.log_every
        wandb.log({
            "loss": total_loss,
            "lr": lr,
            "grad_norm": grad_norm.item(),
            "tok_per_sec": tok_per_sec,
            "step": step,
        })
        t0 = time.time()
 
    # Validate
    if step > 0 and step % cfg.val_every == 0:
        val_loss = evaluate(model, val_ds, cfg.val_batches)
        wandb.log({"val_loss": val_loss, "step": step})
        model.train()
 
    # Checkpoint
    if step > 0 and step % cfg.ckpt_every == 0:
        save_checkpoint(model, optimizer, step, cfg)
 
# Final save
save_checkpoint(model, optimizer, total_steps, cfg, final=True)
open(f"{cfg.ckpt_dir}/DONE", "w").close()

That's it. The whole loop is ~50 lines. Everything not shown (evaluate, save_checkpoint, maybe_load_checkpoint) is standard boilerplate covered in S055.

3 · The pre-flight checklist (2 minutes, saves you 3 hours)

Before you type python train.py, walk through this:

The overfit-one-batch test is the highest-value 30 seconds you can spend before a training run. It catches ~80% of the bugs that would silently ruin an 8-hour run.

4 · What a healthy loss curve looks like

Steps 1–100: loss drops from ~10.4 (random init, = log(32000)) to ~7.5. This is the "the model discovers that letters exist" phase. Very steep.

Steps 100–1000: loss drops from 7.5 to ~5.5. The "grammar and common words" phase. Steep but decelerating.

Steps 1000–3000: loss drops from 5.5 to ~4.5. The "syntax + collocations" phase. Curve visibly bending.

Steps 3000–7000: loss drops from 4.5 to ~3.8. The "long-range structure and facts" phase. Bending gets gentle.

Final loss around 3.7–3.9 for a 100M model at 2B tokens on a mixed English corpus. Anything below 3.5 and you're either overfitting or evaluating on training data. Anything above 4.5 and something is wrong.

4.1 The classic pathologies

We diagnose each of these in S056. For now: watch for smooth monotonic drop in log-space. Anything else, pause.

5 · Numbers to expect for a 100M A100 run

At MFU 40%, bf16, seq_len 1024, effective batch 256:

• Throughput: ~24000 tokens/sec on an A100 40GB.
• Steps to 2B tokens: ~7630.
• Wall clock: ~2.7 hours.
• **Cost @ $1.10/hr:** ~$3.
• Final val loss: 3.7–3.9.
• Final val perplexity: ~40–50 (= exp(loss)).

Sanity check your throughput: at step 100 open wandb, look at tok_per_sec. If it's <15000, something is off — check compile enabled, dtype is bf16, and micro_batch isn't so small that overhead dominates.

6 · Sampling — does it actually work?

Halfway through, sample from it to see what the model has learned:

model.eval()
prompt = torch.tensor(tokenizer.encode("The president visited"), device=device).unsqueeze(0)
with torch.no_grad(), torch.autocast(device_type="cuda", dtype=dtype):
    out = model.generate(prompt, max_new_tokens=50, temperature=0.8, top_k=50)
print(tokenizer.decode(out[0].tolist()))

Expected outputs at different stages:

• Step 500: "The president visited a the on and to of the a . the of the of a of the on the" — token-frequency baseline, no structure yet.
• Step 2000: "The president visited the White House on Tuesday for a meeting with Prime Minister of the country ." — grammar emerging, semantic bleed.
• Step 5000: "The president visited the White House on Tuesday to discuss trade with the Japanese prime minister ." — plausible short sentences.
• Step 7500 (done): "The president visited the White House on Tuesday to discuss trade policies with the Japanese Prime Minister Fumio Kishida ." — coherent short paragraphs. Not GPT-4. But real.

You have a language model. It's small, it hallucinates freely, but it's yours.

7 · What to do while it runs

Two hours forty minutes is longer than you'd like to just stare at a graph. My personal routine:

• First 30 min: watch wandb constantly. Anything weird → kill and diagnose.
• Next 60 min: step away. Eat, walk. Check phone-wandb every 15 min.
• Last 60 min: back at the desk. Start writing the eval script for S054. Sample every 500 steps and diff.
• Final 5 min: be present. Manually run the "am I really done" checklist as DONE file appears.

8 · War stories

9 · The 2025 story — what "good" looks like at 100M

When your training script has been running for 90 minutes and the loss has plateaued at 3.2, the question is: is that good? The answer requires knowing what other people got.

9.1 · The NanoGPT speedrun leaderboard

KellerJordan/modded-nanogpt on GitHub is an ongoing community speedrun: train GPT-2 (124M) on FineWeb to validation loss 3.28 as fast as possible on 8×H100. As of mid-2025 the record has fallen from Karpathy's original ~45 minutes (May 2024) to under 3.5 minutes (May 2025), a 12× speedup driven by:

• Muon optimizer (~20% faster convergence than AdamW on this task).
• FP8 mixed precision via Transformer Engine (~30% throughput on H100).
• Value residual / U-net attention shortcut — skip connections between attention layers.
• Untied embeddings with a shared projection matrix (contradicts your 100M config — at very small scale, tied hurts).
• Various dataloader / tokenization micro-optims.

Even if you don't adopt any of these, following the leaderboard for 3 months teaches you where the current perf frontier is and what the recent research consensus is on "what actually matters for training small LLMs."

9.2 · Reference val-loss numbers at 100M

If your run lands at val loss 3.3–3.6 on FineWeb-Edu after 2B tokens: normal, ship it. If it lands at 2.9–3.2: you probably over-trained or your data has quality leakage from your val set (check with S049's dedup). If it lands above 4.0: you have a bug — go to S056.

9.3 · Perplexity intuition at this scale

Val loss 3.4 = perplexity exp(3.4) ≈ 30. Meaning: on average, the model is as confused about the next token as if it were choosing uniformly among 30 candidates from your 32k vocab. Not smart. It cannot follow instructions, cannot answer questions, cannot count. What it can do: generate syntactically plausible English, complete common phrases, produce weakly-coherent code fragments. That is the 100M baseline. It is not GPT-4 in miniature; it is a fluent gibberish machine, and that is exactly what it should be.

Compare: Llama-3-8B has val perplexity ~4–5 on similar data; GPT-4-class models are around 2–3. The gap between 100M and 8B is roughly one order of magnitude of perplexity, which is why you can't skip the scale.

9.4 · What to do while your run is running — the productive-babysitting protocol

• First 500 steps: watch wandb like a hawk. This is where 90% of bugs show up (loss stuck, gradient NaN, throughput off by 3×).
• Steps 500–3000: loss should be sliding smoothly. Now is the time to sanity-check the val metric and generation quality (sample once every 500 steps).
• Steps 3000–6000: loss curve is boring. Use this time to write your S054 eval harness — lm-eval-harness config, HellaSwag/ARC-easy/PIQA loading, plotting code. When training finishes, everything is ready.
• Last 500 steps: don't touch anything. Watch for the LR decay to finalize the loss smoothly. Confirm your checkpoint-save cadence just fired.

Further reading (9.x)

• Karpathy llm.c writeup — github.com/karpathy/llm.c/discussions/677
• modded-nanogpt leaderboard — github.com/KellerJordan/modded-nanogpt
• Muon optimizer — kellerjordan.github.io/posts/muon
• FineWeb-Edu dataset card — huggingface.co/datasets/HuggingFaceFW/fineweb-edu
• Muennighoff "Scaling Data-Constrained LMs" (over/under-training) — arxiv:2305.16264

10 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §3 (pre-flight checklist) and §9 (2025 numbers) in 24 hours. Revisit the full session on day 7.

Next session (S054): loss curves + evals — how to interpret what you just produced, and whether it's actually usable.

Sticky note (keep on your desk): "Init loss = log(vocab). Overfit-one-batch first. Val loss 3.3–3.6 = healthy at 100M/2B."

Legacy recall drills (kept for spaced-review continuity)

Stretch

Modify the config to a 300M model at Chinchilla-optimal (~6B tokens). Recompute: how many steps? How many A100-hours at MFU 40%? What's the total cost? Now do it for 1B model at 20B tokens. Compare the ratios. Is the compute-vs-cost linear? Why or why not?

In your own words

"The pre-flight checklist matters because ______."

Spaced-review pointer

• S050 (architecture) — the model config.
• S052 (renting GPU) — you're now on that rented pod.
• S018 (mixed precision) — bf16 is doing the heavy lifting.

Next-session teaser

Session 054 is loss-curve interpretation, plus building the eval harness (perplexity on held-out, HellaSwag, ARC-easy, MMLU-tiny) that lets you compare your model to Llama or Phi-1. This is the "did I make anything useful?" session.

Bring back tomorrow

• Config: the full 100M TrainConfig (know every knob).
• Numbers: ~7630 steps, ~2.7 hours, ~$3, final loss ~3.8.
• Checklist: the 6-item pre-flight.

Previous: ← DL S052 · Next: DL S054 →