Fully Sharded Data Parallel (FSDP) and ZeRO: How Memory Sharding Eliminates Redundant Model States in Distributed Training

Training large language models across distributed GPU clusters introduces a fundamental memory bottleneck. In traditional Distributed Data Parallel (DDP) setups, every GPU maintains an identical copy of model weights, optimizer states, and gradients while processing independent data batches. As models scale from billions to hundreds of billions of parameters, static model states quickly exceed single-device High Bandwidth Memory (HBM) capacity.

The Zero Redundancy Optimizer (ZeRO) paradigm, developed by Microsoft Research and DeepSpeed, and its native PyTorch counterpart, Fully Sharded Data Parallel (FSDP), eliminate this memory redundancy. By partitioning model states across data-parallel ranks and materializing full tensors on demand, ZeRO and FSDP allow clusters to scale model capacity proportionally to the total aggregate GPU memory without requiring complex manual model refactoring.

The Static Memory Footprint of Distributed Training

In mixed-precision training using 16-bit floating-point formats (FP16 or BF16) paired with the standard Adam optimizer, static model memory consumption is dominated by three components:

• Model Parameters: 2 bytes per parameter in FP16 or BF16 representation.
• Gradients: 2 bytes per parameter in FP16 or BF16 representation.
• Optimizer States: 12 bytes per parameter under standard AdamW. This comprises 4 bytes for the FP32 master weight copy (to preserve numerical precision during weight updates), 4 bytes for the FP32 first momentum vector, and 4 bytes for the FP32 second momentum variance vector.

Summing these components yields 16 bytes per parameter in baseline static state memory.

For a 70-billion-parameter model, static memory alone requires:

Static Memory = 70 * 10^9 parameters * 16 bytes/parameter = 1,120 GB

Because an 80 GB NVIDIA A100 or H100 GPU can only hold a fraction of this state, traditional DDP cannot run a 70B model even with a batch size of 1. Beyond static states, additional dynamic memory is consumed by forward activation tensors, KV caches, temporary workspace buffers, and memory fragmentation.

The ZeRO Partitioning Stages

The ZeRO paper by Rajbhandari et al. (2019) introduced three distinct stages of memory optimization that systematically remove state duplication across $N$ data-parallel ranks.

ZeRO-Stage 1: Optimizer State Partitioning (P_os)

In ZeRO-1, model weights and gradients remain fully replicated on every GPU, but the 12 bytes per parameter of Adam optimizer states are partitioned equally across all $N$ data-parallel workers.

• Memory footprint: 2 bytes (weights) + 2 bytes (gradients) + (12 / N) bytes (optimizer states).
• Memory reduction: For large worker counts ($N \ge 64$), state memory drops from 16 bytes per parameter down to approximately 4 bytes per parameter, achieving a 4x reduction.
• Communication volume: Identical to standard DDP. At the end of the backward pass, gradients are gathered and reduced to their corresponding partition via Reduce-Scatter, local optimizer steps are executed on the sharded states, and updated parameters are distributed across workers via All-Gather. Total communication volume remains 2x the model parameter size per step.

ZeRO-Stage 2: Gradient Partitioning (P_os+g)

ZeRO-2 extends partitioning to gradients. As backpropagation traverses layers in reverse order, each GPU computes gradients but only retains the gradient slice corresponding to its assigned optimizer state partition (2 / N bytes per parameter). Gradients for other partitions are immediately freed.

• Memory footprint: 2 bytes (weights) + (2 / N) bytes (gradients) + (12 / N) bytes (optimizer states).
• Memory reduction: At scale, memory drops to approximately 2 bytes per parameter, yielding an 8x reduction over baseline DDP.
• Communication volume: Communication volume remains identical to DDP and ZeRO-1 (2x model size per step). Reduce-Scatter operations are executed immediately as each layer finishes backpropagation, overlapping network transfer with upstream backward computation.

ZeRO-Stage 3: Parameter Partitioning (P_os+g+p)

ZeRO-3 shards all three core components: optimizer states, gradients, and model weights. Each GPU holds only a 1 / N slice of every tensor.

• Memory footprint: 16 / N bytes per parameter in total static memory.
• Memory reduction: Memory scales linearly with cluster size (1 / N). On a 64-GPU cluster, static memory for a 70B model drops from 1,120 GB to 17.5 GB per GPU, comfortably fitting within standard GPU VRAM alongside activation memory.
• Communication volume: ZeRO-3 introduces an additional communication phase. During the forward pass, an All-Gather operation reconstructs the full weights for a specific layer right before its computation, and the unsharded weights are discarded immediately afterward. During the backward pass, weights are gathered a second time via All-Gather to compute gradients, followed by a Reduce-Scatter on the computed gradients. This increases total communication volume to 3x model size per step (a 1.5x increase over DDP).

Summary of State Partitioning Dynamics

