Training mechanics

Your PyTorch loop runs. The first batches look ordinary. Then the loss becomes NaN, or training loss keeps falling while validation loss turns upward. The final checkpoint exists, but nobody can explain why that particular model was selected.

The five training calls are correct. The run is not yet a controlled experiment.

Why this matters

Training mechanics are the operating decisions around the loop: how examples form updates, how large each update may be, when its size changes, which evidence selects a checkpoint, and what you record when something breaks.

Plain-English mental model

Step 1 — Count updates, not only epochs

A batchBatchA small group of training examples processed together in one forward and backward pass. Batches make computation practical and provide an estimate of the average training signal; their size affects memory use, speed, and the noisiness of updates. is the group of examples used for one forward and backward calculation. In the ordinary loop, one batch produces one optimiser update.

An epochEpochOne pass through the training examples as defined by a data loader. An epoch is an accounting unit, not proof that a model has learned enough or will work on new data. is one pass through the training loader. It measures data passes, not a fixed amount of learning. Changing batch size changes the number of updates inside the same epoch.

For N training examples and batch size B, when the last smaller batch is kept:

batches per epoch = ceiling(N / B)
updates = batches per epoch × epochs

With 1,000 examples and batch size 100, one epoch contains ten batches and ordinarily ten updates. Six epochs contain 60 updates.

from math import ceil

examples = 1_000
batch_size = 100
epochs = 6

batches_per_epoch = ceil(examples / batch_size)
updates = batches_per_epoch * epochs

print('batches per epoch:', batches_per_epoch)
print('updates:', updates)

Batch size affects memory use, the amount of batch-to-batch gradient variation, and the number of updates. It does not make a split valid, repair leakage, or guarantee better generalisation.

Step 2 — Choose an update rule and a step size

An optimiserOptimiserAn algorithm that changes trainable parameters using their gradients and settings such as a learning rate. In PyTorch an optimiser owns the update rule; calling backward calculates gradients, while calling step applies an update. turns gradients into parameter updates. A learning rateLearning rateA positive setting that controls the size of an optimiser’s parameter update. Too large can overshoot or diverge; too small can make training impractically slow. It must be evaluated with the full training and validation process. controls the update scale.

Plain stochastic gradient descent, usually shortened to SGD, uses the current batch gradient. Momentum keeps a moving average of recent gradients, which can reduce repeated zigzagging. Adam keeps moving averages of gradients and squared gradients, then adapts the update scale by coordinate. AdamW applies weight decay separately from that adaptive gradient update.

These are different starting hypotheses, not a universal ranking. Optimiser state also matters when you resume: momentum and Adam-family methods remember earlier gradients.

For a single weight, the simplest update is:

new weight = old weight − learning rate × gradient

If weight = 1.0 and gradient = 0.5, a rate of 0.1 produces 0.95. A rate of 1.0 produces 0.5. The larger move is not automatically better. It may move quickly, overshoot, oscillate, or make later values non-finite.

Step 3 — Warm up, then decay for a stated reason

A learning-rate scheduleLearning-rate scheduleA rule that changes an optimiser’s learning rate during training, such as reducing it after validation progress plateaus. It is a hyperparameter strategy that needs protected validation evidence. changes the rate according to a declared rule.

Warmup begins below the chosen rate and increases over early updates. This can make early updates less abrupt while activations and optimiser state are still settling.

Decay reduces the rate later. Large early steps can make progress; smaller later steps can refine rather than repeatedly cross a narrow region.

Warmup and decay are not automatic repairs. Their duration and shape are more hyperparameters. Log the rate used for every update, and compare the scheduled run with a clear baseline.

This pure-Python schedule warms up for three updates and then decays linearly to one tenth of the starting rate:

def scheduled_rate(step: int, total_steps: int, peak: float, warmup: int) -> float:
    if step <= warmup:
        return peak * step / warmup

    decay_progress = (step - warmup) / (total_steps - warmup)
    return peak * (1.0 - 0.9 * decay_progress)


for step in range(1, 9):
    rate = scheduled_rate(step, total_steps=8, peak=0.08, warmup=3)
    print(step, f'{rate:.3f}')

In PyTorch, call most schedulers after optimizer.step(). A metric-driven scheduler such as ReduceLROnPlateau receives a validation measure. Never give it final test loss: that would use the test set to configure training.

Step 4 — Read curves as evidence, not decoration

A learning curve records a measure across training progress. Plot training and validation measures against optimiser updates or epochs, and log the learning rate beside them.

• Training and validation losses both fall: the run is improving on these measures so far.
• Training loss falls while validation loss rises: the run has passed its best validation point, the split is unrepresentative, or evaluation differs from training in another important way.
• Both curves stay flat: the model may not receive useful gradients, the rate may be too small, targets or masks may be wrong, or the task may not be learnable from these inputs.
• Loss jumps or becomes NaN: inspect data, arithmetic, gradients, and the most recent configuration change before adding a safeguard.

The lab below uses a deterministic one-parameter illustration. Its training and validation sets deliberately prefer slightly different values, so the best validation point can occur before the final update. The exact numbers do not describe a real neural network.

Try three controlled comparisons:

1. Keep the defaults and note the best validation epoch.
2. Set a high rate and choose Constant rate. Observe whether the run becomes unstable.
3. Return to the default rate. Change only batch size or accumulation count and inspect the update count before interpreting the curves.

The chart uses a solid training line and a dashed validation line, so its meaning does not depend on colour.

Step 5 — Accumulate gradients when one batch does not fit

