FSDP › Worked Example: Memory Savings on a 7B Model

To build intuition for the magnitudes involved, let's trace through the memory accounting for a concrete scenario: training a 7 billion parameter model on 8 A100-40GB GPUs.

With standard DDP, every GPU holds the complete model state:

• Parameters: $4 \times 7 \times 10^9 = 28$ GB
• Gradients: $28$ GB
• Optimizer states (Adam): $56$ GB
• Total model state: $112$ GB

This already exceeds the 40 GB GPU capacity by $2.8\times$. DDP cannot run this configuration at all.

With FULL_SHARD, each GPU holds $1/8$ of each component:

• Parameters (shard): $28 / 8 = 3.5$ GB
• Gradients (shard): $3.5$ GB
• Optimizer states (shard): $7$ GB
• Total steady-state: $14$ GB

During the forward and backward passes, a single FSDP unit's all-gather buffer temporarily consumes an additional $28/M$ GB where $M$ is the number of FSDP units. For a 32-layer model wrapped at the layer level, each all-gather buffer is $28/32 \approx 0.875$ GB, bringing peak memory to about $14.875$ GB per GPU, well within the 40 GB limit.

The remaining memory budget of approximately 25 GB can accommodate activations from long sequence lengths and larger batch sizes. Activations for a typical transformer layer scale as $O(B \times S \times d)$ where $B$ is batch size, $S$ is sequence length, and $d$ is hidden dimension. For a 7B model ($d = 4096$), a batch of 4 sequences of length 2048 in bfloat16 generates roughly $4 \times 2048 \times 4096 \times 2 \approx 67$ MB per layer before considering the activations needed for attention score computation. With activation checkpointing (recomputing rather than storing activations), this budget extends further.

This example illustrates why FSDP enabled more open-source large model training in 2023 and 2024. Before FSDP became widely available and well-documented, training 7B models required either expensive 80 GB GPUs (A100-80GB or H100), proprietary systems like Google's TPU pods, or careful hand-engineering with DeepSpeed. With FSDP on standard 40 GB A100 GPUs, which were broadly available through cloud providers, the research community could train and fine-tune these models on accessible hardware.