Distributed Training: FSDP & ZeRO

The SFT pipeline chapter ended with a trustworthy single-device training recipe. This chapter handles the next limit: full-weight updates may be correct for the task, but one GPU can't hold the states and activations needed to execute them. Fully Sharded Data Parallel (FSDP) and ZeRO (Zero Redundancy Optimizer) reduce repeated model-state storage across data-parallel ranks, then make communication part of the training cost.

Start with the memory contract. One useful sizing convention, used in the ZeRO analysis, assumes low-precision parameters and gradients plus an FP32 master parameter copy and FP32 Adam moments.^[1] Under that specific recipe, each parameter costs 16 bytes of model-state memory before activations, temporary buffers, or allocator overhead:

• 2 bytes for the BF16 (bfloat16, short for Brain Floating Point) / FP16 (16-bit half precision) weight
• 2 bytes for the BF16 / FP16 gradient
• 4 bytes for the FP32 (32-bit single precision) master copy of the weight
• 4 bytes for Adam's first moment (momentum)
• 4 bytes for Adam's second moment (variance)

For a 1B model, that is 16 GB of model states alone. Change the optimizer, parameter dtype, master-weight policy, or offload policy and the byte count changes. Always write down the recipe before using a bytes-per-parameter number.

Scale this recipe to 7B parameters and it reaches 112 GB of model states. At 70B parameters, it exceeds 1 TB before activations. If you need full-parameter training at that scale, replicated data parallelism alone won't make it fit.

FSDP (Fully Sharded Data Parallel) and ZeRO (Zero Redundancy Optimizer, a DeepSpeed technique that shards model states across GPUs) don't make memory disappear. They remove redundant copies of model states, then pay the price in communication.

The memory wall in numbers

For a 70B parameter model under the FP16/BF16 parameter-and-gradient plus FP32-master-and-Adam accounting convention, the model-state math looks like this:

• Weights: $70 \times 10^9 \times 2 \text{ bytes} \approx$ 140 GB
• Gradients: $70 \times 10^9 \times 2 \text{ bytes} \approx$ 140 GB
• Adam state + FP32 master weights: $70 \times 10^9 \times (4 + 4 + 4) \text{ bytes} \approx$ 840 GB

That is about 1.12 TB of model states before activations, temporary buffers, or allocator overhead. It is an assumption-bound calculation, not a promise about every optimizer implementation.

The arithmetic is small enough to turn into a quick check. This example uses decimal GB to match the table below.

Why DDP hits a limit

Before we look at the fix, recall how ordinary Distributed Data Parallel (DDP) works. DDP places a full copy of the model on every GPU, splits the training batch across them, and synchronizes gradients with an all-reduce (a collective operation that sums values across all GPUs and returns the result to every GPU) after each backward pass. It works well for models that fit on one GPU, because adding more GPUs improves throughput without changing the per-GPU memory requirement.

The problem appears when the model itself is too large for one GPU. DDP replicates the full model states on every rank (process/GPU), so adding more GPUs improves throughput but doesn't reduce the per-GPU memory footprint. FSDP and ZeRO attack that specific problem by sharding model states across the data-parallel group.

Think of DDP like every warehouse node keeping a full copy of the same giant catalog. Coordination is easy, but the storage is wasteful. Every GPU keeps a complete copy of the parameters, gradients, and optimizer states.

ZeRO changes that trade-off. Instead of duplicating the full catalog on every rank, it partitions the catalog across ranks and reconstructs only the pieces needed for computation.^[1] In logistics terms, ZeRO is like a distributed warehouse network: each regional hub stores only a fraction of the catalog and pulls the remaining items on demand when an order arrives.

• DDP: Replicate everything. Simpler communication pattern, much higher memory use.
• ZeRO: Shard model states. Lower memory use, more communication.

ZeRO stages: removing redundancy one layer at a time

ZeRO (Zero Redundancy Optimizer) is easiest to understand as an incremental evolution. Each stage removes one more redundant copy of the model states.

