ZeRO Optimization › ZeRO in the PyTorch Ecosystem: FSDP

ZeRO's design philosophy was sufficiently influential that PyTorch incorporated a native implementation: Fully Sharded Data Parallel (FSDP), which implements ZeRO Stage 3 semantics directly in the PyTorch distributed training API. FSDP became part of PyTorch in version 1.11 and was substantially improved in subsequent releases.

FSDP's approach is slightly different from DeepSpeed's ZeRO Stage 3 in implementation, but the core idea is the same: each parameter is permanently owned by a subset of processes, and parameters are gathered on demand before computation. FSDP adds some features that make it easier to use within PyTorch's existing model definition patterns:

• Auto-wrapping: FSDP can automatically identify which submodules to shard based on a minimum parameter count threshold, without requiring manual annotation.
• Mixed sharding strategies: Different submodules can use different sharding strategies (full shard, hybrid shard, or no shard), allowing for hierarchical sharding that maps to the within-node/between-node topology.
• Activation checkpointing integration: FSDP integrates with PyTorch's native gradient checkpointing via checkpoint_wrapper, making the combination straightforward.

FSDP is the standard tool for large-model training in PyTorch-native workflows, while DeepSpeed ZeRO remains dominant in environments that use the DeepSpeed infrastructure or require the ZeRO-Offload and ZeRO-Infinity extensions. The two systems have largely converged on similar feature sets, and the choice between them often comes down to ecosystem fit rather than fundamental capability differences.