The training loop

Problem

A model starts out knowing nothing: its weights are random (or, in our example, all zero). We want to turn it into one that predicts well, using a large pile of examples. No formula does that in one go. We need a procedure that makes the model a little better, many thousands of times in a row, without breaking along the way.

That procedure has to answer several practical questions. How do we score a prediction with one number? How do we use that score to change the weights? How big should each change be, and should that size stay fixed? What if one unlucky batch produces an enormous gradient? What if the numbers get too small for the chip's low-precision format? The training loop is the answer to all of them, and nearly every deep learning project, from a classroom classifier to a frontier language model, runs the same loop with different settings.

Intuition

Think of learning to throw darts. Throw a few, see how far off you are (the loss), work out which way to adjust your arm (the gradient), and adjust a bit (the optimizer step). You adjust a lot at first when you are far off and less as you improve (the learning-rate schedule). If one throw goes wildly wrong, you do not rebuild your whole stance because of it (gradient clipping). Each of those habits is one line in the loop.

The order matters in two places: clip after backward() and before step(), and call the scheduler after the optimizer.

Mechanics

The loss: cross-entropy over logits

A language or classification model ends in a vector of logits, one raw score per class. Cross-entropy turns the logits z into probabilities with softmax and scores the correct class y: ℓ = −log softmax(z)_y = −z_y + log Σ_c exp(z_c). Its gradient with respect to the logits is the beautifully simple p − onehot(y): the predicted probabilities minus the answer (derived on the loss functions page, with the same numbers for the cat, dog, bird example). PyTorch's CrossEntropyLoss takes logits of shape [N, C] (or [N, C, d₁, …]) and integer targets of shape [N], and calls F.cross_entropy (lines 1398-1407).

Two arguments matter for language models:

• ignore_index (default -100): any position whose target equals it contributes nothing to the loss and nothing to the gradient, and reduction="mean" divides by the number of positions that were not ignored. This is how padding and masked prompt tokens are skipped. If every target is ignored the mean is 0/0, which is NaN.
• label_smoothing (default 0.0): the target becomes a mixture of the true class and a uniform distribution (the idea comes from Rethinking the Inception Architecture, cited in the docstring at lines 1293-1296). With smoothing ε the loss is (1 − ε)·(−log p_y) + ε·(mean over classes of −log p_c). It stops the model chasing probability 1.0 for the label.

Backward and the gradient buffer

loss.backward() runs backpropagation and adds each parameter's gradient into its .grad attribute. Adding, not overwriting, is why the loop calls optimizer.zero_grad() every iteration (by default it sets the gradients to None, set_to_none=True). It is also what makes gradient accumulation possible: call backward() on several micro-batches before one step().

Gradient clipping

Occasionally a batch produces a gradient hundreds of times larger than usual, and one step of that size can wreck the weights. clip_grad_norm_(parameters, max_norm) looks at all gradients as one long vector, computes its length (_get_total_norm), and multiplies every gradient by min(max_norm / (total_norm + 10⁻⁶), 1) (lines 164-168). The direction is preserved; only the length is capped. Gradients that are already short are left alone, because the factor is clamped to 1. The function modifies .grad in place and returns the pre-clip total norm, which is worth logging: a spike in it is an early warning.

The optimizer step

optimizer.step() reads each .grad and updates the weights. For language models the usual choice is AdamW with decoupled weight decay, default weight_decay=1e-2; its update rule, state and the decay-versus-L2 difference are on the optimizers page. The numbers on this page use that page's optimizer_step helper.

The learning-rate schedule

A constant learning rate is rarely best. Early on, Adam's running averages are noisy and a large rate can throw the model somewhere bad, so the rate warms up linearly from near zero. Later, a smaller rate lets the weights settle, so the rate decays, most often along a half cosine down to a floor:

η_t = η_min + ½(η_max − η_min)(1 + cos(π·t′/T))

where t′ counts steps since warmup ended and T is the number of decay steps. In PyTorch this is built from two schedulers chained with SequentialLR: LinearLR for the warmup and CosineAnnealingLR for the decay. At the milestone, SequentialLR.step restarts the next scheduler at its own epoch 0 (lines 1186-1197), and the cosine's closed form is the formula above (lines 1456-1461). The constructor of every scheduler overwrites the optimizer's lr, so build the scheduler once, after the optimizer, and call scheduler.step() once per optimizer step.

The peak and floor are typical for a small language model, not prescriptions. Many recipes set the floor at 10% of the peak.

Mixed precision

Modern GPUs run matrix multiplications much faster in 16-bit numbers. Mixed precision keeps the weights in 32-bit but runs the heavy operations in 16-bit. torch.autocast (device_type, dtype=None, enabled=True, constructor at lines 222-228) is a context manager: inside it, eligible operations run in the lower precision. The docstring warns against also calling .half() or .bfloat16() on the model or inputs (line 62).

The float16 format has a catch: its smallest positive numbers are tiny but not tiny enough, so small gradients silently become exactly 0 (underflow). The fix is loss scaling: multiply the loss by a large constant before backward(), so every gradient is that many times bigger, and divide it out again before the update. PyTorch's GradScaler automates it, starting at init_scale=2**16 (lines 121-129):

