Start with a baseline that meets the need
Choose a short dataset with ordinary inputs, edge cases, and decision boundaries. Freeze the model, weights, preprocessing, and train or eval mode. A comparison between two models or two batches does not let you attribute their difference to precision.
Record the output that matters to your application, not just the loss. For a classifier, that may include scores and decisions; for a regression, the error and extreme values. First check whether NaN or inf is already present in the baseline. An incorrect FP32 run does not become a reliable foundation just because it has more bits.
Set tolerance, minimum quality, and absence of non-finite values before the test. PyTorch reminds us that floating-point computation does not guarantee identical results across devices or execution paths.
Distinguish autocast, numeric format, and GradScaler
autocast chooses the type of certain operations according to their compute policy. It does not convert the entire program into a single format. With this usage, avoid manually converting the whole model with half(). The current documentation recommends torch.autocast or torch.amp.autocast; the older torch.cuda.amp interfaces are deprecated.
GradScaler acts on the scale of the loss and gradients during training. It is not used as an inference accelerator, since inference does not perform a backward pass. FP16 has a narrower numeric range than BF16; a model designed for BF16 can overflow in FP16. A repeated drop in the scale therefore does not establish that the problem is resolved.
Choose the format based on the model's constraints and the operations actually used, then verify support on the target. A card's commercial name or a PyTorch preparation preference does not prove that your custom operator has the desired kernel.
Scroll the table to read all columns.| Choice | Role | Check required |
|---|---|---|
| FP32 reference | Project's point of comparison | Finite outputs and expected quality |
| FP16 autocast | Some operations in reduced precision | Numeric range and gradients |
| BF16 autocast | A different range/precision trade-off | Available operators and quality |
| GradScaler | Gradient scale handling | Updates actually applied |
Put the training steps in the right order
The proposed snippet assumes a model and optimizer that have already been built, an input and a target on the same GPU, and a scalar loss. It has not been executed and does not constitute validation of an offering. The autocast context wraps the forward pass and the loss; the backward pass takes place after it closes. The scaler is created once for the training session, not for each batch.
To inspect or clip the gradients, first remove their scale factor with unscale_. The official AMP examples specify doing this only once per optimizer and after accumulating the gradients intended for its update. The clipping threshold of 1.0 below is an illustrative value to choose for your project, not a universal recommendation.
Here the guards stop the diagnostic if the loss, gradients, or total norm are not finite. These CPU reads are intrusive: do not time this snippet. Accumulation, multiple optimizers, and a scheduler each require their own definition of the update.
import torch
# Preconditions: model, optimizer, loss_fn, inputs, and targets exist.
# The model and inputs are on the same CUDA/HIP device.
dtype = torch.float16 # Choice to validate; BF16 is another option to try.
scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))
# Place inside your loop, with scaler kept between batches.
optimizer.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=dtype):
prediction = model(inputs)
loss = loss_fn(prediction, targets)
if not bool(torch.isfinite(loss).item()):
raise FloatingPointError("Non-finite loss: stop the diagnostic")
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
if any(p.grad is not None and
not bool(torch.isfinite(p.grad).all().item())
for p in model.parameters()):
raise FloatingPointError("Non-finite gradient: stop the diagnostic")
torch.nn.utils.clip_grad_norm_(
model.parameters(), max_norm=1.0, error_if_nonfinite=True,
)
scaler.step(optimizer)
scaler.update()Worked example: two closely related decisions are not interchangeable
Suppose a service that chooses the class with the highest score. On an educational input, the reference produces two very close scores: 1.0000 and 1.0003. Another numeric path could change their order or produce a tie. These numbers illustrate a decision boundary; they are not measured FP16 or BF16 outputs.
Proper verification has two levels. Compare the scores with explicit tolerances, then compare the decision and the tie-breaking rule applied. A small difference in absolute value can change the action chosen. Conversely, a visible numeric difference can remain inconsequential for a task whose threshold lies far from the observed scores.
Record identifiers, reference outputs, the AMP trial and its impact on the decision. Set the acceptance rule before reading the results. Do not widen the tolerance to make a troublesome case disappear; outputs at different scales may require distinct criteria.
Scroll the table to read all columns.| Criterion | Reference | AMP trial | Decision |
|---|---|---|---|
| Finished outputs | To check | To check | Reject unexplained non-finite values |
| Numerical deviation | Retained values | Deviation to compute | Tolerance defined before the trial |
| Application decision | Class or action | Class or action | Review changes |
| Quality on the fixed set | To measure | To measure | Meet the project threshold |
Interpreting NaNs and skipped updates
When non-finite values appear, look for the first step that produces them: input, intermediate output, loss or gradient. Replay the same case as a reference, then locally disable autocast around the suspect operation while also checking the dtype of its inputs. Running an entire training in FP32 can serve as a comparison, but it does not automatically localize the problem.
The scaler may skip an update when gradients contain inf or NaN. A loop that keeps running has therefore not necessarily performed as many updates as iterations. Record this behavior during diagnosis. Do not blindly advance a learning policy assumed to follow the actual updates.
A finite loss does not guarantee finite gradients. Conversely, a one-off incident is not enough to declare a training unusable: examine its frequency, progress and quality. The AMP recipe provides a method to isolate autocast and scaling separately when one of the two is suspect.
Preparing a reproducible rollback
Before the trial, keep the reference configuration, the weights, the optimizer state and a consistent checkpoint. If your training uses a scaler, its state is also part of the resume. Document the dtype and any regions left in FP32. Resuming with a different policy is an experimental change to identify, not an implicitly equivalent continuation.
Return to the previous configuration if outputs become non-finite, if quality falls outside the set criterion, or if updates stop progressing in a usable way. Keep the case that prompted this rollback. After a change, restart the comparison on the same set before extending the duration.
The Kernodeck checkpoints exercise verifies a CPU resume without AMP. Reuse its comparison method, adding the states actually consumed by your loop.
Measuring gains only after numerical validation
After validation, measure memory and time without the detailed diagnosis. Keep shapes, batch, model and quality. Scalar reads, synchronizations and profilers can alter the durations; remove intrusive checks from the final measurement.
On ROCm, the PyTorch device name remains cuda and the corresponding interfaces are reused. This guarantees neither the same kernels nor identical results to NVIDIA. Verify the backend and the project's operators on the target. This guide announces no fixed memory reduction, speed multiplier or Kernodeck preparation compatibility.