Data parallelism, ZeRO & FSDP

1

Concept

One accelerator can process only a limited batch and hold only a limited model state. Data parallelism increases training throughput by giving several workers identical model replicas and different examples. If rank $r$ computes gradient $g_r$, a collective reduction forms the mean or sum used by every replica. Starting from equal weights and applying equal gradients keeps replicas synchronized.

The global batch is the combined work across ranks and gradient-accumulation steps. If there are $P$ ranks, local microbatch $b$, sequence length $T$, and $a$ accumulation steps, the global batch contains $PbaT$ tokens, assuming no padding loss. Scaling worker count without accounting for global batch changes optimization, not only speed.

Replication is memory-expensive. Mixed-precision Adam training may keep low-precision parameters, gradients, higher-precision master parameters, and two optimizer moment tensors. Exact bytes depend on implementation and dtype, but optimizer state can exceed the memory for inference weights. Ordinary data parallelism duplicates all of it on every rank.

ZeRO decomposes that redundancy in stages. Stage 1 shards optimizer state. Stage 2 also shards gradients. Stage 3 also shards parameters. Each rank owns only a fraction of the persistent state, while collectives make the necessary pieces available for computation. The names describe memory ownership, not a new optimization algorithm.

Fully Sharded Data Parallel (FSDP) implements the fully sharded pattern around modules. Before a module’s forward pass, ranks all-gather its parameter shards to materialize the full parameters needed locally. After use, parameters can be resharded. During backward, communication reconstructs parameters as needed and reduce-scatter distributes gradient shards to their owners. Communication can overlap with computation when ordering and network permit.

Wrapping policy matters. If the sharded unit is too large, peak all-gather memory is large and overlap is limited. If it is too small, collectives become numerous and latency dominates. Transformer-block boundaries are common because they offer substantial compute between communications. Prefetching the next block can hide latency but increases simultaneous memory.

Checkpointing becomes distributed too. A full state-dict gather may exceed host memory. Sharded checkpoints let ranks write their owned tensors, accompanied by metadata that reconstructs the global names and shapes. Resuming under a different world size requires a format and loader that support resharding. Saving only weights is insufficient for exact training continuation; optimizer, scheduler, scaler, random-number, and data-loader states matter.

Failures must be collective. If one rank encounters bad data while others enter an all-reduce, the job can hang rather than throw a clean exception. Timeouts, coordinated error propagation, rank-aware logs, and reproducible batch IDs are operational necessities. Network topology also matters: fast links inside a node and slower links across nodes may favor hierarchical collectives.

The 3D visual for this lesson should distinguish ownership from temporary materialization. Show persistent parameter, gradient, and optimizer shards as colored slices on ranks. Then animate an all-gather around one block, local compute on different data, and reduce-scatter of gradients. Leaving the full model displayed on every rank would misrepresent the memory benefit.

The stable model is replicated computation over different data with increasingly sharded state. Data parallelism buys throughput. ZeRO and FSDP remove redundant memory. They do not provide free capacity: the price is communication, orchestration, and more demanding checkpoint and failure semantics.