Mixed Precision Training

Using 16-bit floating-point (FP16) or Brain Floating Point (BF16) halves the memory required per parameter compared to standard 32-bit (FP32). This enables:

• Training larger models or using larger batch sizes on the same GPU hardware.
• Storing more optimizer states and activations in memory.
• A direct reduction in GPU memory bandwidth pressure, which is often a training bottleneck.

Modern AI accelerators like NVIDIA Tensor Cores and Google TPUs are optimized for lower-precision matrix operations. Using FP16/BF16 allows these specialized units to perform many more FLOPS (Floating Point Operations Per Second). This results in:

• Faster forward and backward passes during training.
• Significant reductions in per-iteration training time.
• More efficient utilization of expensive hardware.

A pure lower-precision training can fail due to numerical underflow (values rounding to zero) and overflow (values exceeding range). Mixed precision maintains stability via:

• Master Weights: Keeping a master copy of weights in FP32 for precision.
• Loss Scaling: Automatically scaling the loss value before backward pass to preserve small gradients.
• Precision Casting: Performing sensitive operations (e.g., weight updates, normalization layers) in FP32.

BF16 (bfloat16) is an alternative 16-bit format that preserves the same exponent range as FP32 (8 bits) while reducing mantissa bits (7 bits). This design offers a key advantage:

• Greatly reduces the risk of overflow/underflow compared to FP16.
• Often requires less or no loss scaling.
• Becoming the preferred format on hardware that supports it (e.g., TPUs, NVIDIA Ampere+ GPUs).

Mixed precision training is a foundational technique that combines effectively with other scaling methods:

• Gradient Checkpointing: Saves memory on activations, which are now stored in lower precision.
• Model Parallelism & ZeRO: Further reduces memory footprint when partitioned weights are in lower precision.
• Quantization-Aware Training (QAT): Serves as a precursor for deploying quantized models for inference.

Floating-point precision defines the format for representing real numbers in computing, balancing range and accuracy. Lower precision types like FP16 (16-bit) and BF16 (bfloat16) use fewer bits, reducing memory footprint and accelerating computation but with a risk of numerical overflow/underflow. FP32 (32-bit single-precision) is the traditional standard, offering greater numerical stability for critical operations like weight updates. Mixed precision training strategically uses both: low precision for most tensor operations and high precision for a master copy of weights and sensitive calculations.

Gradient checkpointing is a memory optimization technique that trades compute for memory. During the forward pass, it selectively saves only a subset of layer activations (the checkpoints) instead of all. During the backward pass, it re-computes the unsaved intermediate activations from the nearest checkpoint. This dramatically reduces memory consumption, often by a factor of 4-5x, enabling the training of larger models or longer sequences. It is frequently used in conjunction with mixed precision training to manage total memory overhead.

ZeRO is a family of memory optimization techniques for distributed data-parallel training, implemented in libraries like Microsoft's DeepSpeed. It partitions the three primary model states—optimizer states, gradients, and parameters—across GPUs to eliminate memory redundancy.

• ZeRO Stage 1: Partitions optimizer states.
• ZeRO Stage 2: Partitions gradients as well.
• ZeRO Stage 3: Partitions all model parameters. ZeRO enables the training of models with trillions of parameters. When combined with mixed precision, ZeRO-Offload can further move optimizer states and gradients to CPU memory.

Quantization is a model compression technique that reduces the numerical precision of a model's weights and activations, typically from 32-bit floats to 8-bit integers (INT8) or 4-bit integers. Unlike mixed precision training, quantization is primarily used for inference to reduce model size and latency. However, Quantization-Aware Training (QAT) simulates lower precision during training to preserve accuracy. Related techniques like QLoRA combine 4-bit quantization with parameter-efficient fine-tuning, enabling fine-tuning of large models on consumer hardware.

Automatic Mixed Precision (AMP) is an implementation tool that automates the application of mixed precision training. Libraries like NVIDIA's Apex (PyTorch) and TensorFlow's native AMP handle:

• Cast selection: Automatically choosing FP16 or FP32 for each operation.
• Loss scaling: Automatically scaling the loss to prevent gradient underflow in FP16.
• Master weight updates: Maintaining a master copy of weights in FP32 for precision. AMP allows developers to implement mixed precision with minimal code changes, abstracting away the manual management of casts and scaling.

Tensor Cores are specialized processing units on modern NVIDIA GPUs (Volta architecture and later) designed to perform mixed-precision matrix operations at extremely high throughput. They are optimized for operations like D = A * B + C, where A and B are FP16 matrices, C and D can be FP16 or FP32. This hardware acceleration is the primary driver for the speedups (up to 3x) seen in mixed precision training. Utilizing Tensor Cores requires operations to meet specific size alignment criteria (e.g., matrix dimensions multiples of 8).