ZeRO stage 1: shard optimizer states

Each GPU keeps optimizer state for only $1/N$ of the parameters. Parameters and gradients are still replicated.

• Memory savings: Optimizer state drops from 840 GB total to roughly $840/N$ GB per GPU.
• Communication: Similar aggregate gradient-synchronization volume to the replicated baseline; optimizer ownership changes.

ZeRO stage 2: shard optimizer + gradients

Stage 2 also shards gradients, so each GPU owns only the gradient shard that matches its optimizer shard.

• Memory savings: Optimizer states and gradients both scale down by roughly $1/N$.
• Communication: Gradient sync is typically implemented with reduce-scatter (a collective that sums values across GPUs and returns each GPU's shard) instead of all-reduce. The aggregate volume stays in the same ballpark as DDP, but the result is already partitioned.

ZeRO stage 3: shard everything

Stage 3 shards parameters as well. No rank owns a full persistent copy of the model states anymore.

• Memory savings: Parameters, gradients, and optimizer states are all partitioned. For the 70B example on 256 GPUs, the model-state footprint falls to about 4.4 GB per GPU before activations.
• Communication: Parameters now have to be materialized on demand for computation.
  1. Forward: all-gather (collect shards from all GPUs to rebuild the full tensor) the current layer or FSDP unit, run compute, then reshard.
  2. Backward: all-gather that unit again, compute gradients, then reduce-scatter gradient shards.
• Trade-off: Memory scales almost linearly with the data-parallel group size, but communication becomes a first-class bottleneck.

Key insight: ZeRO doesn't reduce the aggregate model-state bytes; it eliminates persistent redundant copies. A 70B model still needs roughly 1.12 TB of model-state memory in aggregate, but ZeRO-3 partitions its steady-state storage across GPUs. Temporary materialized units and communication buffers still add to each rank's peak. The trade-off is communication: you pay in network bandwidth for what you save in persistent per-GPU state.

Memory per GPU (70B model, 256 GPUs)

The table below assumes BF16/FP16 weights and gradients with Adam states in FP32. The totals are for model states only. Activation memory is listed separately because it depends heavily on sequence length, hidden size, micro-batch size, checkpointing, and attention implementation.

This is why ZeRO-3 or FULL_SHARD is often paired with activation checkpointing (recomputing activations during the backward pass to save memory, also called gradient checkpointing). Sharding fixes model-state memory, but activations can still dominate the actual training footprint.

The same table should be generated from a recipe rather than copied into a sizing document:

Communication overhead analysis

ZeRO-3 buys memory by turning parameter access into communication. That makes the interconnect part of your training system design, not an afterthought. Just as a distributed warehouse system only works if the trucking network between hubs can handle replenishment traffic, ZeRO-3's memory savings depend on whether your GPU interconnect can keep up with all-gather volume.

Communication primitives you need to know

Those names stop blurring once you track one tiny tensor through each collective. Read each panel left to right: same four ranks, different "after" state.

The operations become concrete on a two-rank gradient vector. All-reduce leaves the summed vector on both ranks; reduce-scatter leaves one summed slice per rank; all-gather reconstructs a parameter vector from shards.

The 1.5x heuristic

People often summarize ZeRO-3 communication as "about 1.5x DDP." That's a useful heuristic, but it's only a heuristic.

For a model with $P$ parameter bytes:

1. DDP does one gradient all-reduce per step, which is commonly modeled as about $2P$ bytes per rank.
2. ZeRO-3 / FULL_SHARD does roughly two parameter all-gathers plus one gradient reduce-scatter, which is commonly modeled as about $3P$ bytes per rank.

That gives the familiar $3P / 2P = 1.5x$ rule of thumb. The harder part in practice isn't only byte volume, but how that volume is scheduled. DDP tends to use a smaller number of large collectives. ZeRO-3 and FSDP break communication into many layer-level or unit-level operations. If those collectives are small and your network has high latency, throughput can fall off quickly.

ZeRO-3 is most attractive when memory is the hard constraint. If the model already fits and the network is weak, DDP or ZeRO-2 style sharding can be faster.

FSDP (Fully Sharded Data Parallel)

FSDP is PyTorch's native sharded data-parallel stack.^[2] The classic API is FullyShardedDataParallel (often called FSDP1). Current PyTorch exposes a DTensor-based API via torch.distributed.fsdp.fully_shard (often called FSDP2), and its tutorial marks FSDP1 deprecated while giving a migration guide. Treat FSDP1 as an existing-code path and start new designs from fully_shard.^[3][4]

The main reason people choose FSDP is operational simplicity. It stays close to PyTorch's native distributed stack, native profilers, and native checkpointing tools. PyTorch's general torch.compile guidance is to compile the highest-level train or evaluation step, or the top-level module. If a distributed wrapper such as FSDP causes issues, compile the inner module instead. If you're compiling FSDP1, current docs also require use_orig_params=True.^[5][6]

FSDP solves model-state memory. It doesn't solve activation memory. Long sequences and large micro-batches still push activations up, which is why FSDP is often paired with activation checkpointing.

How full sharding materializes one unit

For FULL_SHARD, or for an FSDP2 configuration that reshards after forward, one transformer-sized communication unit follows this lifecycle:

1. Group: Parameters are grouped at transformer-block granularity.
2. Shard: Each unit's parameters are sharded across the process group.
3. Execution:
   • Before forward pass of a unit: All-gather full parameters.
   • After forward pass: Discard full parameters (keep only shard).
   • Before backward pass: All-gather full parameters again.
   • After backward pass: Reduce-scatter gradients and update the sharded optimizer state.

This sequence diagram shows one wrapped unit in a two-GPU FSDP job.

The following code shows a practical FSDP1 setup. You will still meet this API in existing codebases, so it is worth reading even though new designs should begin with fully_shard. The important details are:

1. Wrap at the transformer-block level instead of wrapping the whole model as one giant unit.
2. Create the optimizer after FSDP wraps the model.
3. Don't confuse backward_prefetch with limit_all_gathers. The first helps overlap communication with compute. The second is mainly a CPU-side rate limiter that caps in-flight all-gathers and peak memory.^[6]

This is distributed setup code, not a single-process notebook cell. It assumes torch.distributed.init_process_group() has already run under torchrun and that each rank has selected its local CUDA device.

FSDP2: the newer DTensor-based API

FSDP2 moves away from the flat-parameter wrapper model and toward per-parameter sharding with DTensor. The point isn't only elegance. It changes how composable and inspectable the system is.

For new code, the structural pattern is bottom-up application of fully_shard, then optimizer construction over the resulting DTensor parameters:^[3]

This snippet shows structure rather than a runnable local demo: fully_shard must execute inside a distributed process group on suitable accelerators. The learning rate remains an SFT baseline to evaluate, not a distributed-training requirement.

Key technical improvements

1. Cleaner composition: FSDP2 uses DeviceMesh, which makes it much easier to compose data-parallel sharding with tensor parallelism in the same topology.

2. Cleaner state dict story for large runs: Sharded state dicts are first-class, and distributed checkpointing avoids the old pattern of gathering everything onto rank 0 for saving.

3. Cleaner handling of frozen parameters: Per-parameter sharding is a better fit for PEFT and LoRA-style workflows where only a small subset of weights should update.

4. Different scheduling controls: fully_shard is no longer a wrapper configured through FSDP1's auto_wrap_policy, backward_prefetch, use_orig_params, or limit_all_gathers arguments. Apply it bottom-up, then use FSDP2 methods such as set_modules_to_forward_prefetch or set_modules_to_backward_prefetch when explicit prefetch control is needed.^[3]

The practical takeaway is simple: FSDP1 still shows up in existing codebases, while current PyTorch's FSDP2 documentation provides the migration path and fully_shard API for new work.

The most important FSDP choice is usually sharding granularity. Use transformer blocks as communication units. If one root unit owns the entire model, every all-gather becomes too large and peak memory shoots back up.

This simplified calculation isolates why granularity matters. It models only the transient full-parameter materialization, not activations or prefetch overlap:

FSDP1 sharding strategies reference

Checkpointing, resume, and fault recovery

Large training runs fail in the middle all the time. A node reboots. A preemptible instance disappears. A job gets rescheduled. If your checkpoint only contains model weights, you didn't save the training run. You saved an export artifact.^[7]

For a real resume point, the checkpoint usually needs at least:

In a sharded system, the safest default is to save sharded checkpoints, not to gather everything onto rank 0 first. FSDP2 and distributed checkpoint APIs are pushing in that direction because the "collect all weights on one machine and write a giant file" pattern eventually becomes the bottleneck or the failure point.^[3][7]

DeepSpeed goes one step further with Universal Checkpointing, which is meant to make checkpoint artifacts more portable across different parallelism layouts instead of tying restore logic to the exact original topology.^[8] That matters once you start changing world size between runs or resume on a different cluster shape.

The operational rule is simple:

1. Save checkpoints by training step and consumed tokens, not vague names like latest-final.
2. Test one real resume path early in the run, before trusting a week-long job to it.
3. Verify that resumed loss on the next batch is close to the non-resumed run.

If you can't answer "what exact state do we restore, and how do we know resume is faithful?", the training system isn't production-ready yet.

A distributed checkpoint manifest also has to record the sharding topology it was written under, even when a portable conversion path exists:

DeepSpeed ZeRO

DeepSpeed is Microsoft's distributed training runtime that introduced ZeRO and provides heterogeneous-memory offload through ZeRO-Offload and ZeRO-Infinity.^[1][9] It's got a steeper learning curve than native PyTorch, but it remains a strong choice when memory pressure is the real blocker or when the rest of your stack already depends on Megatron-DeepSpeed style training.

One important DeepSpeed-specific detail is that ZeRO-3 depends on consistent module execution order across ranks. Dynamic routing patterns such as MoE can deadlock parameter all-gathers if different ranks enter different submodules, which is why DeepSpeed exposes the idea of ZeRO-3 leaf modules.^[10]

The following distributed setup snippet shows how to configure and initialize a model using DeepSpeed. Run it through the DeepSpeed launcher in a real training job. The configuration dictionary defines the ZeRO stage, offloading settings, optimizer, and training batch parameters. The learning rate is an example SFT starting point to evaluate, and the bucket sizes are tuning knobs rather than canonical values:

ZeRO-Infinity: CPU and NVMe offloading

DeepSpeed extends ZeRO-3 with ZeRO-Infinity, which allows offloading sharded parameters and optimizer states to CPU RAM or even NVMe SSDs.^[9]

This makes some otherwise impossible training runs feasible, but it doesn't make offload free. Once parameters or optimizer state spill into CPU or NVMe, throughput becomes constrained by PCIe, host memory bandwidth, storage latency, and how well the runtime can overlap data movement with compute.

DeepSpeed's feature matrix is broader, but it isn't "all combinations are valid." Current docs say AutoTP training supports ZeRO stages 0, 1, and 2, while PipelineModule isn't compatible with ZeRO-2 or ZeRO-3.^[10][11]

Turn compatibility constraints into a config gate, not a surprise after scheduling a large job:

FSDP vs DeepSpeed comparison

Both frameworks solve the same core problem, but they optimize for different operational constraints. FSDP optimizes for staying close to standard PyTorch. DeepSpeed optimizes for a wider set of large-scale runtime features.

When to use which?

• Use FSDP when you want the PyTorch-native path, the model fits in aggregate GPU memory once sharded, and standard PyTorch profiling and checkpointing matter more than exotic runtime features.
• Use DeepSpeed when CPU or NVMe offload is non-negotiable, or when you're already in a DeepSpeed or Megatron-DeepSpeed stack whose TP/PP requirements fit the current stage-compatibility rules.

Beyond data parallelism: 3D parallelism

For very large jobs, data-parallel sharding alone isn't enough. As you push the data-parallel group wider, global batch size, optimizer dynamics, and cross-node communication all start fighting you. The standard answer is to combine multiple forms of parallelism in the same training topology.^[12]

That's what people mean by 3D parallelism: data parallelism for replica groups, tensor parallelism (splitting weight matrices within a layer across GPUs) inside layers, and pipeline parallelism (splitting sequential layer groups across different devices) across layer ranges.

The diagram below shows how data, tensor, and pipeline parallelism combine to distribute a model across a large cluster. Solid arrows show pipeline flow through layer ranges. Dotted arrows show synchronization between matching data-parallel replicas.

Tensor parallelism (TP)

Tensor parallelism splits the work inside a layer. Instead of putting a full weight matrix on one GPU, it partitions the matrix across GPUs.

• How it works: In a transformer MLP, the first projection is often column-sharded and the second projection row-sharded. Each rank computes a partial result, then the ranks synchronize those partials.
• Bandwidth requirement: This synchronization happens inside the layer, so TP strongly prefers very fast intra-node links such as NVLink or NVSwitch (high-bandwidth GPU interconnects).

Pipeline parallelism (PP)

Pipeline parallelism splits the model across depth by assigning different layer ranges to different devices or nodes.

• How it works: Stage 0 might own layers 1-12, stage 1 owns 13-24, and activations flow from one stage to the next.
• Main trade-off: You reduce memory pressure per rank and avoid TP-style per-layer collectives across the whole model, but you introduce pipeline bubbles and need micro-batching to keep stages busy.

The escalation path should be deliberate.^[13]

• Start with single-node or data-parallel sharding when the model mostly fits and you want the simplest debug story.
• Add tensor parallelism when single layers are too wide for one device or when matmul shards need to stay on NVLink-class links.
• Add pipeline parallelism when layer depth itself must be partitioned across devices.

A common placement keeps TP on fast intra-node links, uses PP to partition depth across broader topology boundaries, and applies data parallelism across replica groups. The right degrees still depend on available links, memory, model shape, and global-batch constraints.

In current PyTorch, a DeviceMesh is the abstraction used to express multiple device dimensions, which lets FSDP2 compose with tensor parallelism on the same hardware (often called FSDP + TP, or 2D parallelism). Megatron Core also documents context parallelism, which splits the sequence dimension for long-context workloads. Combining data, tensor, pipeline, and context parallelism creates a four-axis topology; runtimes don't always use the same "3D" or "4D" label for that composition.^[3][13]

Each topology degree multiplies into the required world size. Calculate it before reserving hardware:

Debugging & profiling distributed training

Distributed training fails in ways that single-GPU training doesn't. Jobs hang instead of crash. One bad rank can stall everybody else. Performance regressions often come from communication scheduling, not math kernels.

Common failure modes

1. NCCL timeouts or hangs: Often one rank died earlier or entered a different collective pattern. TORCH_NCCL_BLOCKING_WAIT=1 can turn a silent hang into a visible failure.
2. Startup OOM: Large models can OOM before training starts if you materialize full parameters too early. Lazy or meta-device initialization matters.
3. Wrong gradient handling: With an FSDP1 sharded strategy, use fsdp_model.clip_grad_norm_() because it computes across sharded gradients. FSDP2's DTensor-based parameters instead support torch.nn.utils.clip_grad_norm_(model.parameters(), ...) in the current tutorial.^[6][4]

Profiling tools

To optimize throughput (tokens/sec), you must identify if you're compute-bound or communication-bound.

• torch.profiler: Look at ncclAllGather, ncclReduceScatter, and gaps where GPUs wait for communication. For FSDP1, tune wrapping granularity and backward_prefetch; use limit_all_gathers=True for peak-memory rate limiting, not as the overlap knob. For FSDP2, use its module prefetch APIs when implicit scheduling is insufficient.^[6][3]
• DeepSpeed Flops Profiler or DeepSpeed runtime logs: Use them to break step time into compute vs communication and to see whether offload is the bottleneck.

For hard distributed bugs, capture one clean repro run with NCCL_DEBUG=INFO and TORCH_DISTRIBUTED_DEBUG=DETAIL instead of trying to reason from high-level symptoms alone.

Training observability signals that matter

Once the job runs, you still need to know why it is slow or unstable. Good training dashboards separate data, compute, communication, and memory instead of showing only one loss curve.^[13]

If you only watch training loss, a run can look healthy while throughput quietly collapses. Distributed training needs systems telemetry, not only ML telemetry.

Training-system bottleneck map

Strong distributed-training debugging connects the abstraction to the bottleneck. A useful answer names the split, the communication primitive, and the symptom you would measure.

A practical debugging loop is:

1. Reproduce on the smallest rank count that still fails.
2. Confirm every rank enters the same collective sequence.
3. Capture profiler traces and separate compute, communication, and idle time.
4. Check peak memory before and after forward, backward, and optimizer step.
5. Change one axis at a time: sharding granularity, micro-batch size, checkpointing, or placement.

This is the bridge between ML and distributed systems. The right answer is rarely "use FSDP" or "use DeepSpeed." The stronger answer is "model states don't fit, so I shard them; then I measure whether communication or activation memory became the new bottleneck."

A worked sizing exercise

A common sizing prompt is: "How many 80 GB GPUs do you need to train a 70B model with ZeRO-3?" Walk through it step by step.

Step 1: count the model states.

With mixed-precision Adam, each parameter carries:

• 2 bytes (BF16 weight)
• 2 bytes (BF16 gradient)
• 12 bytes (FP32 master weight + Adam moments)

Total: 16 bytes per parameter.

For 70B parameters: $70 \times 10^9 \times 16 \text{ bytes} = 1{,}120 \text{ GB}$ of model states.

Step 2: compute the model-state lower bound.

If you looked only at model states, the absolute lower bound on 80 GB GPUs would be:

$$\frac{1{,}120 \text{ GB}}{80 \text{ GB/GPU}} = 14$$

So 14 GPUs is the math-floor for model states alone.

The same calculation is worth making executable:

Step 3: add overhead.

The math above is for model states only. Real training also needs:

• Activations (depends on sequence length and batch size)
• Communication buffers for all-gather and reduce-scatter
• CUDA allocator overhead and fragmentation

There isn't one honest fixed headroom percentage: a context-length or micro-batch change can dominate it. Measure or estimate the remaining peak memory for the specific workload, then compare candidate topologies:

Step 4: answer the question.

The clean answer is: 14 GPUs is only the raw ZeRO-3 lower bound for model states under this 16-byte recipe. It isn't a runnable cluster recommendation. A practical count requires a measured or modeled activation/buffer peak and an acceptable throughput result. In the example above, an 18 GB non-state peak rules out 16 GPUs but fits at 24; a longer sequence could invalidate that answer again. Activation checkpointing, a smaller micro-batch, tensor parallelism, or pipeline parallelism may change the result.

Common pitfalls

Misidentifying which axis is split

Symptom: A team says they "added ZeRO-3 for tensor parallelism," but profiler traces still show whole-layer matmuls on each rank and the real change is only lower model-state memory.

Cause: ZeRO and FSDP shard model states across data-parallel ranks. They do not split the matrix multiply itself. Tensor parallelism splits layer computation, and pipeline parallelism splits model depth.

Fix: Check what moved in the trace or memory profile before naming the technique. If weights, gradients, and optimizer states shrank per rank, that is sharding. If one layer's matmul is split across devices, that is tensor parallelism.

Wrapping the whole model as one FSDP unit

Symptom: You enable ZeRO-3, but peak memory is still unexpectedly high and you get OOM during the forward pass.

Cause: You wrapped the entire model as a single FSDP unit instead of wrapping individual transformer layers. When the whole model is one unit, FSDP all-gathers every parameter at once. That defeats the purpose of layer-by-layer sharding.

Fix: Use auto_wrap_policy to wrap at the transformer-block level. The code example earlier shows how to do this with transformer_auto_wrap_policy.

Thinking FSDP1 limit_all_gathers=True creates overlap

Symptom: You set limit_all_gathers=True and expect faster training, but throughput doesn't improve.

Cause: limit_all_gathers is a CPU-side rate limiter. It caps how many all-gathers are in flight at once to control peak memory. It isn't an overlap or performance knob.

Fix: In FSDP1, use backward_prefetch and communication-unit granularity for overlap. In FSDP2, use its prefetch methods when needed; it doesn't expose FSDP1's limit_all_gathers argument.

Treating SHARD_GRAD_OP as identical to ZeRO-2

Symptom: You switch from ZeRO-2-style reasoning to FSDP SHARD_GRAD_OP, then wonder why peak memory is higher than expected even though gradients are sharded.

Cause: The names are similar, but the parameter lifetime differs. PyTorch's SHARD_GRAD_OP keeps parameters unsharded through forward and backward, then reshards after backward.^[6] That changes peak memory behavior even when the high-level mental model sounds like ZeRO-2.

Fix: Confirm the actual parameter lifetime in docs and profiler traces before sizing memory. If you need the lowest parameter footprint, compare against FULL_SHARD rather than assuming SHARD_GRAD_OP behaves like a DeepSpeed stage name.

Ignoring activation memory

Symptom: The model-state math says everything should fit, but you still OOM.

Cause: Activations can be larger than model states for long sequences. A single forward pass with a long context can generate gigabytes of intermediate tensors.

Fix: Pair ZeRO-3 or FSDP with activation checkpointing. Trade compute for memory by recomputing activations during the backward pass instead of storing them.

Assuming every DeepSpeed feature composes with ZeRO-3

Symptom: You configure AutoTP with ZeRO-3 and the job crashes or hangs.

Cause: DeepSpeed has documented compatibility limits. AutoTP currently supports ZeRO stages 0, 1, and 2, but not 3. PipelineModule isn't compatible with ZeRO-2 or ZeRO-3.

Fix: Check the stage-compatibility matrix before you design your training topology. Don't assume that because DeepSpeed supports both features, they work together.

What you should be able to defend

By the end of this lesson, you should be able to explain and operate these parts of a distributed training setup:

• How weights, gradients, FP32 master weights, and Adam moments create the memory wall.
• Why DDP improves throughput but doesn't reduce per-GPU model-state memory.
• What ZeRO-1, ZeRO-2, and ZeRO-3 shard.
• How FSDP1 FULL_SHARD, SHARD_GRAD_OP, NO_SHARD, and HYBRID_SHARD trade memory for communication, and why new code should start with FSDP2 fully_shard.
• Why FSDP2's fully_shard is different from the FSDP1 wrapper API.
• How all-reduce, reduce-scatter, and all-gather show up in training traces.
• When to use FSDP, DeepSpeed ZeRO, offload, tensor parallelism, and pipeline parallelism.
• How to size a 70B-style training job before accounting for activations and communication buffers.
• Which profiler and observability signals reveal communication, memory, checkpointing, and dataloader bottlenecks.

Final selection workflow

Use this sequence when sizing a real run:

1. Compute model-state memory first. If the bytes-per-parameter math already blows past one GPU, DDP is not your answer.
2. Ask what is actually limiting memory. If sharding fixes model states but activations still dominate, add activation checkpointing or reduce sequence and micro-batch before changing frameworks.
3. Measure the new bottleneck after sharding. If step time now shifts into all-gather and reduce-scatter, network quality and wrapping granularity matter more than the headline ZeRO stage.
4. Choose the operational stack. Prefer FSDP when you want native PyTorch tooling and the run fits once sharded. Prefer DeepSpeed when CPU or NVMe offload, or an existing DeepSpeed stack, is the real requirement.
5. Escalate topology only when needed. Add tensor parallelism when layers are too wide for one device, pipeline parallelism when depth must be split, and always verify stage-compatibility before combining features.

Quick decision reference