Mixed Precision Training

The master copy of the model's parameters is maintained in full 32-bit floating-point (FP32) precision. This serves as the single source of truth. During training:

• Forward and backward passes use lower-precision copies (e.g., FP16/BF16).
• The FP32 master weights accumulate the small gradient updates with high numerical stability.
• This prevents the vanishing gradient problem where tiny updates in low precision could underflow to zero.

The bulk of tensor operations—matrix multiplications, convolutions, and activation functions—are performed using half-precision (16-bit) formats to leverage hardware acceleration.

• FP16 (float16): Offers a 5-bit exponent and 10-bit mantissa. Common on NVIDIA Tensor Cores but has a limited dynamic range (~65,504), risking overflow/underflow.
• BF16 (bfloat16): Uses an 8-bit exponent (matching FP32) and a 7-bit mantissa. It preserves the dynamic range of FP32, making it more stable for training and the preferred format on Google TPUs and newer NVIDIA/AMD GPUs.
• This reduces memory bandwidth and increases computational throughput by 2-8x on supported hardware.

A critical technique to prevent gradient underflow when using FP16. Small gradient values can fall below the minimum representable range of FP16 (~5.96e-8). The system:

1. Scales up the loss value by a large, constant factor (e.g., 128, 1024) before backpropagation.
2. This amplifies the entire gradient chain, bringing values into FP16's representable range.
3. The optimizer applies the scaled gradients to the FP32 master weights.
4. The weights are then correctly updated as the scaling factor is effectively divided out. Dynamic loss scaling algorithms automatically adjust this factor during training.

An abstraction layer, such as NVIDIA's Apex AMP or PyTorch's torch.cuda.amp, that automates the precision casting and loss scaling. It handles:

• Operator Casting: Automatically decides which operations require FP32 for stability (e.g., reductions, softmax) and which can safely use FP16/BF16.
• Gradient Management: Applies loss scaling and safely unscales gradients before the optimizer step.
• Context Managers: Provides simple Python decorators or context managers (e.g., autocast, GradScaler) to wrap the training loop, minimizing code changes.

