Data parallelism, ZeRO & FSDP

01 · Concept

Concept

Qwen3.8-27B ships as bf16 safetensors totalling about 54 GB. That is a checkpoint-storage fact, not a training-memory recipe and not proof that the full serving stack fits a nominal 80 GB device: KV cache, activations, runtime workspaces, and fragmentation also consume memory. Training adds persistent and transient categories whose dtypes must be stated rather than inferred from the released checkpoint.

Lesson 9.3 owns the course-wide derivation of this model’s memory budget. This lesson imports one explicit scenario from it: bf16 parameters and gradients plus fp32 master weights and two fp32 AdamW moments, totalling about 432 GB of persistent state for 27 billion parameters before activations and workspaces. It is a common conservative recipe, not a claim about Qwen’s undisclosed training run and not a universal requirement; pure-bf16 and alternative optimizer-state policies exist. The ZeRO arithmetic below is conditional on that scenario so the sharding stages can be compared consistently.

Data parallelism is the starting point and does nothing about this. Give several workers identical replicas and different examples; if rank $r$ computes gradient $g_r$, a collective reduction forms the mean used by every replica, and starting from equal weights with equal gradients keeps them synchronized. With $P$ ranks, local microbatch $b$, sequence length $T$, and $a$ accumulation steps, the global batch holds $PbaT$ tokens — so adding workers changes optimization, not merely speed. Under ordinary replication, every rank still holds the full assumed persistent state.

ZeRO removes the redundancy in three stages, and the effect is worth computing on a concrete cluster. Take 64 ranks and the illustrative 432 GB budget imported from lesson 9.3.

Stage 1 shards optimizer state only. Each rank keeps full bf16 parameters and gradients but $1/64$ of the optimizer: $324/64 = 5.06$ GB. Per-rank persistent state becomes $54+54+5.06 \approx 113$ GB. Still impossible on 80 GB cards.

Stage 2 also shards gradients: $54/64 = 0.84$ GB. Now $54+0.84+5.06 \approx 59.9$ GB. This fits — barely. About 20 GB remains for activations, temporaries, fragmentation, and communication buffers on a model with hidden size 5120 and long sequences, which lesson 5.10 will show is not much.

Stage 3 shards parameters as well, so persistent state is simply the whole budget divided by the world size: $432/64 = 6.75$ GB per rank. Full parameters for a module exist only transiently, materialized just before they are needed.

Fully Sharded Data Parallel is that fully sharded pattern implemented around modules. Before a module’s forward, ranks all-gather its parameter shards to reconstruct the full parameters locally; after use, they can be resharded. During backward, communication reconstructs parameters again as needed, and reduce-scatter distributes gradient shards to their owners. Where ordering and network permit, that communication overlaps with computation.

Wrapping policy decides how well stage 3 actually performs. If the sharded unit is too large, peak all-gather memory is large and overlap is limited; too small, and collectives become numerous and latency dominates. Transformer-block boundaries are the usual choice because they place substantial compute between communications. This model offers an unusually natural alternative unit: its 64 layers are organized as 16 repeating super-blocks of three Gated DeltaNet layers followed by one full-attention layer. Wrapping per super-block gives 16 uniform units of about 1.7 billion parameters each, which is coarse but perfectly balanced; wrapping per layer gives 64 finer units of two distinct shapes. Either is defensible, and only measurement decides.

Two operational consequences follow. Checkpointing becomes distributed: a full state-dict gather may exceed host memory, so ranks write their owned tensors alongside metadata that reconstructs global names and shapes, and resuming at a different world size requires a format and loader that support resharding. Saving weights alone is not enough for exact continuation — optimizer moments, scheduler position, loss scaler, random-number state, and data-loader cursor all matter, and under the illustrative recipe the optimizer state is the largest checkpoint category.

Failures must also be collective. If one rank hits bad data while the others enter an all-reduce, the job hangs rather than raising a clean exception. Timeouts, coordinated error propagation, rank-aware logs, and reproducible batch identifiers are not niceties. Network topology shapes the design too: fast links inside a node and slower links between nodes favour hierarchical collectives, and a stage-3 job whose all-gathers cross a slow fabric can spend more time moving parameters than using them.

The stable model is replicated computation over different data, with increasingly sharded state behind it. Data parallelism buys throughput. ZeRO and FSDP remove redundant memory. Neither creates capacity for free: the price is communication volume, scheduling complexity, and much more demanding checkpoint and failure semantics.