Distributed Training: FSDP & ZeRO

A 70-billion-parameter (70B) training job can fit its persistent model state across 24 80 GB GPUs and still fail during backward, when one layer's parameters are gathered again. Move to 32 GPUs, avoid the out-of-memory error, then lose throughput because ranks wait on the network. Same model, two different bottlenecks: transient memory and scheduled communication.

The mixed-precision chapter left you with a dtype recipe: cheap compute, careful update storage, and a communication dtype you record separately. That recipe doesn't shrink the model. It only says how many bytes each parameter still occupies.

Full-parameter supervised fine-tuning (SFT) on one GPU fails when those bytes, plus activations, no longer fit. Fully Sharded Data Parallel (FSDP) and ZeRO (Zero Redundancy Optimizer) attack the next bottleneck: replicated model states across data-parallel ranks. Communication becomes part of the training cost.

We'll keep this 70B route as a running example: count persistent states, watch a gathered unit, choose a sharding policy, then test topology and resume behavior.

Before choosing a sharding policy, predict which memory is actually missing. One useful mixed-precision sizing convention, used in the ZeRO analysis, assumes low-precision parameters and gradients plus an FP32 master parameter copy and FP32 Adam moments.^{[1]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054} 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's 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.

Sharding that 16 GB evenly across eight data-parallel ranks suggests 2 GB of steady-state model state per rank, not a 2 GB peak allocation. The currently gathered layer, activations, communication buffers, and allocator overhead still need their own memory budget.

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 Distributed Data Parallel (DDP) alone won't make it fit.

The first figure tracks only model states. Activations and transient buffers stay outside the columns because sharding model states doesn't remove those costs.

The memory wall in numbers

Before reading the totals, predict which state dominates. The answer is usually the FP32 optimizer and master-weight inventory, not the low-precision parameter copy.

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's about 1.12 TB of model states before activations, temporary buffers, or allocator overhead. It's an assumption-bound calculation, not a promise about every optimizer implementation.

The arithmetic is small enough to turn into a quick check. Use decimal GB so the computed numbers match this table.

Why DDP hits a limit

Use the running 70B example to test the first instinct: if you add more data-parallel ranks, does each rank's model-state copy get smaller? It doesn't. Replication increases the number of copies while splitting only the input batch.

Ordinary DDP (Distributed Data Parallel) 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 that sums values across GPUs and returns the result to every GPU) after each backward pass. For a model that already fits, extra GPUs buy throughput without changing per-GPU memory.

That stops helping when the model itself is the problem. DDP replicates the full 1.12 TB of 70B model states on every rank (process/GPU). Sixteen GPUs give you sixteen copies, not a smaller copy. FSDP and ZeRO attack that redundancy by sharding model states across the data-parallel group.

ZeRO partitions more of that state at each stage.^{[1]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054} ZeRO-1 shards optimizer states, ZeRO-2 also shards gradients, and only ZeRO-3 shards parameters. Under ZeRO-3, each GPU keeps a fraction of the parameters and gathers the shards it needs for compute.

• 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

Each stage answers one storage question: which model state can stop being replicated while the rank still completes its update?

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: Same modeled volume as DDP. Gradient sync is still an all-reduce-class $2\Psi$ step ($\Psi$ is the parameter-element count). Ownership changes; the byte volume doesn't.

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 modeled volume stays $2\Psi$, matching DDP: reduce-scatter the gradients, then all-gather the updated parameter partition. Each GPU receives only its shard.

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.

ZeRO memory trade-off: 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. You pay in network bandwidth for what you save in persistent per-GPU state.

Memory per GPU (70B model, 256 GPUs)

This table 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.

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

Once state is sharded, the memory win has a matching communication bill. Follow one parameter shard through the collectives before comparing frameworks or networks.

ZeRO-3 doesn't delete those 1.12 TB. It spreads the persistent copies and pays the interconnect every time a block needs its full weights. The network is now part of the training system, not an afterthought.

Communication primitives you need to know