A micro-batch is the small batch that fits through the model once. Gradient accumulationGradient accumulationA technique that adds gradients from several micro-batches before one optimiser step. It creates a larger effective batch without holding all examples in memory at once; mean losses normally need scaling by the accumulation count. combines gradients from several micro-batches before one optimiser update.

If a micro-batch contains 16 examples and you accumulate four micro-batches, one update uses an effective batch size of 64 examples on one process:

effective batch = micro-batch size × accumulation steps

When the loss function returns a mean for each micro-batch, divide each loss by the number of accumulation steps before backward(). Then the stored gradients approximate the mean over the effective batch instead of a sum four times larger.

accumulation_steps = 4
optimizer.zero_grad()

for micro_step, (features, targets) in enumerate(train_loader, start=1):
    loss = loss_fn(model(features), targets) / accumulation_steps
    loss.backward()

    if micro_step % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Handle an incomplete final accumulation group deliberately. In distributed training, the effective global batch also depends on the number of worker processes. Do not compare two runs only by epoch when their update counts or effective batches differ.

Try it — break the scaling on purpose

This browser lab uses two simple gradient values. Run it, then set scale_each_micro_batch = False. The update gradient doubles from the intended average 3.0 to the sum 6.0.

Accumulation reduces activation-memory pressure because each micro-batch is smaller. It does not reproduce every behaviour of one physically larger batch: batch-dependent layers, randomness, distributed communication, and numerical order can differ.

Step 6 — Clip gradients, but keep the evidence

A gradient clipGradient clippingA technique that limits the size of gradients before an optimiser update, commonly by scaling them down when their overall norm exceeds a threshold. It can prevent a rare extreme update from destabilising training, but does not explain why gradients became extreme. limits gradient size before the optimiser update. Norm clipping measures all parameter gradients together and rescales them when their combined norm exceeds a threshold.

loss.backward()
gradient_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0,
    error_if_nonfinite=True,
)
optimizer.step()

print('norm before clipping:', gradient_norm.item())

clip_grad_norm_ changes gradients in place and returns their total norm before clipping. Log that value. If almost every update clips, the threshold may be hiding a scale, loss, rate, or data problem.

With automatic mixed precision, unscale the gradients before inspecting or clipping them. Otherwise the clipping threshold is applied to temporarily scaled values rather than the gradients the optimiser should use.

Clipping can prevent one unusually large update from dominating a run. It cannot explain why the gradient became large, make a NaN finite, or prove the model is learning the intended task.

Step 7 — Debug the first divergence, not the final wreckage

Non-finite means NaN, positive infinity, or negative infinity. Find the first value or update that becomes non-finite. Later failures are often consequences.

Use this order:

1. Preserve the failing configuration, code revision, data version, and first failing batch.
2. Return to the last known-good run and reproduce its result.
3. Check input features, targets, loss, parameters, and gradients with torch.isfinite.
4. Verify shapes, target ranges, masks, loss reduction, training/evaluation modes, and gradient clearing.
5. Try to overfit one tiny batch. If loss cannot fall there, debug the training mechanics before adding more data.
6. Log the learning rate and gradient norm before clipping.
7. Change one factor and rerun the smallest reproduction.

torch.autograd.detect_anomaly() can report the forward operation connected to a failing backward operation and can raise when backward produces NaN. It slows execution, so use it on the small debugging reproduction rather than leaving it on for ordinary training.

Different symptoms suggest different first checks:

• Flat from the start: confirm parameters require gradients, the optimiser owns them, step() runs, and the loss depends on the model output.
• Wild but finite: inspect rate, batch variation, loss scale, gradient norm, and data outliers.
• Sudden non-finite value: inspect the first failing operation, division or logarithm inputs, mixed-precision scaling, and the preceding update.
• Training improves, validation worsens: audit split and evaluation code first; then compare an earlier checkpoint or a controlled generalisation change.

Step 8 — Stop and save with a declared validation rule

Early stoppingEarly stoppingA predefined rule that stops training when a validation measure has not improved for a chosen number of checks. It uses validation data to select duration and must not repeatedly expose the final test set. stops training after a validation measure has not improved for a chosen patience: a number of validation checks you are willing to wait.

Choose before final testing:

• the validation measure;
• whether lower or higher is better;
• the minimum change that counts as improvement;
• the patience;
• how often validation runs.

Save the best validation checkpointCheckpointA saved training artifact containing model state and supporting context such as optimiser state, epoch, configuration, and metrics. A checkpoint makes a run recoverable or investigable; it is not deployment evidence by itself., not merely the final epoch. A recoverable checkpoint normally contains:

• model state;
• optimiser state;
• scheduler state;
• epoch and optimiser-update number;
• best validation measure and selection rule;
• model and preprocessing configuration;
• data and split versions;
• code revision and relevant package versions;
• random seeds where reproducibility requires them.

if validation_loss < best_validation_loss - minimum_change:
    best_validation_loss = validation_loss
    checks_without_improvement = 0
    torch.save(
        {
            'model_state': model.state_dict(),
            'optimizer_state': optimizer.state_dict(),
            'scheduler_state': scheduler.state_dict(),
            'epoch': epoch,
            'update': update,
            'best_validation_loss': best_validation_loss,
            'selection_rule': {
                'minimum_change': minimum_change,
                'patience': patience,
            },
        },
        'best-validation.pt',
    )
else:
    checks_without_improvement += 1

Before the one final test evaluation, restore the selected checkpoint. Repeatedly testing different runs on the same test set turns that test set into configuration data.

Real-world AI engineering example

Common mistakes

Knowledge check

Practice activity

Exercises

Key takeaways

Flashcards

One-page sketchnote summary

batch → backward
accumulate? → clip? → step
validate → log → save best

Resources