ZeRO and FSDP sharding

Deep dive

How it works

In: a model's parameters, gradients and optimizer state, N data-parallel ranks, each with its own micro-batch · Out: the same updates as plain data parallelism · State: per rank a 1/N shard of fp32 master weights, gradients and Adam moments · Cost: 16Ψ/N bytes of model state per GPU; 3Ψ communicated elements per step (ZeRO-3), 1.5x DDP

Flatten each layer's (or unit's) parameters into one buffer padded to a multiple of N. Rank r owns slice r.

forward, per unit:   all-gather bf16 weights (Ψ_unit) → compute → free the non-owned part
backward, per unit:  all-gather again → compute grads → reduce-scatter grads (each rank gets its slice) → free
step:                AdamW on the fp32 master slice (Ψ/N) → cast to bf16 shard
                     ZeRO-1/2: one all-gather of updated bf16 params after the step

Memory per GPU (Ψ params, N ranks, K = 12 optimizer bytes):

Communication per step in elements: DDP all-reduce = reduce-scatter + all-gather = 2Ψ. ZeRO-3 does all-gather (fwd) + all-gather (bwd) + reduce-scatter = 3Ψ. Prefetching unit k+1's all-gather while unit k computes keeps it off the critical path.

Code

ZeRO-2 on threads: reduce-scatter the gradients, run AdamW on the local slice of the fp32 master weights and optimizer state, all-gather the updated weights. The result must equal unsharded AdamW on the averaged gradient, bit for bit.

use std::sync::{Barrier, Mutex};
use std::thread;

/// Shared-memory stand-in for NCCL. A real collective moves only the needed bytes (see ring all-reduce).
struct Comm { n: usize, barrier: Barrier, slots: Mutex<Vec<Vec<f32>>> }

impl Comm {
    fn new(n: usize) -> Self { Self { n, barrier: Barrier::new(n), slots: Mutex::new(vec![vec![]; n]) } }
    fn exchange(&self, rank: usize, buf: Vec<f32>) -> Vec<Vec<f32>> {
        self.slots.lock().unwrap()[rank] = buf;
        self.barrier.wait();
        let all = self.slots.lock().unwrap().clone();
        self.barrier.wait();                                   // nobody overwrites before all have read
        all
    }
    /// Sum over ranks, but rank r receives only shard r.
    fn reduce_scatter(&self, rank: usize, full: Vec<f32>) -> Vec<f32> {
        let s = full.len() / self.n;
        let all = self.exchange(rank, full);
        (rank * s..(rank + 1) * s).map(|i| all.iter().map(|b| b[i]).sum()).collect()
    }
    fn all_gather(&self, rank: usize, shard: Vec<f32>) -> Vec<f32> { self.exchange(rank, shard).concat() }
}

struct AdamW { m: Vec<f32>, v: Vec<f32>, t: i32 }
impl AdamW {
    fn new(n: usize) -> Self { Self { m: vec![0.0; n], v: vec![0.0; n], t: 0 } }
    fn step(&mut self, w: &mut [f32], g: &[f32], lr: f32) {
        let (b1, b2, eps, wd) = (0.9f32, 0.95f32, 1e-8f32, 0.1f32);
        self.t += 1;
        let (bc1, bc2) = (1.0 - b1.powi(self.t), (1.0 - b2.powi(self.t)).sqrt());
        for i in 0..w.len() {
            self.m[i] = b1 * self.m[i] + (1.0 - b1) * g[i];
            self.v[i] = b2 * self.v[i] + (1.0 - b2) * g[i] * g[i];
            w[i] = w[i] * (1.0 - lr * wd) - lr / bc1 * self.m[i] / (self.v[i].sqrt() / bc2 + eps);
        }
    }
}

/// One data-parallel rank. Local loss ½‖w − target‖², so the local gradient is w − target.
fn zero2_rank(rank: usize, comm: &Comm, mut w: Vec<f32>, target: &[f32], steps: usize) -> Vec<f32> {
    let shard = w.len() / comm.n;                              // pad P to a multiple of N in real code
    let mut master = w[rank * shard..(rank + 1) * shard].to_vec(); // fp32 master: my slice only
    let mut opt = AdamW::new(shard);                           // m, v: my slice only
    for _ in 0..steps {
        let g: Vec<f32> = w.iter().zip(target).map(|(w, t)| w - t).collect();
        let mut gs = comm.reduce_scatter(rank, g);             // full gradient freed here (ZeRO-2)
        gs.iter_mut().for_each(|x| *x /= comm.n as f32);
        opt.step(&mut master, &gs, 0.01);
        w = comm.all_gather(rank, master.clone());             // bf16 copy in practice
    }
    w
}

fn main() {
    let (n, p, steps) = (4, 1000, 50);
    let targets: Vec<Vec<f32>> = (0..n).map(|r| (0..p).map(|i| ((i * 7 + r * 13) % 11) as f32).collect()).collect();
    let comm = Comm::new(n);
    let results: Vec<Vec<f32>> = thread::scope(|s| {
        let hs: Vec<_> = (0..n).map(|r| { let (c, t) = (&comm, &targets[r]); s.spawn(move || zero2_rank(r, c, vec![0.0; p], t, steps)) }).collect();
        hs.into_iter().map(|h| h.join().unwrap()).collect()
    });

    let (mut w, mut opt) = (vec![0.0f32; p], AdamW::new(p));    // reference: unsharded AdamW
    for _ in 0..steps {
        let g: Vec<f32> = (0..p).map(|i| targets.iter().map(|t| w[i] - t[i]).sum::<f32>() / n as f32).collect();
        opt.step(&mut w, &g, 0.01);
    }
    assert!(results.iter().all(|r| r == &w));
    println!("ZeRO-2 == AdamW; optimizer floats per rank {} instead of {}", 3 * p / n, 3 * p);
}

