Automatic Mixed Precision (AMP)

The core function of AMP is to automatically manage data types during the forward and backward passes of training. It identifies operations that can safely use FP16 (half-precision) for computation—like matrix multiplications in linear and convolutional layers—and keeps other operations, such as reductions and loss computations, in FP32 (single-precision). This eliminates the need for manual to(dtype) calls, reducing boilerplate code and potential errors. For example, in a standard ResNet-50 training loop, AMP would cast the input tensors and weights of convolutional layers to FP16 for the forward pass.

A critical technique to prevent gradient underflow in FP16. Gradients for some layers can have magnitudes smaller than the minimum positive value representable in FP16 (~5.96e-8), causing them to become zero. AMP applies dynamic or static loss scaling:

• The computed loss is multiplied by a scale factor (e.g., 128, 1024) before backpropagation.
• This shifts gradient values into the FP16 representable range.
• After the backward pass, gradients are unscaled before the optimizer step.
• Dynamic scaling automatically adjusts the scale factor by monitoring for gradient overflow (NaN/Inf values).

AMP maintains a copy of the model's parameters in FP32 precision, known as the master weights. During each training iteration:

1. Weights are cast from FP32 to FP16 for the forward pass.
2. Gradients are computed in FP16.
3. Gradients are cast back to FP32.
4. The FP32 master weights are updated by the optimizer using the FP32 gradients. This ensures the high-precision weight accumulation necessary for stable convergence, as the small updates from stochastic gradient descent can be lost in FP16. The master weights are the source of truth for the model state.

AMP uses internal policy engines to decide which operations to cast. A whitelist contains operations proven numerically safe and beneficial in FP16 (e.g., torch.mm, torch.conv2d). A blacklist contains operations that require FP32 for stability, typically those involving reductions (e.g., torch.sum, torch.softmax) or sensitive functions like torch.exp for large values. All other operations fall into a graylist and use FP16 if input types are FP16, or FP32 otherwise. This policy-based approach ensures critical numerical stability is not compromised for speed.

When using AMP, saving and loading model checkpoints requires handling both the FP32 master weights and the optimizer state (which may also be in FP32). Frameworks like PyTorch's torch.cuda.amp integrate with checkpointing utilities to automatically save and restore this additional state. The optimizer state often includes momentum buffers that must remain in FP32 for correctness. Loading an AMP checkpoint without the proper context will typically restore the model to a standard FP32 mode, requiring re-initialization of the AMP scaler for continued mixed-precision training.

The foundational technique that AMP automates. It strategically uses multiple numerical formats (e.g., FP16 for activations/gradients, FP32 for master weights) during neural network training. The goal is to accelerate computation and reduce memory usage while maintaining model convergence and final accuracy. AMP provides a library to manage this casting automatically, handling the complexities of maintaining a stable training process.

A 16-bit floating-point format defined by the IEEE 754 standard. It is the primary lower-precision type used in AMP for compute-intensive operations. Key characteristics:

• Halves memory footprint and bandwidth compared to FP32.
• Enables faster computation on hardware with native FP16 support (e.g., NVIDIA Tensor Cores, modern NPUs).
• Has a limited dynamic range (≈ 5.96e-8 to 65504), which can lead to numerical underflow for small gradient values—a primary challenge AMP's loss scaling addresses.

An alternative 16-bit format that uses an 8-bit exponent (like FP32) and a 7-bit mantissa. It is increasingly common in AMP frameworks for training. Advantages over FP16:

• Preserves the dynamic range of FP32, drastically reducing underflow risk for gradients.
• Sacrifices mantissa precision for this range, which is often less critical for deep learning robustness.
• Supported natively by hardware like Google TPUs, Intel AMX, and NVIDIA Hopper GPUs, making it a robust choice for automatic precision selection.

A critical technique integrated into AMP to prevent gradient underflow in FP16. The process:

1. The computed loss value is multiplied by a scaling factor (e.g., 128, 1024) before backpropagation.
2. This scales up the entire gradient chain, keeping small values within FP16's representable range.
3. Gradients are unscaled before the weight optimizer step to apply the correct update. AMP typically implements dynamic loss scaling, which automatically adjusts the scale factor by monitoring for gradient overflow.

A more aggressive precision-reduction technique that AMP can complement. While AMP dynamically manages FP16/FP32 during training, QAT simulates INT8 or lower integer precision throughout the forward pass. This allows the model to learn to compensate for quantization error. AMP and QAT can be staged: use AMP for faster initial training, then apply QAT to prepare a model for ultra-efficient integer-only inference on NPUs or edge devices.

The core numerical instability risks that AMP is designed to mitigate.

• Underflow: When a value is smaller than the smallest positive number representable in a format (e.g., a gradient < ~5.96e-8 in FP16). It becomes zero, halting learning. AMP uses loss scaling to combat this.
• Overflow: When a value exceeds the maximum representable number (e.g., > 65504 in FP16). This leads to saturation or infinity. AMP's dynamic loss scaling also detects and reduces the scale factor if overflow occurs, ensuring training stability.