FSDP › The Memory Problem at Scale

Before examining how FSDP works, it helps to understand precisely where memory goes during training. The memory breakdown directly determines which GPU configurations can run which models, and misunderstanding it leads to mysterious out-of-memory crashes in production training runs.

For a model with $P$ parameters stored in float32, memory consumption breaks down across four components:

• Parameters: $4P$ bytes (4 bytes per float32 value)
• Gradients: $4P$ bytes (same shape as parameters, one gradient per parameter)
• Optimizer states: $8P$ bytes for Adam (first and second moment estimates, each 4 bytes per parameter)
• Activations: variable, roughly proportional to batch size and sequence length

The total for a training step without any memory optimization is approximately $16P$ bytes. For GPT-3 with $P = 175$ billion parameters, that amounts to:

$$16 \times 175 \times 10^9 \text{ bytes} \approx 2.8 \text{ terabytes}$$

That figure represents nearly 70 A100 GPUs worth of memory, just for the model state, before counting activations. Even the humble 7B-parameter models that became common in the open-source community after the Llama release require:

$$16 \times 7 \times 10^9 \text{ bytes} \approx 112 \text{ GB}$$

A single A100 with 80 GB of memory cannot hold the full training state for even a 7B model, let alone its activations.

Many people new to large model training expect parameters to dominate memory. In practice, the Adam optimizer states are the largest single component at $8P$ bytes, accounting for half of the $16P$ total. This happens because Adam maintains two running averages per parameter: the first moment (exponential moving average of gradients) and the second moment (exponential moving average of squared gradients). Both are kept in float32 for numerical stability even when training in lower precision, because accumulated rounding errors in these statistics degrade convergence.

When using mixed precision training, the picture shifts somewhat. Parameters and gradients may be stored in float16 or bfloat16 (2 bytes each), but the optimizer states and a master copy of the parameters remain in float32. A common mixed precision accounting gives approximately $2P$ (fp16 params) $+ 2P$ (fp16 grads) $+ 4P$ (fp32 master params) $+ 8P$ (fp32 optimizer states) $= 16P$ bytes total, the same as pure float32. The mixed precision savings come from smaller activation memory (since intermediate activations are in fp16) and from faster matrix multiplications on hardware with Tensor Cores, not from reduced parameter or optimizer state memory.

Standard DistributedDataParallel (DDP) replicates the full model on every GPU. With 8 GPUs, you have 8 complete copies of the model, 8 complete sets of gradients, and 8 complete sets of optimizer states. The total memory consumption is exactly the same per GPU as running on a single device. The benefit of DDP is purely throughput: you can process 8 times as many data samples per step, but you cannot train a model that does not fit on one GPU.

This is the fundamental limitation that sharding addresses. Rather than replicating the model across GPUs, sharding divides it. Each GPU is responsible for a fraction of the total parameter space, and the GPUs communicate to temporarily share parameters when computation requires them.