In PyTorch, ZeRO-3 is FSDP2: call fully_shard(block, mp_policy=MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=fp32)) on each transformer block, then on the root module. The optimizer then sees only local DTensor shards.

Complexity and constants

• 70B params × 16 B = 1.12 TB of model state. ZeRO-3 over 64 GPUs: 17.5 GB each, before activations and gather buffers.
• Peak extra memory under ZeRO-3 is about 2 units of full bf16 weights (current + prefetched). For a 70B model with 80 blocks, about 1.75 GB each.
• Overlap condition, per unit: all-gather time 2Ψ_unit·(N−1)/N / busbw ≤ unit compute time. For a 70B block (about 0.87B params, 1.75 GB bf16) at about 400 GB/s intra-node: about 4 ms. That is hidden when tokens per GPU per step are large (e.g. 8k tokens: about 6·0.87e9·8k ≈ 42 TFLOP ≈ 60 ms), and not hidden at tiny batches.
• Across nodes at about 50 GB/s per GPU, the same gather takes about 35 ms. That motivates HSDP (shard in a node, replicate across) and ZeRO++ (quantized gathers, hierarchical partitions).

Efficiency

• Cost: model state falls from 16Ψ to 16Ψ/N bytes per GPU: 70B parameters are 1.12 TB, 17.5 GB per GPU over 64 GPUs before activations. ZeRO-3 communicates 3Ψ elements per step against 2Ψ for DDP. A 70B block's all-gather takes about 4 ms intra-node at about 400 GB/s, hidden behind about 60 ms of compute at 8k tokens per GPU, but about 35 ms across nodes at about 50 GB/s (computed above). Bound: all-gather bandwidth when compute per unit is small.
• Levers:
  • ZeRO stage 1 → 3 — 120 → 31.4 → 16.6 → 1.9 GB for 7.5B parameters on 64 GPUs (Rajbhandari et al. 2020) — memory-bound training
  • prefetch the next unit's all-gather — keeps the 1.5x traffic off the critical path — large tokens per GPU
  • hybrid sharding (HSDP) — shard in a node, replicate across — slow inter-node links
  • ZeRO++ quantized gathers and hierarchical partition — about 4x less cross-node traffic (Wang et al. 2023) — multi-node
  • CPU/NVMe offload — capacity for PCIe bandwidth of about 25-50 GB/s (ZeRO-Offload, Infinity) — few GPUs
• Naive → standard → state of the art:
  • DDP, full replica per GPU · 16 bytes per parameter on every GPU
  • ZeRO-3 / FSDP with prefetch · 16Ψ/N, 1.5x DDP traffic
  • HSDP or ZeRO++ with per-parameter DTensor sharding (FSDP2) · less cross-node traffic, same memory per shard

Implementation notes

• Wrap at transformer-block granularity. Too fine means many small collectives (latency-bound); too coarse means a large gather buffer.
• reshard_after_forward=False keeps weights gathered from forward to backward: ZeRO-2-like memory for less communication.
• Reduce gradients in fp32 (reduce_dtype) when N is large; gather weights in bf16.
• Gradient clipping needs the global norm: sum squared norms of the local shards, then all-reduce. FSDP's clip_grad_norm_ does this.
• Checkpoint sharded (each rank writes its shard, as with torch.distributed.checkpoint). Gathering a 1 TB state onto rank 0 is slow and can OOM. Resharding to a different N happens at load time.
• ZeRO-Offload / Infinity moves the optimizer state (and the step itself) to CPU or NVMe. It trades PCIe bandwidth (about 25-50 GB/s) for capacity. Useful for fine-tuning on few GPUs.

Pitfalls

• Accessing module.weight outside forward under FSDP returns a shard or DTensor, not the full tensor. Use summon_full_params or full_tensor().
• Shared or tied parameters must live in the same FSDP unit, or they are gathered twice and their gradients split.
• Padding: the flat numel must be divisible by N. Forgetting it corrupts the last rank's slice.
• Mixing ZeRO-3 with gradient accumulation without no_sync-style handling reduce-scatters every micro-batch.
• Optimizers that need whole tensors (Shampoo, Muon's per-matrix orthogonalization, LAMB/LARS trust ratios) break with flat-shard ZeRO. FSDP2's per-parameter DTensor sharding makes this tractable but not free.

Variants and alternatives

In production

DeepSpeed ZeRO 1/2/3, Offload and Infinity; PyTorch FSDP (FlatParameter) and FSDP2 (fully_shard, DTensor); TorchTitan; Megatron-LM distributed optimizer (ZeRO-1); JAX/XLA via sharding annotations (GSPMD); Hugging Face Accelerate.

Implementations

Sources

• Rajbhandari, Rasley, Ruwase, He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020. paper · code
• Ren et al., ZeRO-Offload: Democratizing Billion-Scale Model Training, USENIX ATC 2021. paper
• Rajbhandari et al., ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning, SC 2021. paper
• Zhao et al., PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, VLDB 2023. paper · code
• Wang et al., ZeRO++: Extremely Efficient Collective Communication for Giant Model Training, 2023. paper