FSDP › Sharding Strategies

PyTorch FSDP exposes several sharding strategies through the ShardingStrategy enum. Each strategy offers a different trade-off between memory savings and communication overhead. Choosing the right strategy requires understanding both the memory equations and the communication patterns, because on bandwidth-limited hardware the communication cost can exceed the memory savings benefit.

FULL_SHARD is the most aggressive strategy. It shards parameters, gradients, and optimizer states across all GPUs. This corresponds to ZeRO Stage 3 and delivers the maximum possible memory reduction per GPU.

The communication cost is also highest with this strategy, because all-gather operations must run before every forward and backward pass through each FSDP unit. For a model split into $M$ FSDP units, one complete training step requires:

• $M$ all-gather operations during forward pass
• $M$ all-gather operations during backward pass
• $M$ reduce-scatter operations during backward pass

The total communication volume per step is proportional to $3MP'$ where $P'$ is the average parameters per FSDP unit. Since $MP' = P$ (total parameters), this simplifies to $3P \times \text{bytes\_per\_param}$ in aggregated data moved. Compare this to DDP, which requires one all-reduce of $2P \times \text{bytes\_per\_param}$ per step (factor of 2 comes from the ring-allreduce algorithm needing two passes). FSDP's communication volume is roughly 1.5 times larger than DDP per step, which is the price for the memory savings.

Use FULL_SHARD when:

• The model does not fit in GPU memory without sharding
• You have fast interconnects (NVLink, InfiniBand) that can absorb the communication overhead
• You are training at scale and memory efficiency is the primary constraint

SHARD_GRAD_OP shards gradients and optimizer states but keeps full parameter copies on each GPU during the forward and backward passes. This corresponds to ZeRO Stage 2.

Since parameters are not sharded, no all-gather is needed before each layer's computation. The communication happens only at the end of backward, where a reduce-scatter replaces the all-reduce of standard data parallelism. This reduces communication frequency significantly compared to FULL_SHARD: instead of $2M$ all-gathers plus $M$ reduce-scatters, you get just one reduce-scatter per step.

The memory savings are smaller. Parameters ($4P$ bytes) are not sharded, so each GPU still holds the full parameter tensor. Only gradients ($4P$ bytes) and optimizer states ($8P$ bytes) are split by $N$, saving:

$$4P + 8P - (4P/N + 8P/N) = 12P(1 - 1/N)$$

For large $N$, this approaches $12P$ bytes savings, leaving only the $4P$ parameter bytes per GPU. Compare to FULL_SHARD, which brings per-GPU memory to $(16P/N + 4P/M)$ bytes. The gap between the two strategies narrows as $N$ increases: on 64 GPUs, SHARD_GRAD_OP saves $12P \times 63/64 \approx 11.8P$ bytes, while FULL_SHARD saves nearly $15.75P$ bytes.

Use SHARD_GRAD_OP when:

• The model fits in GPU memory during the forward and backward passes (just the weights)
• You want to reduce optimizer memory usage without the communication overhead of full sharding
• Training throughput is more important than maximum memory reduction

NO_SHARD disables all sharding. FSDP with this strategy behaves exactly like DistributedDataParallel (DDP): each GPU holds the full model, and an all-reduce synchronizes gradients after each backward pass.

This strategy exists primarily so that code written to use the FSDP API can be run without any actual sharding, for debugging or profiling, without changing the training code. It also is a useful baseline for measuring the overhead that sharding introduces, since you can compare throughput with NO_SHARD against SHARD_GRAD_OP or FULL_SHARD on the same hardware with the same code.

HYBRID_SHARD is designed for multi-node training where intra-node bandwidth (via NVLink) is much higher than inter-node bandwidth (via InfiniBand or Ethernet). This asymmetry is extremely common in real clusters: eight GPUs on the same node may be connected by NVLink at 600 GB/s bidirectional, while the nodes connect to each other at 200 Gb/s InfiniBand (roughly 25 GB/s, 24 times slower).

With HYBRID_SHARD, FSDP applies full sharding within each node (using the fast intra-node interconnect) but maintains full parameter replicas across nodes. Gradient synchronization happens between nodes at the end of each backward pass, similar to DDP, but only synchronizing the gradients that each node already has aggregated.

The key benefit is that the expensive all-gather communication for parameter reconstruction stays within the node, using the fast NVLink. Cross-node communication is limited to one gradient all-reduce per step, which is the same communication pattern as standard DDP. This avoids sending all-gather communication across the slower inter-node network, which would be the bottleneck with FULL_SHARD at scale.

The cost is that memory savings are limited to within each node: if each node has 8 GPUs, parameters are sharded across those 8 GPUs but replicated across nodes. For a 70B model on a 4-node cluster of 8 GPUs each (32 GPUs total), HYBRID_SHARD gives you the memory benefit of 8-GPU sharding (saving $7/8$ of parameter memory within each node), while FULL_SHARD would give 32-GPU sharding savings ($31/32$).

Use HYBRID_SHARD when:

• You are training on multiple nodes
• Intra-node bandwidth (NVLink) significantly exceeds inter-node bandwidth
• The model fits in memory when sharded across GPUs within a single node

_HYBRID_SHARD_ZERO2 combines the intra-node sharding of HYBRID_SHARD with ZeRO Stage 2 semantics: parameters are kept full within each node during computation, but gradients and optimizer states are sharded intra-node while being replicated across nodes.

The leading underscore indicates this is a less commonly used variant. It is useful when you want the communication profile of SHARD_GRAD_OP within each node combined with the cross-node replication of HYBRID_SHARD. Concretely, this means no all-gathers for parameter reconstruction during forward and backward, which is beneficial when the bottleneck is communication latency rather than memory.