Memory Sharding in DeepSpeed ZeRO & PyTorch FSDP

Introduction

Welcome back. In our previous lesson, we performed a detailed quantitative analysis of the three stages of the Zero Redundancy Optimizer (ZeRO). We established that ZeRO-1 and ZeRO-2 offer significant memory savings with no communication overhead, while ZeRO-3 achieves a massive, linear reduction in memory at the cost of a 1.5x increase in communication. This covered the what and why of memory partitioning.

Today, we transition from theory to practice and explore the how. Our goal is to understand the implementation of the most powerful sharding strategy—full model state partitioning—as realized in two leading frameworks: Microsoft's DeepSpeed (which introduced ZeRO) and PyTorch's native Fully Sharded Data Parallel (FSDP).

Understanding these implementations is crucial for designing and debugging large-scale training systems. You'll learn about the core algorithm they share and the specific design choices and optimizations that differentiate them, enabling you to make informed decisions when architecting AI systems.

The Core Mechanism: On-Demand Parameter Materialization

In the previous lesson, we saw that ZeRO Stages 1 and 2 still require each GPU to hold a full copy of the model parameters, which becomes the memory bottleneck. ZeRO-3 and FSDP solve this by partitioning the parameters themselves. This raises a fundamental question: how can a GPU perform a forward pass for a layer if it doesn't have the complete weights for that layer?

The answer is a clever, dynamic process of on-demand parameter materialization. Instead of holding the entire model in memory, each GPU retrieves the parameters it needs just-in-time for computation and discards them immediately afterward.

Let's watch a short segment that introduces this concept and contrasts it with the limitations of standard data parallelism.

Now, let's look at the detailed mechanics of how this on-demand materialization works. The process can be broken down into the forward and backward passes.

The diagram below illustrates this workflow. Each GPU (or "rank") holds only a shard of the model's parameters.

Here is the step-by-step process, which applies to both DeepSpeed ZeRO-3 and FSDP:

During the Forward Pass (for each block of layers):

1. Gather Parameters: Right before computation, all GPUs participate in an AllGather collective operation. Each GPU broadcasts its local shard of the parameters for the current layer(s) to all other GPUs. At the end of this step, every GPU momentarily has the full, unsharded parameters for that specific layer block.
2. Perform Computation: With the full parameters now materialized in memory, the GPU executes the standard forward pass for that block.
3. Discard Parameters: Immediately after the computation, each GPU discards the parameter shards it received from its peers, freeing that memory. It only retains its original, assigned shard.

During the Backward Pass (for each block of layers):

1. Gather Parameters: Just like in the forward pass, an AllGather operation re-materializes the full parameters for the current layer block, as they are needed to compute the gradients.
2. Perform Computation: The backward pass computation is executed for that block, producing local gradients.
3. Reduce and Scatter Gradients: Instead of synchronizing the full gradient tensor with AllReduce (as in standard DDP), the frameworks use a ReduceScatter collective. This single, efficient operation achieves two things: it averages the gradients from all GPUs and shards the result, delivering to each GPU only the partition of the averaged gradient that corresponds to its parameter partition.
4. Discard Parameters: The gathered parameters are discarded.
5. Update Weights (Locally): Each GPU now has everything it needs to update its local partition of the model: its shard of the averaged gradients and its shard of the optimizer states. It performs the optimizer step on its local data only.

To see this dynamic process explained in detail, let's watch another excellent video.

This on-demand materialization is the fundamental technique that allows data-parallel training of models that are far too large to fit on a single GPU.

Framework Implementations: DeepSpeed vs. PyTorch FSDP

While both DeepSpeed and FSDP use the same core mechanism, they are distinct libraries with different origins and some unique features.

DeepSpeed ZeRO

DeepSpeed is a comprehensive library from Microsoft for large-scale model training. It introduced the ZeRO optimizer and its three stages. As you recall from the last lesson:

• ZeRO-1: Shards optimizer states.
• ZeRO-2: Shards optimizer states and gradients.
• ZeRO-3: Shards optimizer states, gradients, and model parameters.

