The final exception in an interrupted training run may come from logging. In one run, strict JSON serialization rejected a non-finite float while writing an optimizer audit record. A later run explicitly recorded a non-finite loss and was stopped by a guard. Both exits involved NaN, but the final exception alone could not identify the operation that produced it.
After restoring the latest normal weights, FP32 and AMP forward/backward checks on a few real images were finite. The tested state did not trigger the problem again. However, the failing batch and preceding updates had not been replayed, so the numerical fault remained unresolved. The investigation shifted toward the state needed to reconstruct the failed computation.
1. Distinguish the stage where the exception appears
A training iteration includes input processing, forward execution, loss computation, backward execution, an optimizer update and logging. Each stage needs different checks:
| Stage | What to inspect | What it can establish |
|---|---|---|
| Forward and loss | Inputs, intermediate activations and loss components | Whether non-finite values have already appeared |
| Backward | Gradients and whether they are still scaled | Whether backward execution exposes an anomaly |
| Update | GradScaler state and actual update events | Whether a parameter update ran |
| After update | Parameters, buffers and optimizer state | Whether the updated state is finite |
| Logging | Serialized fields and the original exception | Whether the observation system preserved the failure |
Finite parameters and buffers in the last normal checkpoint describe the saved state only. They do not cover later unsaved changes or establish that optimizer state, the next batch and intermediate activations are normal.
A NaN loss also does not directly implicate AMP or a particular attention implementation. It may originate earlier in the forward pass or within a loss branch. Narrow the search along the computation path.
Finite inputs and parameters do not guarantee finite intermediates. PyTorch's numerical-accuracy documentation gives a norm example: FP32 inputs remain within range, but an intermediate sum of squares can overflow even when the final mathematical result is representable in FP32. Check intermediate ranges and execution precision as well as the checkpoint. Scanning saved parameters for NaN misses this kind of runtime overflow.
The first observed anomaly is a starting point for investigation. A non-finite reduction may receive already-invalid inputs, or may overflow locally. Distinguishing those cases requires corresponding inputs, outputs and relevant local intermediate values.
2. AMP skipped steps and exits on non-finite loss
GradScaler checks an optimizer's gradients. When it finds Inf or NaN, it skips that optimizer's update and adjusts the scale. Advancing the loop counter therefore does not necessarily mean that parameters were updated. Record attempted and actual updates to establish whether training is progressing.
An occasional skipped update caused by gradient overflow is different from training producing a non-finite loss and then exiting. A reduced scale shows that the guard responded. It does not show that the guard caused the anomaly or that the next iteration will recover.
Be explicit about where gradient inspection or clipping occurs. To use an unscaled gradient norm, call scaler.unscale_(optimizer) first. Comparing a scaled norm directly with a norm from ordinary training mixes in the scale factor.
Gradient clipping happens after forward and loss computation. Changing its threshold cannot repair an already non-finite loss computation.
Skipping an optimizer update is not a rollback of the iteration. A training forward may have updated BatchNorm statistics and consumed random numbers. Framework counters or schedulers may also have advanced. GradScaler controls gradient checks and the corresponding optimizer update; it does not undo those side effects. Replay should start from the pre-failure state rather than weights saved after a skipped step under an assumption that everything else is unchanged.
This is also why a successful FP32 run narrows the investigation without proving AMP was the root cause. Changing precision changes representable ranges, rounding and potentially kernels. It shows that the failure depends on the execution conditions. The smallest triggering computation still needs to be found; if the change merely expands the range, an abnormally growing intermediate may still exist.
3. Save the batch actually used in computation
Sample identifiers alone may not reproduce random crops, color transforms and other augmentations. To replay the failed forward, prioritize the actual model-input tensors, processed labels and metadata needed by the loss, such as sizes and masks.
Create independent snapshots so later in-place operations or buffer reuse cannot change the saved content. Record training/evaluation mode, autocast settings, input dtype and runtime environment. A label of “FP16” does not distinguish a full-model .half() conversion from AMP's per-operator precision selection.
Different recovery goals require different state:
| Recovery goal | Main material |
|---|---|
| Replay one forward and loss | Pre-failure model state, actual batch, mode and precision settings |
| Replay backward | Forward conditions plus required gradient-accumulation and scaling state |
| Replay one update | Optimizer and GradScaler state, pre-update gradients and training state |
| Replay the preceding trajectory | An earlier complete recovery point and continuous data/random-process state |
A forward containing random layers also depends on the RNG state before execution. RNG saved when catching an exception has usually already advanced during that execution and cannot be used directly as its starting point. Random streams used by multiple workers for prefetching and augmentation may not be captured by the main-process RNG either.
A failing batch and the nearest checkpoint are useful for exploration, but complete reproduction depends on the intervening steps. It is unsurprising that restoring weights and trying a few other unaugmented images may succeed.
PyTorch's reproducibility documentation separates control of random sources from deterministic algorithm selection. A seed specifies a random stream's starting point; mid-training recovery also needs its current position. Even identical random state does not guarantee bitwise agreement across versions, devices or execution algorithms. A diagnostic report should distinguish replaying one computation from continuing a whole training segment from an earlier point.
cuDNN benchmark and determinism settings serve different purposes. The former affects convolution algorithm search and selection; disabling it does not make every operator deterministic. Changes made to improve reproducibility should be recorded as diagnostic conditions. If the failure disappears, return to the original settings to check it. A changed environment is not evidence that the model was repaired.
4. Preserve the anomaly in the logs
Strict JSON writing commonly uses allow_nan=False. This avoids invalid non-finite JSON numbers, but can raise a new exception precisely when training most needs a diagnostic record.
Represent non-finite values explicitly, for example { "finite": false, "value": "nan" }, and preserve the field name, iteration and original exception. Replacing NaN with zero and continuing to use it in statistics disguises the failure as a normal sample.
A logging failure should not overwrite the original training exception. Preserve the original error, make a best effort to save tensors and a structured summary, and report any secondary recording error separately.
The earliest record in this case did not identify which serialized field first became non-finite. It cannot be classified as a logging-only fault. The later guard could record a non-finite loss and stop the run. This improved observability while leaving the model's numerical cause to be investigated.
Large tensors do not all belong in text logs. Record dtype, shape, finite-element count and minimum/maximum over finite elements, and save a few important tensors separately. Taking the maximum or mean of a tensor containing NaN may reduce the entire summary to NaN and lose the range of the other values. If no finite elements exist, record an empty result explicitly rather than calculating meaningless extrema.
5. Reproduce compiled non-finite outputs separately
In another execution check, the same weights produced finite outputs in eager FP16 but non-finite values under torch.compile. Its trigger differed from the NaN appearing after prolonged training. Keep separate reproductions for the two cases.
Align inputs, mode, caches and parameters before comparison. Derived caches that do not follow a precision change can introduce state differences first; see Precision conversion and invalidation of derived PyTorch caches. Once state is aligned, compare eager and compiled outputs. Check finiteness before calculating error, rather than keeping only an unexplained failed error assertion.
Further investigation can use backend="eager", "aot_eager" and "inductor" separately to distinguish graph capture, automatic-differentiation graph processing and code generation. This is a proposed diagnostic route. The existing record did not isolate a specific compiler stage. Passing after changing the input shape only shows that the triggering conditions changed.
Numerical-accuracy guidance also distinguishes ordinary error from non-finite results. Floating-point arithmetic is not associative. Fusion, reduction order and intermediate precision can introduce small finite differences. Inf or NaN requires separate checks for representable range and invalid operations. The absence of a bitwise-equality guarantee does not explain every compiled difference; failure on a compiled path alone does not establish a compiler defect either.
Calculating relative L2 directly on NaN outputs often returns NaN without showing how much output is affected. A more useful sequence is to compare the counts and positions of non-finite values, compare absolute and relative errors at jointly finite positions, then locate where the anomaly first appears in the output. When the denominator is near zero, inspect absolute error too, so a tiny absolute difference is not mistaken for a large numerical failure.
Module hooks can narrow the search, but functional operations inside a module do not automatically become separate modules. Added checks, forced synchronization and graph splitting can change execution. After finding a clue with probes, verify it using the originally failing execution path.
torch.autograd.detect_anomaly(check_nan=True) is another backward diagnostic tool. Running forward and backward inside its context can associate a failing backward operation with its forward stack and check for NaNs generated during backward. It adds overhead and does not replace forward checks for Inf, NaN and numerical ranges. Use it on a reproducible local example; do not compare throughput with anomaly detection enabled against ordinary training throughput.