Training infrastructure (mixed precision, optimisers, monitoring)

Why this matters

You can have the best architecture, the best data, and the right scaling law — and still fail to train if your infrastructure is wrong. Numerical precision, optimiser choice, learning-rate schedule, gradient clipping, and monitoring are all small decisions whose product determines whether your model converges or diverges.

This lesson is the engineering counterpart to the architectural ones. It covers the realities of getting a training job to actually finish.

Learning objectives

1. Configure bf16 / fp8 mixed precision correctly.
2. Choose and tune AdamW (and know about Lion, Sophia, Muon).
3. Implement warmup + cosine LR schedule.
4. Apply gradient clipping and skip-on-NaN.
5. Monitor a run with WandB or TensorBoard and interpret the dashboards.

1. Numerical precision

bf16 mixed precision

with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    out = model(x)
    loss = loss_fn(out, y)
loss.backward()                  # bf16 gradients accumulate fine
optim.step()                     # AdamW master weights stay in fp32

No loss scaler needed (unlike fp16). bf16 is the universally correct default in 2026.

fp8 training (advanced)

NVIDIA's Transformer Engine (TE) automates fp8 matmuls with per-tensor scaling. DeepSpeed and Megatron support it. ~2× speedup on H100; small quality cost. Used in Llama 3 405B, DeepSeek-V3.

from transformer_engine.pytorch import fp8_autocast
with fp8_autocast(enabled=True):
    out = model(x)

For your home lab: stick to bf16. fp8 is an "I have a Hopper cluster" tool.

2. Optimiser choice

For pretraining ≥1B params: AdamW with (β1, β2) = (0.9, 0.95), eps=1e-8, weight_decay=0.1. This is the de facto setting from GPT-3 onward.

Decoupled weight decay (the W in AdamW)

Apply weight decay outside the gradient update:

θ ← θ - lr * (m / (sqrt(v) + eps) + wd * θ)

Empirically much better than coupling weight decay into the gradient. Loshchilov & Hutter (2019).

3. Learning-rate schedule

The standard recipe is:

lr(step) =
  step / warmup_steps * lr_max               if step < warmup_steps
  lr_min + 0.5*(lr_max - lr_min) * (1 + cos(pi * (step - warmup) / decay_steps))   else

Common values:

• lr_max = 3e-4 for small models, 1.5e-4 to 3e-4 for 7-70B (depends on batch size).
• warmup_steps = 0.5 - 2% of total steps.
• lr_min = 0.1 * lr_max.

Variants:

• Cosine (default).
• Linear decay (used by some Llama trainings).
• WSD (Warmup-Stable-Decay) — keeps lr_max constant for the bulk of training, then linearly decays at the end. Popular in 2024+ because you can extend training without recomputing the schedule.

Learning rate vs batch size

Linear scaling rule: when batch increases by k, lr increases by k (roughly). Breaks at very large batches (gradient noise dominates). Use sqrt(k) if k > 8.

4. Gradient clipping

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Cap the global gradient norm at 1.0 (sometimes 0.5). Without this, one bad batch can produce a NaN that destroys the run.

Skip-on-NaN

A more aggressive version: detect NaN/Inf in loss or gradients, and skip the optimiser step.

if torch.isfinite(loss):
    scaler.step(optim)
else:
    print(f"step {step} non-finite; skipping")

Frontier labs go further: skip when gradient norm is > 5× the running median (anomaly detection).

5. Initialisation, again

For deep networks:

• Linear layers: N(0, 0.02) (GPT-2 convention) or He init.
• Output projection of attention/MLP scaled by 1 / sqrt(2L) (Megatron, GPT-J trick) → keeps activations bounded with depth.
• Final LM head: optionally zero-init for very stable starts (some labs do).
• Embeddings: N(0, 0.02) then sometimes scaled by sqrt(d_model) if embeddings are not tied.

def init(m):
    if isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, std=0.02)
    elif isinstance(m, nn.Embedding):
        nn.init.normal_(m.weight, std=0.02)
model.apply(init)

# Megatron-style residual scaling
for name, p in model.named_parameters():
    if "out_proj" in name or "w2" in name:   # output projections
        p.data.mul_(1.0 / math.sqrt(2 * cfg.n_layers))