• Standard DDP: 2 bytes weights, 2 bytes grads, 12 bytes optimizer states. Total: 16 bytes/param. Communication: 2x model size.
• ZeRO-1 (P_os): 2 bytes weights, 2 bytes grads, 12 / N bytes optimizer states. Total: 4 + 12/N bytes/param. Communication: 2x model size.
• ZeRO-2 (P_os+g): 2 bytes weights, 2 / N bytes grads, 12 / N bytes optimizer states. Total: 2 + 14/N bytes/param. Communication: 2x model size.
• ZeRO-3 / FSDP (P_os+g+p): 2 / N bytes weights, 2 / N bytes grads, 12 / N bytes optimizer states. Total: 16 / N bytes/param. Communication: 3x model size.

PyTorch Fully Sharded Data Parallel (FSDP)

While ZeRO was originally implemented within the DeepSpeed library, Meta and the PyTorch team developed and upstreamed Fully Sharded Data Parallel (FSDP) as a native distributed primitive in PyTorch (torch.distributed.fsdp).

Modular Units and FlatParameter

FSDP structures models into nested, modular execution units (typically wrapping each transformer block or decoder layer). Inside each unit:

1. Parameter Flattening: FSDP flattens all weights and biases within a module into a single 1D tensor (FlatParameter). This contiguous memory layout avoids small, fragmented CUDA memory allocations and allows NCCL collective operations to run at peak interconnect bandwidth.
2. Just-in-Time Materialization: Individual FSDP units unshard their internal parameters just before their forward or backward pass begins, execute standard PyTorch module forward and backward hooks, and immediately release the unsharded memory back to the CUDA caching allocator.

Prefetching and Stream Overlapping

To mitigate the 1.5x communication overhead of full parameter sharding, FSDP uses dedicated CUDA streams to overlap network communication with GPU computation:

• Forward Prefetch: While GPU Stream A executes compute on Layer L, GPU Stream B issues an asynchronous NCCL All-Gather for Layer L+1. By the time computation on Layer L finishes, Layer L+1 weights are already present in local device memory.
• Backward Prefetch: During the reverse pass, while compute runs on Layer L, Stream B simultaneously initiates the All-Gather for Layer L-1 while Stream C handles the asynchronous Reduce-Scatter of gradients for Layer L.

When compute time per layer exceeds network transmission time, the communication overhead is effectively hidden behind computation.

Hybrid Sharding and Cluster Scale Realities

In modern AI compute clusters, network bandwidth is non-uniform:

• Intra-Node Interconnect: GPUs within the same server communicate over high-speed links like NVLink (up to 900 GB/s bidirectional bandwidth per GPU on H100 systems).
• Inter-Node Interconnect: Communication across servers travels over InfiniBand or RoCE (Remote Direct Memory Access over Converged Ethernet), typically limited to 400 Gbps to 800 Gbps (50 to 100 GB/s per node).

Running full ZeRO-3/FSDP sharding across hundreds of nodes can bottleneck training on cross-node interconnect bandwidth during frequent All-Gather and Reduce-Scatter passes.

Hybrid Sharded Data Parallel (HSDP)

To resolve this imbalance, Hybrid Sharding (HSDP in PyTorch, ZeRO-Stage 3 with dp_size sub-groups in DeepSpeed) combines intra-node sharding with inter-node replication:

1. Intra-Node (NVLink Domain): Parameters, gradients, and optimizer states are fully sharded across the 8 GPUs inside each physical server node, maximizing NVLink utilization and eliminating intra-node memory redundancy.
2. Inter-Node (InfiniBand Domain): Nodes replicate the 8-GPU sharded model across the network using standard DDP-style All-Reduce on gradients, avoiding cross-node parameter All-Gathers.

This two-tier architecture keeps high-frequency parameter reconstruction confined to high-speed NVLink domains while maintaining high aggregate cluster throughput.

Practical Implementation Considerations

When deploying FSDP and ZeRO in production training pipelines:

• Activation Checkpointing: Combining FSDP with activation checkpointing (gradient checkpointing) reduces forward activation memory by discarding intermediate activations and recomputing them during backpropagation. This frees memory for larger per-device micro-batch sizes.
• Mixed Precision Configuration: Setting parameters to BF16, forward computation to BF16, and optimizer master weights to FP32 ensures numerical stability without the dynamic loss scaling overhead required by standard FP16.
• Distributed Checkpointing (DCP): Saving unsharded weights from hundreds of ranks can cause severe I/O bottlenecks. Modern frameworks use asynchronous sharded checkpointing, where each rank writes its local parameter shard to distributed storage (such as Lustre or S3), and consolidation is performed offline.

Sources

• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari et al., 2019)
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (Zhao et al., VLDB 2023)
• PyTorch Distributed: Experiences on Accelerating Data Parallel Training (Li et al., 2020)
• DeepSpeed ZeRO Documentation (Microsoft)
• PyTorch Fully Sharded Data Parallel (FSDP) Official Guide