Read each panel left to right: same ranks, different "after" state. The two-rank numbers in the snippet below are the same story with fewer GPUs: all-reduce leaves [11, 22] on both ranks, reduce-scatter leaves one summed slice per rank, and all-gather rebuilds the parameter vector from shards.

The 1.5x heuristic

People often summarize ZeRO-3 communication as "about 1.5x DDP." That's the paper's bound, not a throughput guarantee.^{[1]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}

For a model with $\Psi$ parameter elements (or $P$ parameter bytes, same ratio):

1. DDP implements gradient all-reduce as reduce-scatter plus all-gather, modeled as $2\Psi$ moved per rank per step.
2. ZeRO-1 and ZeRO-2 stay at $2\Psi$: they change who stores optimizer state and gradients, not the modeled byte volume.
3. ZeRO-3 / FULL_SHARD adds a second parameter all-gather so backward can run after the forward copy was freed. Two all-gathers plus one gradient reduce-scatter land at $3\Psi$.

That is the $3\Psi / 2\Psi = 1.5x$ rule. The harder part in production 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, step time can fall off quickly even when the byte ratio looks modest.

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.

Why can ZeRO-3 be slower than DDP even if the back-of-the-envelope byte ratio is only about 1.5x?

Answer

Because the real cost isn't total bytes alone. ZeRO-3 and FSDP break communication into many layer-level all-gathers and reduce-scatters. On a weak or high-latency interconnect, those smaller, more frequent collectives can hurt throughput more than the simple 3P versus 2P byte model suggests.

FSDP (Fully Sharded Data Parallel)

ZeRO names the sharding strategy; FSDP is PyTorch's native way to express much of that strategy in a model and training loop. The choice now becomes an API and composition decision, not another memory formula.

FSDP is PyTorch's native sharded data-parallel stack.^{[2]Reference 2PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel.https://arxiv.org/abs/2304.11277} Current docs start new designs on torch.distributed.fsdp.fully_shard (FSDP2): per-parameter DTensor shards, no FlatParameter wrapper, and a migration guide from FullyShardedDataParallel (FSDP1). The tutorial marks FSDP1 deprecated. You'll still read FSDP1 in older repos; start new work on fully_shard.^{[3]Reference 3torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html[4]Reference 4Getting Started with Fully Sharded Data Parallel (FSDP2)https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html}