6. Monitoring — what to log

A modern dashboard tracks ~15-30 series. The non-negotiables:

For a 1k-GPU run you also chart per-rank metrics so you can find a slow GPU.

WandB minimal example

import wandb
wandb.init(project="nano-gpt", config=cfg.__dict__)
for step, (x, y) in enumerate(loader):
    ...
    if step % 10 == 0:
        wandb.log({"train/loss": loss.item(), "lr": lr, "grad_norm": gn,
                   "tok_per_s": tok_per_s}, step=step)

7. Reproducibility (the unsung skill)

For a paper, others will try to reproduce your run. Bake these in from day 1:

• Pin torch, cuda, flash-attn versions.
• Save a full config (architecture, data hashes, optimiser, seed) per run.
• Save tokeniser version.
• Store the exact git SHA of the training code.
• Deterministic data shuffling: np.random.RandomState(seed + rank).permutation.
• Avoid non-deterministic kernels for paper baselines (torch.use_deterministic_algorithms(True)).

For a production run absolute determinism is usually not worth the cost; reproducibility-by-config is.

8. Putting it all together — a 30M-param recipe you can run

cfg = Config(d_model=512, n_layers=8, n_heads_q=8, n_heads_kv=4,
             d_ff=1536, block_size=1024, vocab_size=50257)
m = GPT(cfg).cuda()

steps = 20_000
warmup = 200
lr_max = 3e-4
opt = torch.optim.AdamW(m.parameters(), lr=lr_max, betas=(0.9, 0.95),
                        weight_decay=0.1)

def get_lr(step):
    if step < warmup:
        return lr_max * step / warmup
    progress = (step - warmup) / (steps - warmup)
    return 0.1 * lr_max + 0.5 * (lr_max - 0.1 * lr_max) * (1 + math.cos(math.pi * progress))

for step in range(steps):
    for g in opt.param_groups:
        g["lr"] = get_lr(step)

    x, y = get_batch(train, cfg.block_size, 32, "cuda")
    with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
        _, loss = m(x, y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    gn = torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
    opt.step()

    if step % 50 == 0:
        print(f"step {step:5d} | loss {loss.item():.4f} | lr {get_lr(step):.2e} | gnorm {gn:.2f}")

Train for 20k steps on a 1B-token corpus and you have a respectable 30M-param language model.

Hands-on lab (3 hours)

infra_lab.ipynb:

1. Wire up your training loop with WandB. Log loss, lr, grad_norm, tokens/s.
2. Implement warmup + cosine LR by hand. Plot the schedule.
3. Add skip-on-NaN. Test by injecting loss = torch.tensor(float('nan')) once.
4. Switch from AdamW to Lion (pip install lion-pytorch). Compare loss curves at the same lr.
5. Switch from fp32 to bf16 autocast. Verify identical training behaviour and ~30% more throughput.
6. Bonus: enable torch.compile(model) and add torch._dynamo.config.cache_size_limit=128. Measure speedup.

Common pitfalls

1. Forgetting to schedule the LR — the model trains, just much worse.
2. Using grad_clip = 1.0 without monitoring → silent gradient saturation.
3. Setting eps=1e-5 in AdamW (Adam default) — for big models 1e-8 is correct.
4. Not logging step_time — you only realise the I/O is slow when training takes 5× longer than estimated.
5. Initialising the bias of LayerNorm/RMSNorm to 0 when convention is 1 (RMSNorm has only weight, init to 1).

Self-check

1. Why bf16 over fp16 for training?
2. What does decoupled weight decay mean?
3. Standard β1, β2 for LLM AdamW?
4. Why warmup the learning rate?
5. What is skip-on-NaN and why does it help?

References

• Loshchilov & Hutter (2019), "Decoupled Weight Decay Regularization."
• Smith et al. (2017), "Don't Decay the Learning Rate, Increase the Batch Size."
• Liu et al. (2023), "Sophia: A Scalable Stochastic Second-order Optimizer."
• Chen et al. (2023), "Symbolic Discovery of Optimization Algorithms" (Lion).
• Karpathy, nanoGPT training tips.
• NVIDIA, Transformer Engine docs.
• Weights & Biases training tips.