Specialized processing units within modern GPUs (NVIDIA's Tensor Cores, AMD's Matrix Cores, Google's TPU MXUs) that perform mixed-precision matrix operations at extreme speeds. They are the physical enablers of the performance gains.

• These cores are designed to compute D = A * B + C, where A and B are in FP16/BF16, and C and D can be in FP16, BF16, or FP32.
• They perform fused multiply-add (FMA) operations, completing a full matrix operation in a single clock cycle, which is significantly faster than using standard CUDA cores for FP32 math.

Optimizers must be adapted to work with the two-weight system (FP16 model copy, FP32 master weights).

• AdamW, SGD, etc.: Their internal states (momentum, variance) are typically maintained in FP32 for stability, even when updating weights derived from FP16 gradients.
• Fused Optimizers: Advanced implementations combine the weight update, unscaling, and master weight copy into a single, fused GPU kernel. This reduces memory access overhead and is a key optimization in frameworks like NVIDIA's Apex.
• The optimizer step is the point where the FP32 master weights are updated and then copied back down to the FP16 model copy for the next forward pass.

Specialized AI training chips are designed with mixed precision as a first-class citizen.

• Intel Habana Gaudi: Features dedicated matrix multiplication engines that efficiently handle BF16, with support for custom data types like FP8 for further optimization. Its SynapseAI software stack integrates with PyTorch and TensorFlow.
• AWS Trainium: Amazon's custom chip, available on AWS EC2 Trn1 instances, is built for high-performance training. It natively supports a range of precisions (BF16, FP16, TF32, FP8) and includes a Neuron SDK that compiles framework models to optimize for these data types, managing precision conversions and scaling automatically.

A process where a neural network is trained or fine-tuned with simulated quantization noise in the forward pass. This allows the model to learn parameters that are robust to the precision loss incurred during subsequent integer deployment, ensuring accuracy is preserved when the model is converted to low-bit formats like INT8.

• Purpose: Bridges the accuracy gap between a full-precision model and its quantized version.
• Process: Insert 'fake quantization' nodes during training to mimic the rounding and clamping effects of true integer arithmetic.
• Contrast with Mixed Precision: QAT prepares a model for inference at low precision, while Mixed Precision optimizes the training process itself using a mix of precisions.

An I/O-aware, exact attention algorithm that recomputes attention scores on-the-fly within fast SRAM (Static Random-Access Memory) to avoid reading and writing the large attention matrix to slow HBM (High Bandwidth Memory). This dramatically speeds up training and inference and reduces memory footprint, which is complementary to Mixed Precision Training's goal of efficient computation.

• Core Innovation: Optimizes memory hierarchy usage, making attention memory-bound rather than compute-bound.
• Synergy with Mixed Precision: Often used in conjunction, as lower precision (FP16/BF16) further reduces the memory traffic and compute that FlashAttention optimizes.
• Impact: Enables training of longer-context models, a critical capability for modern LLMs.

A technique to trade compute for memory during training by selectively discarding intermediate activations in the forward pass and recomputing them during the backward pass. This enables the training of larger models or the use of larger batch sizes with limited GPU memory.

• Mechanism: Only stores activations at certain 'checkpoint' layers. Non-checkpointed activations are recomputed during backpropagation.
• Relationship to Mixed Precision: Both are memory-saving techniques. Mixed Precision reduces the bytes per activation; Gradient Checkpointing reduces the number of activations stored. They are frequently used together to maximize model size within hardware constraints.
• Trade-off: Introduces a significant computational overhead (typically ~30%) in exchange for a substantial memory reduction (~70%).

Empirically derived relationships that predict how a model's performance (loss) scales as a power-law function of its size (parameters), dataset size, and computational budget. These laws guide the efficient allocation of resources during pre-training.

• Key Insight: Performance improves predictably with increased scale, but the optimal balance between model size (N), data size (D), and compute (C) is non-obvious.
• Compute-Optimal Scaling (Chinchilla): The principle that for a fixed compute budget, optimal performance is achieved by scaling model size and training data jointly, often requiring more data than previously thought.
• Connection to Mixed Precision: Mixed Precision Training is a direct enabler of scaling, as it allows for training larger models (more parameters) or using larger batches (more data) within the same memory and compute budget, directly impacting the variables governed by scaling laws.

A parameter-efficient fine-tuning (PEFT) method that freezes the pre-trained model weights and injects trainable low-rank decomposition matrices into transformer layers. It drastically reduces the number of trainable parameters for downstream task adaptation.

• Mechanism: Represents weight updates as ΔW = BA, where B and A are low-rank matrices. Only these small matrices are trained.
• Precision Context: LoRA's trainable adapters are typically kept in full precision (FP32) during fine-tuning, even if the base model uses a lower precision format. This is a form of targeted mixed precision within a fine-tuning workflow.
• Efficiency Synergy: LoRA reduces storage and memory for task-specific weights, while Mixed Precision reduces the memory and compute for the forward/backward passes through the frozen base model.

A 16-bit floating-point format designed by Google for machine learning. It preserves the dynamic range of FP32 (8 exponent bits) while truncating the mantissa/significand (7 bits vs. FP32's 23). This makes it a superior choice over FP16 for Mixed Precision Training.

• Key Advantage over FP16: Much lower risk of underflow and overflow due to its identical exponent range to FP32. Gradient values are more stable.
• Hardware Support: Natively supported on modern AI accelerators (e.g., NVIDIA A100+ Ampere/ Hopper GPUs, Google TPUs, Intel Habana Gaudi).
• Standard Practice: BF16 is increasingly the default lower-precision format for Mixed Precision Training in large-scale model development, often replacing the FP16/FP32 master copy paradigm with a simpler BF16/FP32 approach.