Skip to content
ML to LLMs
Labs / LLM and RAG Lab

Data parallelism and ZeRO memory

Split weights, gradients, and optimizer state across GPUs with DDP, ZeRO-1/2/3, and FSDP; compare memory per GPU and communication.

Theory and worked examples: ZeRO and FSDP

Data parallelism and ZeRO memory

Synthetic Llama-like model shapes, 2,048 tokens per GPU with gradient checkpointing. Estimates of training state, not measurements.

Memory per GPU by strategy0255075100memory per GPU (GB, estimate)DDPweights: 13.48 GBgradients: 13.48 GBoptimizer states: 80.86 GBactivations: 0.82 GB108.6 GBZeRO-1weights: 13.48 GBgradients: 13.48 GBoptimizer states: 10.11 GBactivations: 0.82 GB37.9 GBZeRO-2weights: 13.48 GBgradients: 1.68 GBoptimizer states: 10.11 GBactivations: 0.82 GB26.1 GBZeRO-3 / FSDPweights: 1.68 GBgradients: 1.68 GBoptimizer states: 10.11 GBactivations: 0.82 GB14.3 GB80 GB GPU
weightsgradientsoptimizer statesactivations

What each GPU stores under ZeRO-1: a full bar is a replicated copy, a slice is that GPU's 1/8 shard.

DDP1.00× DDP traffic
ZeRO-11.00× DDP traffic
ZeRO-21.00× DDP traffic
ZeRO-3 / FSDP1.50× DDP traffic

ZeRO-1 on 8 GPUs: 2P+2P+12P/N2P + 2P + 12P/N = 37.9 GB per GPU with activations, which fits in 80 GB.

The optimizer state is split 8 ways at the same communication cost as DDP. The replicated weights and gradients set a floor that more GPUs cannot remove.