Classify the symptom before changing precision
A loss spike, a nonfinite tensor, and a skipped optimizer step are related observations, not interchangeable diagnoses. A spike can come from a hard batch, normalization, a schedule transition, corrupted input, an optimizer-state problem, or reduced-precision range. A gradient scaler may skip an update when gradients contain infinities or NaNs, but that safety action does not prove autocast created the invalid value.
Capture the first bad update with its batch identifiers, learning rate, loss components, gradient norm, scaler value, skipped-step status, and recent checkpoint. Compare against the normal distribution for the same run. Do not immediately lower the learning rate, clip harder, and disable mixed precision together. Multiple simultaneous changes can make the symptom disappear while leaving the cause unknown.
Understand what loss scaling can and cannot do
Float16 gradients can underflow when small values are not representable. Gradient scaling multiplies the loss before backward, then unscales gradients before the optimizer update. PyTorch GradScaler checks gradients and can skip optimizer.step when infinities or NaNs are present, then updates the scale. This protects parameters from one invalid update; it does not repair a forward activation that already became nonfinite or a mathematically invalid loss.
With gradient accumulation, the scale must remain constant across the microbatches contributing to one effective batch, and unscaling should occur only after accumulation is complete. PyTorch’s AMP examples explicitly place unscale, clipping, step, and scale update at the effective-batch boundary. Updating the scaler per microbatch mixes differently scaled gradients and makes recovery of the intended sum impossible.
Find the first nonfinite tensor
Instrument from inputs toward the optimizer. Check input tensors and labels, model outputs, each loss component, selected intermediate activations, unscaled gradients, gradient norm, parameters after the step, and optimizer state. PyTorch torch.isfinite returns a boolean mask for NaN and positive or negative infinity, making it suitable for assertions at chosen boundaries. Start with coarse module boundaries, then narrow the failing region to avoid overwhelming logs and perturbing timing.
Run the captured batch in evaluation and training modes as appropriate, first with the original mixed-precision path, then with autocast disabled for the suspect region, and finally in full precision if memory permits. If the input or full-precision loss is already nonfinite, mixed precision is not the root cause. If only one reduced-precision region fails, inspect its operation, input range, and accumulation behavior before widening more of the model.
Use a controlled diagnostic matrix
Freeze the model checkpoint, optimizer state, batch, seeds, and software environment. Change one axis per replay: autocast on or off, scaler on or off, suspect region forced to float32, optimizer step disabled, and data batch replaced with a known-good batch. PyTorch’s AMP recipe recommends disabling autocast or GradScaler individually and forcing a suspect forward region to float32 when investigating infinities or NaNs.
Interpret the matrix cautiously. Failure before backward implicates data, model forward, or loss computation. Finite loss followed by nonfinite unscaled gradients points toward backward amplification or an operation’s gradient. Finite gradients followed by bad parameters points toward clipping, optimizer math, or optimizer state. A failure that follows one batch in full precision suggests a data or objective issue. A failure that disappears in float32 only localizes sensitivity; it does not by itself identify why values reached that range.
| Observation | Leading possibilities | Next controlled check |
|---|---|---|
| Nonfinite input or label | Data decode or preprocessing | Replay raw sample pipeline |
| Finite input, bad loss in full precision | Objective or data edge case | Inspect loss components |
| Finite loss, bad unscaled gradient | Backward overflow or amplification | Module-level gradient checks |
| Finite gradient, bad parameter after step | Optimizer state or update | Disable step and inspect moments |
A worked ledger prevents false progress
Consider an illustrative run where attempted update 4,200 starts with scale 65,536 and the scaler uses a backoff factor of 0.5. The loss is finite, but the scaler's gradient checks find infinity after unscaling, so scaler.step skips the underlying optimizer step and scaler.update reduces the scale to 32,768. Parameters and a learning-rate schedule defined over successful optimizer updates should remain at committed update 4,199. The next batch succeeds at scale 32,768, and only then does committed update count become 4,200.
If logging increments the step and scheduler before knowing whether the optimizer updated, the run can silently advance schedules without changing weights. Track attempted batches separately from committed optimizer updates and skipped updates. The values are illustrative. Occasional early skips may be part of scaler calibration, but repeated skips, a steadily falling scale, or clustered failures on the same inputs require diagnosis rather than being averaged away.
Recover from a known-good state
Once parameters or optimizer moments contain NaNs, continuing after reducing the scale is not reliable. Restore the newest checkpoint whose model, optimizer, scheduler, and scaler states pass finite checks. Replay the failing batch under diagnostic instrumentation. If a specific reduced-precision operation is responsible, keep only that region in float32 or choose a supported wider-range format, then rerun a representative window.
Quarantine bad data only when the sample itself violates a declared data contract and the decision is recorded. Skipping every batch that triggers instability can bias training and conceal a model defect. Likewise, resetting optimizer state may appear to recover loss while materially changing the run. Label that action as a warm restart and require a renewed quality comparison.
Set a promotion rule and acknowledge limits
A measurable recovery rule can require zero nonfinite inputs, losses, unscaled gradients, parameters, and optimizer tensors over a replay window; no unexplained skipped updates after the scaler’s declared warm-up; and loss plus gradient-norm distributions within a chosen tolerance of a full-precision diagnostic. For example, a team might require 1,000 clean committed updates and less than a 2 percent median throughput loss from the targeted float32 region. Those thresholds are illustrative, not universal.
Then run the normal validation suite and a longer representative segment. Finite arithmetic does not guarantee equivalent convergence, and a short full-precision replay is not a gold standard for final quality. Mixed precision also interacts with model architecture, loss normalization, optimizer hyperparameters, batch composition, and hardware kernels. Capacity measurements should use the precision mode that passed this acceptance test, including any float32 regions and their memory impact, rather than an assumed nominal datatype.