• scaler.scale(loss) multiplies the loss by the current scale (line 193).
• scaler.unscale_(optimizer) divides the gradients back down, and is the call to make before clipping so the clip threshold applies to real-size gradients (lines 302-374).
• scaler.step(optimizer) unscales if needed, and skips the update if any gradient is infinite or NaN (lines 375-483).
• scaler.update() multiplies the scale by backoff_factor=0.5 after a skipped step, and by growth_factor=2.0 after growth_interval=2000 consecutive good steps (lines 484-560).

Train mode, eval mode and teacher forcing

model.train() and model.eval() flip a flag on every submodule (train, eval). It changes how Dropout and BatchNorm behave (see normalization) and does nothing else; it does not turn gradient tracking off. For that, wrap evaluation in torch.no_grad().

Sequence models are trained with teacher forcing: at every position the model is given the true previous tokens, not its own earlier guesses, and asked for the next one. That lets all positions be scored in one parallel pass. At generation time there is no truth to feed in, so the model consumes its own output one token at a time, which is why training and inference differ (see the transformer page).

Worked example

Implementation

import math
import torch
from torch import nn
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR

model = MyModel().cuda()
loss_fn = nn.CrossEntropyLoss(ignore_index=-100, label_smoothing=0.0)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)

warmup, total = 100, 1000
sched = SequentialLR(
    opt,
    [LinearLR(opt, start_factor=1 / warmup, total_iters=warmup - 1),
     CosineAnnealingLR(opt, T_max=total - warmup, eta_min=3e-5)],
    milestones=[warmup],
)
scaler = torch.amp.GradScaler("cuda")          # float16 only

model.train()
for x, y in loader:                            # x: [B, T] token ids, y: [B, T]
    opt.zero_grad(set_to_none=True)
    with torch.autocast("cuda", dtype=torch.float16):
        logits = model(x)                      # [B, T, V]
        loss = loss_fn(logits.flatten(0, 1), y.flatten())
    scaler.scale(loss).backward()
    scaler.unscale_(opt)                       # gradients are real-size again
    norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(opt)                           # skipped if a gradient is inf/NaN
    scaler.update()
    sched.step()

model.eval()
with torch.no_grad():
    val_loss = loss_fn(model(xv).flatten(0, 1), yv.flatten())

With bfloat16 there is no GradScaler: its exponent range matches float32's, so gradients do not underflow the way float16 ones do. Pass dtype=torch.bfloat16 to autocast and drop the scaler lines. The same loop in plain Python, as used for the numbers on this page:

import math

def clip(grads, max_norm):
    total = math.sqrt(sum(g * g for g in grads))
    coef = min(max_norm / (total + 1e-6), 1.0)
    return [g * coef for g in grads], total

def lr_at(step, peak, warmup, total, floor):
    if step < warmup:
        return peak * (step + 1) / warmup
    t = (step - warmup) / (total - warmup)
    return floor + (peak - floor) * (1 + math.cos(math.pi * t)) / 2

# per step: loss -> gradients -> clip -> AdamW update with lr_at(step, ...)

Tradeoffs

• Clipping threshold. A threshold of 1.0 is a common default. Set too low it turns every step into a fixed-length step and slows learning; set too high it never fires. Log the pre-clip norm and choose a value just above its typical size.
• Warmup length. Too short and early steps can diverge, especially with Adam, whose second-moment estimate is poor in the first iterations (see Liu et al., On the Variance of the Adaptive Learning Rate and Beyond). A few hundred to a few thousand steps is typical for language models; large-batch training also warms up (Goyal et al.).
• Decay shape. Cosine to a floor is a popular default, but linear decay and constant-then-decay schedules work too. Cosine needs the total number of steps in advance, which makes extending a run awkward.
• Label smoothing. It can improve calibration and generalization on classification, but it flattens the model's confidence, which hurts if you later rely on exact probabilities.
• float16, bfloat16 or float32. float16 is precise but narrow and needs the scaler; bfloat16 is coarse but wide and needs none; float32 is the safe, slow baseline.
• Pitfalls. Forgetting zero_grad() accumulates gradients across steps. Calling scheduler.step() before optimizer.step() skips the first learning rate and warns. Clipping before unscale_ clips scaled gradients against an unscaled threshold. Forgetting model.eval() leaves dropout on during validation. Averaging losses over micro-batches with unequal numbers of kept tokens biases the gradient.

Check yourself

Connections

References

• Kingma and Ba. Adam: A Method for Stochastic Optimization. ICLR, 2015 (arXiv 1412.6980).
• Loshchilov and Hutter. Decoupled Weight Decay Regularization. ICLR, 2019 (arXiv 1711.05101).
• Loshchilov and Hutter. SGDR: Stochastic Gradient Descent with Warm Restarts. ICLR, 2017 (arXiv 1608.03983; the cosine schedule).
• Goyal et al. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. 2017 (arXiv 1706.02677; linear warmup).
• Liu et al. On the Variance of the Adaptive Learning Rate and Beyond. ICLR, 2020 (arXiv 1908.03265).
• Pascanu, Mikolov and Bengio. On the Difficulty of Training Recurrent Neural Networks. ICML, 2013 (arXiv 1211.5063; gradient norm clipping).
• Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR, 2016 (arXiv 1512.00567; label smoothing).
• Micikevicius et al. Mixed Precision Training. ICLR, 2018 (arXiv 1710.03740; loss scaling).
• PyTorch v2.15.0: loss.py, adamw.py, lr_scheduler.py, autocast_mode.py, grad_scaler.py, clip_grad.py.