People pick FSDP when they want PyTorch-native profilers, checkpointing, and DeviceMesh composition. torch.compile guidance is still to compile the highest-level train or eval step, or the top-level module, and to fall back to the inner module if a distributed wrapper causes issues. FSDP1 compile also needs use_orig_params=True. FSDP2 always uses original parameters, so that flag is gone.^{[5]Reference 5Where to apply torch.compile?https://docs.pytorch.org/docs/main/user_guide/torch_compiler/compile/programming_model.where_to_apply_compile.html[6]Reference 6FullyShardedDataParallelhttps://docs.pytorch.org/docs/stable/fsdp.html}

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

Predict the peak for one wrapped block: a small shard sits at rest, a full unit appears for forward, and another full unit may appear for backward when the first copy was freed.

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.

FSDP2: fully_shard on DTensors

If you construct the optimizer before sharding, it sees the wrong parameter representation. Predict the safe order: shard blocks, shard the root, then create the optimizer over the resulting DTensor parameters.

For new code, apply fully_shard bottom-up: each transformer block becomes one communication group, then the root call shards leftover embeddings and heads. Build the optimizer after that, over the resulting DTensor parameters.^{[3]Reference 3torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html[4]Reference 4Getting Started with Fully Sharded Data Parallel (FSDP2)https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html}

reshard_after_forward is the FSDP2 name for the old sharding-strategy trade-off. True frees the full unit after forward and all-gathers again in backward (ZeRO-3 / FULL_SHARD behavior). False keeps the unsharded parameters until backward finishes (ZeRO-2-like / SHARD_GRAD_OP). The default is True for non-root modules and False for the root, because backward starts at the root and that extra all-gather would be wasted.

MixedPrecisionPolicy continues the previous chapter's dtype split. param_dtype is the unsharded compute and all-gather dtype. reduce_dtype is the gradient reduce-scatter dtype, and it can differ. Persistent sharded parameters stay in the original dtype, which is what the optimizer steps. Current docs are explicit: this setup doesn't add a second FP32 master copy on top of already-high-precision shards.^{[3]Reference 3torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html[4]Reference 4Getting Started with Fully Sharded Data Parallel (FSDP2)https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html}

This snippet is structure, not a local notebook cell. fully_shard needs a process group on real accelerators. The learning rate is an SFT starting point to evaluate.

Write the two recipes down before you size a cluster. The ZeRO paper's 16-byte convention stores low-precision parameters and gradients plus an FP32 master. FSDP2 mixed precision keeps the high-precision shard as the optimizer parameter, so you don't add that extra master on top.

That 840 GB figure still isn't "fits on one GPU." It assumes the module was built in FP32, so the persistent shard is four bytes per parameter before gradients. The number collides with the earlier Adam+master column for a different reason: 4-byte parameters plus 8-byte moments, with no extra master on top. Sharded gradients still exist after backward; this comparison only isolates the extra-master question. Don't paste the ZeRO 16-byte recipe onto an FSDP2 MixedPrecisionPolicy run without checking which tensors are persistent.

FSDP2 also handles mixed frozen and trainable parameters in one communication group without the extra memory FSDP1 often spent on that case. That's the hinge into LoRA: most weights can stay frozen on purpose, and the sharded runtime shouldn't pretend they're all updating.

FSDP1 you'll still meet

Older code wraps FullyShardedDataParallel. The same three rules still matter:

1. Wrap transformer blocks, not the whole model as one unit.
2. Create the optimizer after wrapping.
3. Don't treat limit_all_gathers=True as an overlap knob. backward_prefetch is the overlap control; limit_all_gathers caps in-flight all-gathers to control peak memory.^{[6]Reference 6FullyShardedDataParallelhttps://docs.pytorch.org/docs/stable/fsdp.html}

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.

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:

Sharding strategy names

Checkpointing, resume, and fault recovery

An out-of-memory fix that can't resume faithfully is still an incomplete training system. Treat the checkpoint as the state needed to take the next optimizer step, not as a copy of weights for inference.

Large training runs fail in the middle all the time. Nodes reboot, preemptible instances disappear, and jobs get rescheduled. If your checkpoint only contains model weights, you didn't save the training run. You saved an export artifact.^{[7]Reference 7Checkpointing in torchtune.https://meta-pytorch.org/torchtune/stable/deep_dives/checkpointer.html}

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]Reference 3torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html[7]Reference 7Checkpointing in torchtune.https://meta-pytorch.org/torchtune/stable/deep_dives/checkpointer.html}

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]Reference 8Universal Checkpointing with DeepSpeed: A Practical Guide.https://www.deepspeed.ai/tutorials/universal-checkpointing/} 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

If native sharding solves persistent GPU memory but the run still needs host or NVMe capacity, the unresolved question is where those states should live and what transfer cost the step can tolerate.

DeepSpeed is Microsoft's distributed training runtime that introduced ZeRO and provides heterogeneous-memory offload through ZeRO-Offload and ZeRO-Infinity.^{[1]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054[9]Reference 9ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning.https://arxiv.org/abs/2104.07857} It has 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]Reference 10Training APIhttps://deepspeed.readthedocs.io/en/latest/training.html}

This distributed setup snippet configures and initializes a model with 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]Reference 9ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning.https://arxiv.org/abs/2104.07857}

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 through 3, while PipelineModule isn't compatible with ZeRO-2 or ZeRO-3.^{[10]Reference 10Training APIhttps://deepspeed.readthedocs.io/en/latest/training.html[11]Reference 11Pipeline Parallelismhttps://deepspeed.readthedocs.io/en/latest/pipeline.html}

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.

Which stack to start with

• Use FSDP when you want the PyTorch-native path, DeviceMesh composition, and standard profilers and checkpoints. GPU memory after sharding, or FSDP2's CPU offload, is enough.
• Use DeepSpeed when NVMe spillover is non-negotiable, you need ZeRO-Infinity's host/storage hierarchy, or you're already in a DeepSpeed or Megatron-DeepSpeed stack whose TP/PP requirements fit the current stage-compatibility rules.

