Training loop

What the loop does, in one loop

With the data tokenized and the model built, training is a standard supervised loop: sample a batch, run the forward pass, compute a loss, backprop, step the optimizer, repeat. What differs from a first-course PyTorch loop are the engineering details that keep it stable and fast at language-model scale.

The conceptual shape (actual code lives in full_llm_pipeline.ipynb):

for step, batch in enumerate(loader):
    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
        logits, aux_loss = model(batch.input_ids)
        loss = F.cross_entropy(
            logits[:, :-1].reshape(-1, V),
            batch.input_ids[:, 1:].reshape(-1),
        )
        loss = loss + 0.01 * aux_loss     # MoE load-balancing

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)
    scheduler.step()

The rest of this page walks through each piece.

Optimizer — AdamW

AdamW is Adam with decoupled weight decay. Classical Adam rolls L2 regularization into the gradient, which interacts badly with Adam's adaptive per-parameter learning rates (large-gradient parameters effectively get less decay). AdamW applies the decay directly to the weights:

w  ←  w  −  η · m_hat / (sqrt(v_hat) + ε)   # Adam step
w  ←  w  −  η · λ · w                       # decay, separate from gradient

In code:

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    weight_decay=0.01,
)

lr=3e-4 is a reasonable starting LR for Transformer pretraining. Tune for your model size and batch size.

Paper: Loshchilov & Hutter, Decoupled Weight Decay Regularization, ICLR 2019. https://arxiv.org/abs/1711.05101.

Schedule — cosine annealing

Instead of a constant learning rate, we smoothly decay lr along a cosine half-period from lr_max at step 0 to lr_min at the last step. Intuitively: take big optimization steps early when the loss landscape is simple, small steps late when you are polishing.

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=PRETRAIN_EPOCHS,
)

Paper: Loshchilov & Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, ICLR 2017. https://arxiv.org/abs/1608.03983.

In practice, modern pretraining often prepends a short warmup phase (linearly ramp lr from 0 to lr_max over the first few thousand steps). v0 keeps it simple and omits warmup; adding it is a single-line change.

Mixed precision — bf16 AMP

Training in fp32 uses twice the memory and twice the memory bandwidth of bf16. bfloat16 matches fp32's range (same 8-bit exponent) but with a shorter 7-bit mantissa, so it is the standard choice for large Transformer training.

• Forward + backward run under torch.amp.autocast("cuda", dtype=torch.bfloat16), which casts compatible ops to bf16.
• Optimizer state stays in fp32 (the m, v buffers of Adam) for numerical stability.
• GradScaler handles the edge cases when fp16 ops are still in play; with pure bf16 it is effectively a no-op but we keep it in for CUDA portability.
• Disabled on CPU — the dev config runs on CPU in plain fp32.

Papers: Micikevicius et al., Mixed Precision Training, ICLR 2018. https://arxiv.org/abs/1710.03740. Kalamkar et al., A Study of BFLOAT16 for Deep Learning Training, 2019. https://arxiv.org/abs/1905.12322.

Gradient clipping

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

Caps the total gradient L2-norm at 1.0 before the optimizer step. If a bad batch (a rare token, a noisy sample) produces a huge gradient, clipping keeps a single step from blowing up the training trajectory. Standard practice for LM training.

MoE auxiliary loss

The MoE router returns an auxiliary load-balancing loss alongside the prediction logits. We add it to the cross-entropy loss with weight 0.01:

loss = ce_loss + 0.01 * aux_loss

Without it, the router collapses to a single expert. With it, expert utilization stays roughly uniform.

Distributed training at scale

The dev config trains fine on one GPU (or slowly on CPU). The 70B reference config is designed for 8× A100-80GB with either:

• FSDP (PyTorch's Fully Sharded Data Parallel) — shards parameters, gradients, and optimizer state across GPUs. https://pytorch.org/docs/stable/fsdp.html.
• DeepSpeed ZeRO-3 — the other mainstream option for the same sharding strategy. https://github.com/microsoft/DeepSpeed.

Both collapse the per-GPU memory footprint enough that a 70B-param model fits on 8 cards. Neither is exercised in v0.1.0; wiring either one up is future work.

Next stage: Checkpointing — saving state so a long run survives a crash.