TL;DR: For a new PyTorch-native codebase, start by evaluating FSDP2 alongside the model's required parallelism and checkpoint path. FSDP2 and DeepSpeed ZeRO-3 both support full state sharding, but buffer lifetimes, dtypes, offload and integration differ. Keep a working DeepSpeed stack unless a measured requirement justifies migration; CPU offload alone is not an exclusive DeepSpeed feature.
How to approach it
First establish model support, optimizer state, device mesh and restart requirements. Trace a sharded group's lifetime in each candidate, then compare complete updates and peak memory on equal work. Include a real save-resume test before choosing.
A strong answer
FSDP (Fully Sharded Data Parallel) is PyTorch's built-in way to shard parameters, gradients and optimizer state across data-parallel ranks; FSDP2 is its current API. The estimated stored state fits in memory, but training runs out of memory during its first parameter all-gather. The missing term is the full parameter group needed for execution. ZeRO and FSDP explains state ownership; this choice requires understanding when each implementation materializes those groups.
PyTorch 2.14 FSDP2, checked September 2026, represents sharded parameters with DTensor (PyTorch's tensor type that records how an array is split across devices) and groups communication through the fully_shard application. Per-parameter representation does not mean one network collective per parameter. Its resharding, prefetch and mixed-precision policies affect both traffic and memory.
DeepSpeed's ZeRO documentation describes stage-3 partitioning and offload options. FSDP2 also documents a CPU offload policy for parameters, gradients and optimizer state. NVMe-specific capabilities, optimizer support and framework integrations require separate checks; “needs offload” is not enough to choose a library.
Use an illustrative matched recipe, not a claim about either library's default: 30B parameters, 256 fully sharding ranks, 16 bytes of static training state per parameter, BF16 weight and gradient communication, one forward and one backward materialization with resharding between them.
static state/rank = 30e9 × 16 / 256 = 1.875 GB
full communicated weight payload W = 30e9 × 2 = 60 GB
ring-equivalent sends/rank for two gathers and one reduce-scatter:
3 × (255/256) × 60 = 179.296875 GB
at an assumed achieved 50 GB/s: payload service time ≈ 3.586 s
sanity: this is all calls combined, before considering overlap or startup.
Keeping parameters after forward, changing reduction dtype or accumulating gradients differently changes the count. Quantized collectives introduce their own representation and numerical contract. The two products therefore need not send equal bytes merely because both support stage-3-style sharding.
Suppose one gathered parameter group occupies 4 GB. Keeping that full group alongside 1.875 GB of stored shards already uses more than 5 GB, before counting activations. Prefetching a second group can raise the live set further. Use the allocation trace instead of inferring peak memory from state divided by ranks.
| Requirement | Acceptance evidence |
|---|---|
| Parallel layout | Correct gradients for the required layout |
| Performance | Matched complete-update time, including tails |
| Memory | Peak during gather, update and checkpoint |
| Resumption | Save, restart and verify intended state |
| Operations | Supported configuration and diagnosable failures |
Begin with FSDP2 when its PyTorch integration covers the new codebase's needs. Reverse that default if a required DeepSpeed feature wins the actual acceptance test or existing operational knowledge makes migration more costly than the measured benefit. No universal five-percent performance tie is established by the shared sharding idea.
Inspect p99 update duration and per-group arrival/completion in the profiler. A short total communication estimate does not prove it hides: a late gather can block the next layer even when total compute time is longer.
What interviewers probe next
- “Does full sharding make peak state exactly 16N/G?” No. That is one static recipe; full groups, activations and workspaces are extra.
- “Does a device mesh guarantee every combination works?” No. Test the specific model, operators and installed library versions.
- “Can a checkpoint load at another world size?” Verify the supported format and resharding path with the state you actually save.
Common mistakes
- Calling offload exclusive to one implementation.
- Assuming identical communication regardless of dtypes and resharding policy.
- Migrating a working training stack without testing recovery.
Key takeaways
- The toy static state is 1.875 GB per rank, far below possible execution peaks.
- The declared communication schedule sends about 179.3 GB per rank.
- Integration, measured updates and resumption decide the tool choice.