You already know the model-state math fits once sharded. What question decides whether you stay with FSDP or reach for DeepSpeed?

Answer

Ask what requirement remains unsolved after sharding. If you mainly want native PyTorch tooling and standard checkpointing, stay with FSDP. If you still need NVMe offload, ZeRO-Infinity, or the rest of the stack already depends on DeepSpeed-specific runtime features, DeepSpeed becomes the better operational fit.

Beyond data parallelism: 3D parallelism

When sharding alone leaves either one layer or one sequence too large, name the dimension that must be split before adding another parallel axis. That prediction keeps a topology from becoming a list of acronyms.

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]Reference 12Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.https://arxiv.org/abs/1909.08053}

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.

Solid arrows are activations moving through pipeline stages. Dotted arrows are the matching data-parallel replicas syncing. The figure is a topology, not a recommendation to turn every axis on.

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]Reference 13Parallelism Strategies Guide.https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html}

• 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]Reference 3torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html[13]Reference 13Parallelism Strategies Guide.https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html}

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

Debugging & profiling distributed training

Distributed training turns a local symptom into a coordination question. Before changing a config, decide whether the first failure is a rank mismatch, a memory peak, a data stall, or a communication schedule.

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 clipping: 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]Reference 6FullyShardedDataParallelhttps://docs.pytorch.org/docs/stable/fsdp.html[4]Reference 4Getting Started with Fully Sharded Data Parallel (FSDP2)https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html}

Profiling tools

Throughput (tokens/sec) only moves after you know whether the step is compute-bound or communication-bound. A credible comparison records accelerator model and count, topology, model and sequence shape, micro-batch and accumulation, dtype, sharding and prefetch settings, warmup and steady-state window, exact baseline, and correctness or convergence checks.

Report the boundary you changed with the result. "FSDP is faster" is not reproducible evidence until another engineer can reconstruct the workload and verify that both runs reached the same loss or task quality.

• 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]Reference 6FullyShardedDataParallelhttps://docs.pytorch.org/docs/stable/fsdp.html[3]Reference 3torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html}
• 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's slow or unstable. Good training dashboards separate data, compute, communication, and memory instead of showing only one loss curve.^{[13]Reference 13Parallelism Strategies Guide.https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html}

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

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. Weak answers stop at "use FSDP" or "use DeepSpeed." Stronger answers say "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

The earlier table explains why sharding helps. Now use the same recipe to answer a hardware question, then separate a mathematical lower bound from a runnable plan.

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 raw model-state math says a 70B-style run needs about 14 GPUs with ZeRO-3, but you only have sixteen 80 GB cards. Why is that still a shaky plan?

Answer

Because 14 GPUs is only the lower bound for model states. On 16 GPUs, model states still land around 70 GB per card, leaving little room for activations, communication buffers, and allocator overhead. ZeRO-3 shrinks memory by sharding optimizer states, gradients, and parameters, but it doesn't remove the need for headroom.

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

Each failure below has a recovery path: identify what was split, observe the first failing boundary, then change one memory or topology axis and rerun the smallest faithful case.

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 don't 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's sharding. If one layer's matmul is split across devices, that's 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: Apply fully_shard to each transformer block, then the root. In FSDP1 code, use auto_wrap_policy the same way. The code examples earlier show both.

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, and FSDP2's reshard_after_forward=False, keep parameters unsharded through forward and backward, then reshard after backward.^{[6]Reference 6FullyShardedDataParallelhttps://docs.pytorch.org/docs/stable/fsdp.html[4]Reference 4Getting Started with Fully Sharded Data Parallel (FSDP2)https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html} That changes peak memory 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 or reshard_after_forward=True rather than assuming the ZeRO-2-like name matches DeepSpeed stage 2.

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 every ZeRO stage

• Symptom: You configure PipelineModule with ZeRO-3 and the job fails during initialization.

• Cause: DeepSpeed has documented compatibility limits. AutoTP supports ZeRO stages 0 through 3, but 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.