FSDP or DeepSpeed ZeRO-3: which would you pick for a new training codebase today, and why?

TL;DR: Both are ZeRO stage 3: shard weights, gradients and optimizer state across data-parallel ranks, all-gather weights per layer, reduce-scatter gradients. The choice is about surfaces, not math. FSDP2 shards each parameter on dimension 0 as a DTensor, which composes with tensor parallelism, context parallelism, torch.compile and distributed checkpointing without wrappers; it is the default for a new PyTorch codebase. DeepSpeed still wins when you need ZeRO-Offload or ZeRO-Infinity to train on GPUs too small for the sharded state, or when an existing config-driven stack already runs on it.

How to approach it

Neutralize the false framing first: they implement the same algorithm, and any memory or throughput argument at the level of "ZeRO-3 vs FSDP" is really about buffer management and prefetch, which are tunable in both. Then compare the surfaces: how a model is wrapped, how parameters are represented, what else composes with each, and how checkpoints are written. Give the default recommendation and the two conditions that override it. If the interviewer works at a company with an existing stack, ask which one it is before recommending a migration.

A strong answer

A typical situation: a team is starting a 30B pretraining codebase on 256 H100s and has engineers who have used both. The static training state is 30e9 × 16 B = 480 GB, so full sharding is required either way; at n = 256 that is 1.9 GB per rank plus activations, and either library handles it. The decision is about what the codebase will need next year.

What FSDP2 changed. FSDP1 flattened every parameter in a wrapped module into one contiguous FlatParameter and sharded that buffer. It worked, but the flat parameter hid the individual tensors, so anything that wanted to look at a parameter (a custom optimizer, a per-layer learning rate, a quantization hook, a tensor-parallel sharding) had to work around it. FSDP2 (fully_shard in torch.distributed.fsdp) shards each parameter on its own dimension 0 and represents the shard as a DTensor with a device mesh and a placement. The consequences are the ones the choice turns on:

• Tensor parallelism and context parallelism compose by adding mesh dimensions, so a 2D or 3D layout is one DeviceMesh with FSDP on one axis and TP on another, no second library.
• torch.compile traces through it, because there is no runtime flattening to break the graph.
• Distributed checkpointing writes DTensors natively and reshards on load to a different world size.
• Mixed precision is per parameter group, so an fp8 or bf16 policy can differ across layers.
• Memory is slightly higher than FSDP1 at the same settings, because per-parameter all-gathers are less contiguous; in practice it is inside noise once prefetch is on.

What DeepSpeed offers. ZeRO-3 with a JSON config that turns on sharding, prefetch bucket sizes, activation checkpointing, and communication settings without touching the model code. ZeRO++ adds quantized weight all-gathers (int8 or fp8 on the wire) and hierarchical partitioning that keeps a secondary copy of weights inside each node, so the cross-node all-gather traffic drops by up to 4×. ZeRO-Offload moves optimizer state to host memory, and ZeRO-Infinity moves it to NVMe, which lets a 70B fine-tune run on a single node at a fraction of the throughput. DeepSpeed-MoE and its inference engine are separate tools that share the config surface.

The per-step traffic is the same in both by construction:

ZeRO-3 / FSDP traffic per rank per step, model of N params in bf16 (2 B):
  forward  all-gather of weights:      (n−1)/n × 2N bytes
  backward all-gather of weights:      (n−1)/n × 2N bytes
  backward reduce-scatter of grads:    (n−1)/n × 2N bytes
  total ≈ 3 × 2N = 6N bytes per rank (for large n)

for N = 30e9: 6 × 30e9 = 180 GB per rank per step
  at 50 GB/s per NIC: 3.6 s
compute per rank per step at 16k tokens per rank:
  6 × 30e9 × 16,384 ÷ (989e12 × 0.4) ≈ 7.4 s

sanity: communication is half the compute time, so prefetch has to hide it; that is the same
        problem in either library, and both hide it with one-layer-ahead all-gathers.

ZeRO++'s hierarchical partitioning is the one implementation feature that changes this number: with a node-local secondary copy, the cross-node all-gathers are replaced by intra-node ones over NVLink and the NIC only carries the reduce-scatter. FSDP2's equivalent is hybrid sharding (HSDP), which shards within a group and replicates across groups; it reduces cross-group traffic to an all-reduce of gradients at the cost of replicating weights per group, which is a different trade with the same intent.

Decision: FSDP2 for a new PyTorch codebase, because a pretraining stack acquires tensor parallelism, context parallelism, compile and distributed checkpointing within a year, and FSDP2 gives all of them on one device mesh. The conditions that reverse it: the run needs offload to fit on hardware that cannot hold the sharded state even at maximum n, or the organization already operates a DeepSpeed or Megatron-DeepSpeed stack with the operational knowledge that comes with it, in which case the migration cost is real and the technical gap is small.

The reversal condition: an existing DeepSpeed codebase with working pipeline and MoE support and no appetite for a migration. Both implement the same algorithm, so the choice is operational, and ZeRO and FSDP covers what they share. TORCH_LOGS on the first run is where a misconfigured shard shows up. Model Memory Footprint is where the per-rank arithmetic comes from.

What interviewers probe next

• "Which one is faster?" Neither, at equal tuning; measured differences of 5 to 10% in either direction come from prefetch and bucket settings, and a candidate who quotes a benchmark without saying the settings has not run one.
• "How do you get tensor parallelism with DeepSpeed?" Megatron-DeepSpeed, which pairs Megatron's TP and PP with ZeRO for the data-parallel axis; the coupling is at the code level rather than the mesh level.
• "What does hybrid sharding buy?" At n = 512, full sharding all-gathers over 511 remote ranks per layer; HSDP with groups of 8 all-gathers over NVLink and all-reduces gradients across the 64 groups, which moves the heavy traffic onto the fast link at the cost of 8× replicated weights.
• "How does FSDP2 handle a parameter that must not be sharded?" fully_shard accepts an ignored_params set, and a DTensor placement of Replicate() on the FSDP mesh dimension does the same at the tensor level.

Common mistakes

• Arguing memory: "ZeRO-3 uses less memory than FSDP". The static state is 16 B/param divided by n in both; differences are in transient buffers and are settings.
• Not knowing FSDP2 exists and describing FSDP1's flat parameters and auto_wrap_policy as the current design.
• Recommending a migration off a working DeepSpeed stack for a fine-tuning team that will never add tensor parallelism.
• Forgetting offload as DeepSpeed's remaining differentiator.

Key takeaways

• Same algorithm: shard all three states, all-gather weights per layer twice per step, reduce-scatter gradients once; about 6N bytes per rank per step.
• FSDP2 shards per parameter as DTensors on a device mesh, which is why TP, CP, compile and checkpoint resharding compose with it.
• DeepSpeed keeps ZeRO++ quantized collectives and Offload/Infinity as its distinct features.
• New PyTorch codebase: FSDP2. Existing DeepSpeed stack or a need for offload: stay.