Distributed training: data parallelism, sharding and failure recovery

About this article: AI-assisted writing, reviewed against the cited sources. This is source research, not a lab validation of the procedures. Check the stated versions and your own recovery requirements before making changes.

Data parallelism does not make the model smaller. Each worker holds a full copy of the weights, a full copy of the gradients, and a full copy of the optimiser state, and the workers only agree on the gradient. Sharding is what removes those copies. In the ordinary collective and torchrun recovery path described here, a worker failure does not leave the others to finish that collective: the collective is waiting for every rank, and the launcher documented by PyTorch stops the survivors and starts the group again. The steps since the last checkpoint are gone.

An AI writing system produced this article. No training job was launched. The speeds and memory figures below are the ZeRO paper's, on the hardware that paper names. No NetLeverage training system is described.

Why it matters

Rajbhandari and colleagues put a number on the gap that surprises people who have only looked at a weight file. Mixed-precision Adam, for a model of $\Psi$ parameters, needs 16$\Psi$ bytes for the model state alone: 2$\Psi$ for fp16 weights, 2$\Psi$ for fp16 gradients, and 12$\Psi$ for the optimiser, which keeps an fp32 copy of the weights, an fp32 momentum, and an fp32 variance. GPT-2 at 1.5 billion parameters is then at least 24 GB, against 3 GB to hold the fp16 weights. A flattened fp32 buffer of those gradients is another 6 GB. The file you copy to a laptop is not the thing that fills the GPU.

On the 32 GB GPUs of that paper, plain data parallelism ran out of memory above about 1.4 billion parameters. Adding a second GPU does not help that problem. It adds a second copy.

How it works

What a data-parallel step actually synchronises

In the ZeRO paper's description of data parallelism, the parameters are replicated. A mini-batch is split across processes. Each process runs forward and backward on its own samples. The gradients are averaged, and each process applies that same average to its own copy. The copies stay equal because they started equal and then saw the same update.

PyTorch's DistributedDataParallel design note, which the 2.5 documentation still carries and which that page says was written for the implementation as of version 1.4, fills in the mechanics and warns that the implementation has moved on. At construction, rank 0 broadcasts its state_dict so every replica starts from the same weights. Each process builds a reducer. Gradients are packed into buckets, filled in roughly the reverse of parameters(), because backward tends to finish the last layer first. A hook fires when a gradient is ready. When a bucket is complete, the reducer starts an asynchronous all-reduce that writes the mean into param.grad. It issues those all-reduces in bucket-index order, not in the order buckets happened to become ready. The note is explicit that a mismatched all-reduce order across processes produces wrong results or a backward pass that hangs. After the backward pass, the optimiser is a local optimiser. Nothing else is communicated. Replicas match because the start state matched and the averaged gradient matched.

find_unused_parameters exists for a loss that does not touch every parameter. The reducer still reduces every bucket. The flag only stops it waiting forever for a gradient that will never come, and the note says the autograd walk that finds those parameters has a cost, so it should be off unless the model actually skips parameters.

That design has two limits the rest of this article is about. The model state is replicated, so the largest model is the model that fits on one device. And the all-reduce has a participant list: a process that never arrives does not get averaged around. The step stops.

What the three ZeRO stages partition

ZeRO's model-state half, ZeRO-DP, keeps the compute pattern of data parallelism and partitions the state instead of replicating it. The stages stack.

The first two stages are where the bytes are. With $K=12$, the optimiser is three quarters of the 16$\Psi$ bytes, and the fp16 gradient is another eighth. Partitioning parameters is the stage that makes per-device memory fall with every added GPU, at the price of gathering a parameter when it is needed and discarding it afterwards.