When we talk about full memory sharding, we are referring to ZeRO-3. In practice, using DeepSpeed often involves simply adding a configuration file to your training script that specifies which ZeRO stage to use.

The key takeaway is that DeepSpeed provides a mature, battle-tested implementation of the ZeRO stages, accessible as configuration options.

PyTorch FSDP

Fully Sharded Data Parallel (FSDP) is PyTorch's native implementation of the ZeRO-3 paradigm. It was co-designed with PyTorch's core components, such as the autograd engine and memory allocator, for tight integration and high performance. FSDP introduces a few key concepts that are important to understand from a system design perspective.

Let's summarize the key FSDP concepts from the paper:

1. Wrapping Policy and FSDP Units: FSDP requires you to "wrap" your model's layers. This wrapping partitions the model into FSDP Units. The AllGather and ReduceScatter operations work on a per-unit basis. The size of these units creates a trade-off:

   • Fine-grained units (e.g., wrapping each layer individually) result in lower peak memory, as only a small part of the model is materialized at once. However, this can lead to higher communication overhead due to many small, less efficient AllGather calls.
   • Coarse-grained units (e.g., wrapping large blocks of layers) have higher peak memory but can be more efficient by issuing fewer, larger communication calls.

2. FlatParameter: To optimize communication, FSDP coalesces all parameters within a single FSDP unit into a single, contiguous 1D tensor called a FlatParameter. This ensures that communication primitives like AllGather and ReduceScatter operate on a single, large chunk of data, which is much more efficient for the underlying networking hardware (e.g., NCCL) than issuing many small calls.

3. Sharding Strategies: FSDP offers two primary strategies that are vital for system architects:

   • Full Sharding (FULL_SHARD): This is the direct equivalent of ZeRO-3. All model states are sharded across all available GPUs in the world. This provides the maximum memory savings but sends all communication traffic across the entire network.
   • Hybrid Sharding (HYBRID_SHARD): This is a more sophisticated strategy designed for multi-node clusters. It allows you to define subgroups of GPUs for sharding. A common use case is to shard within a node (e.g., across 8 GPUs connected by high-speed NVLink) and replicate across nodes (connected by slower InfiniBand or Ethernet). This minimizes expensive inter-node communication for the AllGather operations, restricting it to the faster intra-node interconnect. Gradient synchronization still happens across all nodes, but the volume of parameter communication is drastically reduced. This is the same concept as "hierarchical partitioning" in DeepSpeed's ZeRO++.

To solidify the distinction between these frameworks, the following video provides a concise comparison.

In modern practice, high-level libraries like Hugging Face Accelerate provide a unified interface, allowing you to switch between DeepSpeed and FSDP often with just a change in a configuration file. This abstraction makes it easier to experiment with these powerful techniques without deep code changes.

Conclusion

In this lesson, we demystified how full memory sharding is implemented in practice. We moved beyond the quantitative benefits of ZeRO-3 to the underlying mechanisms that enable it.

Key Takeaways:

• Both DeepSpeed ZeRO-3 and PyTorch FSDP implement full model state partitioning to enable data-parallel training of models that exceed single-GPU memory.
• The core mechanism is on-demand parameter materialization: parameters for a given layer are gathered via AllGather just before computation and discarded immediately after.
• Gradients are efficiently synchronized and sharded using a ReduceScatter operation, avoiding the need to store the full gradient tensor on every device.
• FSDP introduces explicit concepts like FlatParameter for communication efficiency and a HYBRID_SHARD strategy, which is critical for optimizing performance on multi-node clusters by aligning the communication pattern with the network topology.
• While FSDP is equivalent to DeepSpeed Stage 3 in its sharding goal, its native integration and hybrid sharding model are key design features within the PyTorch ecosystem.

Preview of the Next Lesson:

So far, our memory optimization efforts have focused entirely on the model states: parameters, gradients, and optimizer states. However, another significant portion of GPU memory is consumed by activations—the intermediate outputs of each layer stored during the forward pass for use in the backward pass. In the next lesson, we will explore activation checkpointing, a powerful technique that trades a small amount of re-computation for substantial memory savings by not storing all activations.