FSDP › Mixed Precision with FSDP
Language AI HandbookAll chapters FSDP › Mixed Precision with FSDPOne section · Sections in this chapter7 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
FSDP has first-class support for mixed precision training through the MixedPrecision policy. This lets you specify different dtypes for parameters, gradients, and buffer tensors independently.
Why Mixed Precision Matters at Scale
On modern GPUs, bfloat16 matrix multiplications run significantly faster than float32 because the hardware has dedicated Tensor Core units optimized for reduced-precision arithmetic. On A100 GPUs, bfloat16 matrix multiplications deliver up to 312 TFLOPS, compared to 77.6 TFLOPS for float32, a 4x speedup. Beyond raw throughput, reduced-precision tensors are half the size of float32, which matters both for the all-gather buffer (which temporarily holds full-precision parameters) and for activations accumulated during the forward pass.
The challenge with reduced precision is numerical stability: float16 has a limited exponent range (maximum value ), causing overflow during training when gradients or activations become large. Bfloat16 solves this by using the same exponent range as float32 (8 bits), sacrificing mantissa precision instead (7 bits vs. 23 bits for float32). For most training scenarios, bfloat16 provides the speed benefits of fp16 without the overflow instability.
The standard mixed precision configuration for large model training stores reduced-precision parameters during computation (to speed up matrix multiplications and save memory on the all-gather buffer) while keeping full-precision optimizer states for numerical stability.
import torch
from torch.distributed.fsdp import MixedPrecision
# BF16 mixed precision policy
bf16_policy = MixedPrecision(
param_dtype=torch.bfloat16, # parameters stored and communicated in BF16
reduce_dtype=torch.float32, # gradients reduced in FP32 for numerical stability
buffer_dtype=torch.bfloat16, # buffers (e.g., layer norm running stats) in BF16
)Setting reduce_dtype=torch.float32 ensures that gradient accumulation happens in float32, preventing the loss of small gradient values that would occur if you accumulated in bfloat16 (where the minimum representable positive value is much larger than in float32). The parameters and all-gather buffers are in bfloat16, giving you the memory and compute benefits, while the gradients that drive the optimizer update retain float32 precision.
BF16 is generally preferred over FP16 for large model training because its exponent range matches FP32, eliminating the need for loss scaling that FP16 requires. Loss scaling is a workaround where gradients are multiplied by a large constant before the backward pass to prevent fp16 underflow, then divided back before the optimizer step. Managing loss scaling adds code complexity and can cause instability if the scale factor is poorly chosen.