Their arithmetic for a full stack: a trillion-parameter model with an Adam-style optimiser in 16-bit training needs about 16 TB for optimiser state, gradients, and parameters. Split over 1,024 GPUs that is 16 GB of model state per GPU, which would fit in a 32 GB device of that generation with room left over. They are careful about what was run. The full set of optimisations would still have left a trillion-parameter run longer than a year on the hardware they had, so the system they implemented and timed, ZeRO-100B, is the second stage plus the residual-memory work, not the parameter partition. On 400 V100 GPUs that system trained 100 billion parameter models at over 38 teraflops per GPU and over 15 petaflops in aggregate, which they report as more than ten times the training speed of the prior state of the art at that size. From 64 to 400 GPUs they saw superlinear speedup: doubling the GPUs more than doubled performance, because a larger data-parallel degree frees memory and a larger batch then fits. They also report training up to 13 billion parameters with no model parallelism and no pipeline parallelism, past the 1.4 billion parameter point where the data-parallel baseline ran out of memory.

The other half of memory is what they call residual state: activations kept for the backward pass, temporary buffers, and fragmentation. Activation checkpointing, they say, is not enough once the model is large. ZeRO-R partitions activations that model parallelism had been replicating, and it can offload activations to CPU. It also sizes the temporary buffers, and it allocates with tensor lifetime in mind so a free byte count is not split into holes too small to use. A failed allocation with memory still free is the symptom they are describing.

Model parallelism is not deleted by this. They give two cases where they still want it. Combined with ZeRO-R, it cuts activation memory on very large models. On smaller models, a pure data-parallel batch can get large enough to hurt convergence, a point they attribute to earlier work rather than measure again. ZeRO and model parallelism together have a theoretical per-device reduction of $N_d\times N_m$. Their sketch is 16-way model parallelism inside a node and 64-way data parallelism across nodes, on 1,024 GPUs. That is a fit argument, not the 400-GPU measurement.

How PyTorch sharding maps onto those stages

The PyTorch 2.15 FullyShardedDataParallel class page describes four strategies, and it says the class is inspired by ZeRO stage 3.

FULL_SHARD partitions parameters, gradients, and optimiser state. Parameters are all-gathered before the forward, resharded after it, all-gathered again before the backward, and resharded after that. Gradients are synchronised and sharded with a reduce-scatter after the backward. The optimiser step is local, on the shard that rank owns.

SHARD_GRAD_OP partitions gradients and optimiser state. Parameters are unsharded for the forward and stay unsharded until the backward finishes, which removes one all-gather and uses more memory. The advanced tutorial states the same trade: less communication, a higher memory footprint, and a gather on the next forward.

NO_SHARD replicates parameters, gradients, and optimiser state and all-reduces gradients. The class page calls this a pattern similar to DistributedDataParallel.

HYBRID_SHARD uses full sharding inside a node and replicates parameters across nodes, so the expensive gathers stay on the faster links.

The FSDP2 getting-started tutorial, which states that FSDP1 is deprecated, describes the same split in one paragraph: outside forward and backward the parameters are sharded; they are all-gathered before each; the backward reduce-scatters gradients; the optimiser updates sharded parameters with sharded gradients and so the optimiser state stays sharded. It calls this a decomposition of data-parallel all-reduce into reduce-scatter plus all-gather. fully_shard is applied to each layer and then to the root, so that during one layer the other layers remain shards. The optimiser is constructed after fully_shard, because the parameters it must see are DTensor shards, placed on dimension 0. Adam and gradient clipping accept those tensors. A full state dictionary is assembled through the DTensor APIs. DTensor permits sharded state dictionaries without first all-gathering a full model. That is not a claim that the distributed checkpoint save protocol performs no coordination: its own API documents collectives. A complete sharded checkpoint needs all required shards and metadata.

The class page's older helper full_optim_state_dict gathers a full optimiser dictionary, and with rank0_only=True only rank 0 receives the populated dictionary. A file written from any other rank under that setting is not a checkpoint.

What to choose and why

The 2.5 torchrun page is the failure contract. If $k$ of $n$ workers fail, all workers are stopped and restarted, up to the configured maximum. An agent or a node failure is a local worker-group failure; the job manager either fails the whole job or replaces the node. On a failure or a membership change, every surviving worker is killed immediately. The page tells you to checkpoint, and to set the interval from how much work you can afford to lose. RANK is not stable across restarts. A checkpoint that means "the file for rank 3" is not a checkpoint after the restart unless the slice is identified by its shard, not by the rank number that happened to hold it. World size is not stable either if the job is elastic (min different from max). The same page requires the same number of local workers on every node, recommends NCCL for GPU training, and says the first error by timestamp on each node is only a heuristic root cause.

