Fully Sharded Data Parallel (FSDP) — Memory-Efficient Training Tool

As deep learning models continue to grow in size, from millions to billions of parameters, the limitations of traditional parallel training strategies become increasingly apparent. One of the most well-known strategies like this is Distributed Data Parallel (DDP).

Distributed Data Parallel (DDP) has long been the standard approach for scaling deep learning workloads across multiple GPUs. It achieves parallelism by replicating the full model on each GPU, distributing mini-batches of data to each device, and then averaging gradients across GPUs during backpropagation. While effective for small to medium-sized models, DDP’s design introduces significant memory overhead; every GPU holds a complete copy of the model, along with duplicated gradients and optimizer states. As models grow larger, this redundancy becomes a serious bottleneck, limiting memory efficiency and forcing practitioners to make trade-offs in model complexity or hardware usage.

This is where Fully Sharded Data Parallel (FSDP) comes into play, which is a PyTorch’s native implementation of a ZeRO Stage 3-like system.

Fully Sharded Data Parallel (FSDP) is a data parallel training technique that combines the advantages of data and model parallelism for distributed training by sharding the model parameters, gradients, and optimizer states across all available GPUs, significantly reducing the memory usage and enabling the training of large-scale models that would otherwise exceed memory limits. Additionally, FSDP offers the option to offload parameters and optimizer states to the CPU when they are not actively needed, enabling even larger models to be trained on limited GPU memory.

Working of FSDP

FSDP Training

One of the key innovations of FSDP is how it manages memory during training. Each GPU holds only a shard of the model’s parameters. During the forward pass, FSDP gathers only the necessary parameters from other GPUs using an All Gather operation, performs the computation, and then immediately releases the full parameters to free up memory. Some of the key benefits of using the All Gather are:

1. No GPU is sitting idle
2. This highly reduces the overhead communication as GPUs do not communicate with each other.

This process is repeated in the Backward pass, parameters are re-gathered using the All Gather operation, gradients are computed, and then the averaged gradients are redistributed using the Reduce Scatter operation after Backward pass, so that each GPU ends up with an updated shard of the weights. This approach ensures maximum memory efficiency, making it possible to train very large models with significantly lower GPU memory requirements.

Zero Redundancy Optimizer (ZeRO)

ZeRO (Zero Redundancy Optimizer) by DeepSpeed is a set of strategies designed to reduce memory usage during distributed training. Instead of replicating the entire model, gradients, and optimizer states across GPUs, ZeRO partitions these elements across devices.

ZeRO Stage 1: Optimizer States

The optimizer states are sharded across all GPUs. Each GPU holds and updates only its assigned portion of the optimizer states, significantly reducing redundancy.

ZeRO Stage 2: Optimizer States + Gradients

In addition to optimizer states, gradients are also partitioned across GPUs. This further reduces memory usage, as each device now stores only a subset of both gradients and optimizer states.

ZeRO Stage 2: Optimizer States + Gradients + Parameters

This is the most memory-efficient stage, and the foundation upon which FSDP is built. In this stage, model parameters themselves are sharded across GPUs along with gradients and optimizer states. Parameters are gathered only when needed (e.g., during forward and backward passes), and released afterward to maximize memory savings.

CPU Offloading and Activation Checkpointing

FSDP enhances memory efficiency by integrating two powerful techniques: activation checkpointing and CPU offloading.

Activation checkpointing addresses one of the key sources of memory consumption during training, intermediate activations generated in the forward pass. Instead of storing these activations, which can be substantial in deep neural networks, FSDP discards them and re-computes them as needed during the backward pass. This approach trades a moderate increase in computation for a significant reduction in memory usage, allowing deeper and more complex models to be trained without running out of GPU memory.

CPU offloading, a feature that moves model parameters and optimizer states to CPU memory when they are not actively being used. By storing these components on the CPU and fetching them only when needed, FSDP further reduces the strain on GPU memory.

In conclusion, FSDP significantly reduces memory overhead compared to traditional approaches like DDP. This enables practitioners to train larger models or increase batch sizes without compromising performance, making FSDP a powerful tool for modern large-scale deep learning workloads.

References

1. https://kye.hashnode.dev/explaining-zero-fsdp-to-non-ml-engineers
2. https://engineering.fb.com/2021/07/15/open-source/fsdp/
3. https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/

I’d love to hear your thoughts, ideas, or experiences with FSDP. Let’s connect on LinkedIn and share knowledge, collaborate, and explore more about building smarter AI systems together.

👉 **Connect with me on LinkedIn**

Looking forward to learning and growing with this amazing community!