The restarted training script must load the last complete checkpoint and repeat lost steps; torchrun does not implement that restore logic for the application. It does not resume the collective that was in flight. As an engineering requirement, checkpoint coordination must prevent a mixture of states from different optimiser steps. The FSDP2 tutorial's resume shape is the plain one: load, then train. Save on a cadence the lost-work budget allows, and save a full state or a real sharded state, not one shard renamed as the model.

The failure behaviour above is specifically torchrun's documented contract. TorchFT is a separate library for fault-tolerant DDP/HSDP that can change process-group membership at training-step boundaries without restarting the whole job. Its LocalSGD and DiLoCo modes are experimental. Neither path was tested here. The cited DDP and torchrun editions are 2.5, whereas the FSDP class page is 2.15; these sources are not a single release contract.

Risks and limits

• The DDP design note is dated. The 2.5 page says the note describes version 1.4 and that the implementation evolves. Bucket order, the mean reduction, the rank-0 broadcast, and the hang on mismatched all-reduce order are that note. Treat a newer PyTorch as a thing to read, not as a silent copy of 1.4.
• ZeRO's trillion-parameter paragraph is a memory budget, and the paper says the run would have taken more than a year. The measurement they publish is ZeRO-100B: second-stage partitioning plus residual-state work, 100 billion parameters, 400 V100s, the speeds above.
• The 4× and 8× figures are the paper's analysis for mixed-precision Adam, where $K=12$. A different optimiser has a different $K$. The communication claims (same volume, then 50% more at the third stage) are the paper's as well.
• Superlinear speedup in that paper came from larger batches fitting, between 64 and 400 GPUs. The same paper warns, by citation, that a large enough batch slows convergence. More GPUs are not free accuracy.
• FSDP1 and FSDP2 are both in the documents opened here. The FSDP2 tutorial says FSDP1 is deprecated. The class page still documents FullyShardedDataParallel and the four strategies. Compatibility and any needed conversion must be checked against the chosen state-dictionary and checkpoint APIs.
• torchrun kills survivors on purpose. A restart without a checkpoint repeats from the start of the run. Ranks move. Elastic jobs also change world size.
• No job was launched for this article.
• No NetLeverage product behaviour is claimed.

Sources

1. Rajbhandari, Rasley, Ruwase and He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, arXiv:1910.02054 — https://arxiv.org/html/1910.02054 (opened 10 November 2026).
2. PyTorch 2.5, Distributed Data Parallel design note — https://docs.pytorch.org/docs/2.5/notes/ddp.html (opened 10 November 2026; the page dates the design to v1.4).
3. PyTorch 2.5, torchrun (Elastic Launch) — https://docs.pytorch.org/docs/2.5/elastic/run.html (opened 10 November 2026).
4. PyTorch 2.15 (served by the stable alias on the review date), FullyShardedDataParallel — https://docs.pytorch.org/docs/2.15/fsdp.html (opened 10 November 2026).

5. PyTorch, Getting Started with Fully Sharded Data Parallel (FSDP2) — https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html (opened 10 November 2026).

6. PyTorch, Advanced Model Training with Fully Sharded Data Parallel — https://docs.pytorch.org/tutorials/intermediate/FSDP_advanced_tutorial.html (opened 10 November 2026). SHARD_GRAD_OP memory/communication trade-off.

7. TorchFT README — https://github.com/meta-pytorch/torchft (opened 10 November 2026). Per-step fault tolerance, a separate recovery design.
8. PyTorch 2.15, distributed checkpoint — https://docs.pytorch.org/docs/2.15/distributed.checkpoint.html (opened 10 November 2026). Sharded checkpoint coordination and collectives.

This article was originally published on 10 November 2026.