Data parallelism and ZeRO optimization

What is Data Parallelism?

How Standard Data Parallelism Works

The training loop on each GPU (rank $r$ ):

Forward pass: Process mini-batch $\mathcal{B}_r$ through local model copy $\theta_r$
Backward pass: Compute local gradients $\nabla_{\theta_r} \mathcal{L}(\mathcal{B}_r)$
AllReduce: Synchronize gradients across all GPUs: $\bar{\nabla}_\theta = \frac{1}{N} \sum_{r=1}^{N} \nabla_{\theta_r} \mathcal{L}(\mathcal{B}_r)$
Update: Apply optimizer: $\theta \leftarrow \theta - \eta \bar{\nabla}_\theta$
Broadcast: All GPUs now have identical updated parameters

WHY does this work? Because the gradient of the full dataset is the average of gradients from each subset (linearity of expectation). Each GPU computes a noisy estimate of the true gradient, and averaging reduces variance.

Memory Consumption Analysis

Formula Memory Per GPU in Standard Data Parallelism

For a model with $\Phi$ parameters (each 4 bytes for fp32, 2 bytes for fp16): $M_{\text{gpu}} = \underbrace{2\Phi}_{\text{params (fp16)}} + \underbrace{2\Phi}_{\text{gradients}} + \underbrace{12\Phi}_{\text{optimizer states}} + \underbrace{M_{\text{activations}}}_{\text{batch-dependent}}$

Total model state: $16\Phi$ bytes per GPU (for mixed precision + Adam optimizer)

Derivation from first principles:

Parameters (fp16): $2\Phi$ bytes (half precision for compute efficiency)
Gradients: $2\Phi$ bytes (same size as parameters, fp16)
Optimizer states (Adam):
- Momentum: $4\Phi$ bytes (fp32 for numerical stability)
- Variance: $4\Phi$ bytes (fp32)
- Master weights: $4\Phi$ bytes (fp32 copy for updates)
- Total: $12\Phi$ bytes WHY fp32 for optimizer? Mixed precision training uses fp16 for speed but maintains fp32 optimizer states to prevent underflow in tiny gradient updates. The update $\theta_{\text{fp32}} \leftarrow \theta_{\text{fp32}} - \eta \cdot \text{gradient}$ needs precision for small $\eta$ .

The ZeRO Optimization

ZeRO Stages: Progressive Partitioning

ZeRO-1: Optimizer State Partitioning

What: Each GPU stores only $\frac{1}{N}$ of the optimizer states.

How:

Partition parameters into $N$ groups: $\theta =[\theta_1, \theta_2, \ldots, \theta_N]$
GPU $r$ stores optimizer states only for partition $\theta_r$
After AllReduce of gradients, GPU $r$ updates only $\theta_r$ using its local optimizer states
AllGather updated parameters so all GPUs have the full model

Memory savings: $M_{\text{ZeRO-1}} = 2\Phi + 2\Phi + \frac{12\Phi}{N} = 4\Phi + \frac{12\Phi}{N}$

WHY does this work? Optimizer states (momentum, variance) for parameter $\theta_i$ only depend on gradient $\nabla \theta_i$ , not on other parameters. So GPU $r$ can independently update its assigned partition.

ZeRO-2: Gradient Partitioning

What: Additionally partition gradients—each GPU only stores $\frac{1}{N}$ of gradients.

How:

During backward pass, each GPU computes full local gradients
Use ReduceScatter instead of AllReduce: GPU $r$ receives only the averaged gradient for partition $\theta_r$
GPU $r$ updates $\theta_r$ , then AllGather

Memory savings: $M_{\text{ZeRO-2}} = 2\Phi + \frac{2\Phi}{N} + \frac{12\Phi}{N} = 2\Phi + \frac{14\Phi}{N}$

WHY ReduceScatter? It's AllReduce + partition in one operation. Instead of every GPU getting the full averaged gradient (AllReduce), each GPU gets only its assigned slice of the averaged gradient.

Worked example ZeRO-2 Communication Pattern

Setup: 4 GPUs, each computes gradients $g_r = [g_r^1, g_r^2, g_r^3, g_r^4]$ for 4 parameter partitions.

ReduceScatter operation:

GPU 0 receives: $\bar{g}^1 = \frac{1}{4}(g_0^1 + g_1^1 + g_2^1 + g_3^1)$
GPU 1 receives: $\bar{g}^2 = \frac{1}{4}(g_0^2 + g_1^2 + g_2^2 + g_3^2)$
GPU 2 receives: $\bar{g}^3 = \frac{1}{4}(g_0^3 + g_1^3 + g_2^3 + g_3^3)$
GPU 3 receives: $\bar{g}^4 = \frac{1}{4}(g_0^4 + g_1^4 + g_2^4 + g_3^4)$

Why this step? Each GPU only needs the averaged gradient for its parameter partition to perform the update. Storing all gradients would be wasteful.

ZeRO-3: Parameter Partitioning

What: The ultimate step—partition even the model parameters. Each GPU only stores $\frac{1}{N}$ of parameters.

How:

Forward pass: GPU $r$ $r$ needs layer $l$ $l$ 's parameters
- If $\theta_l$ is on GPU $r$ : use directly
- Otherwise: AllGather $\theta_l$ from owner GPU, compute, then discard
Backward pass: AllGather parameters again when needed for gradient computation
Update: Same as ZeRO-2 (ReduceScatter gradients, local update, AllGather updated params)

Memory savings: $M_{\text{ZeRO-3}} = \frac{2\Phi}{N} + \frac{2\Phi}{N} + \frac{12\Phi}{N} = \frac{16\Phi}{N}$

This is the $N\times$ reduction! All model state components are partitioned.

Formula ZeRO-3 Trade-off: Communication Overhead

Communication volume per training step: $C_{\text{ZeRO-3}} = 2 \times \underbrace{\frac{N-1}{N} \times 2\Phi}_{\text{AllGather params}} + \underbrace{\frac{N-1}{N} \times 2\Phi}_{\text{ReduceScatter grads}} = \frac{6\Phi(N-1)}{N}$

Compared to standard data parallelism: $C_{\text{standard}} = \frac{2\Phi(N-1)}{N} \text{ (AllReduce gradients only)}$

ZeRO-3 requires 3× more communication (factor of $\frac{6\Phi(N-1)/N}{2\Phi(N-1)/N} = 3$ ).

WHY? We AllGather parameters twice (forward + backward) and ReduceScatter gradients once, versus one AllReduce in standard DP. The trade-off: 3× communication for $N\times$ memory savings.

Worked example ZeRO-3 in Action: 175B Parameter Model

Given: $\Phi = 175 \times 10^9$ , $N = 64$ GPUs (8 nodes × 8 GPUs)

Memory per GPU: $M_{\text{gpu}} = \frac{16 \times 175 \times 10^9}{64} = 43.75 \text{ GB}$ Now feasible! Fits in an80 GB A100 with room for activations.

Communication per step: $C = \frac{6 \times 175 \times 10^9 \times 2 \times 63}{64} = 2.06 \text{ TB}$

With NVLink at ~300 GB/s per GPU, communication time $\approx 7$ seconds per step. If computation is 10seconds, we achieve $\frac{10}{10+7} \approx 59\%$ efficiency.

Why this step? This shows ZeRO-3 is practical only with fast interconnects. On slow networks (10 Gbps Ethernet ~1 GB/s), communication would dominate and destroy efficiency.

Advanced: ZeRO-Infinity and ZeRO-Offload

Memory hierarchy: $\text{GPU HBM} \xrightarrow[\text{5 TB/s}]{\text{AllGather}} \text{GPU HBM} \quad \text{or} \quad \text{CPU RAM} \xrightarrow[\text{30 GB/s}]{\text{PCIe}} \text{GPU} \quad \text{or} \quad \text{NVMe} \xrightarrow[\text{7 GB/s}]{\text{PCIe}} \text{CPU}$

Strategy: Keep only the currently-computed layer's parameters on GPU. Prefetch next layer's parameters from CPU while computing current layer.

Recall Explain to a 12-year-old

Imagine you and7 friends are building identical LEGO castles (that's your neural network). You each need the instruction manual (model parameters), and you're building different rooms (processing different data batches).

Standard way: Everyone has their own full copy of the 1000-page manual. When you each finish a room, you compare notes: "I think tower should be 10 blocks high," "I think 11," etc. You average the suggestions and everyone updates their manual to say "10.5 blocks." This works, but you've printed 8,000 pages total!

ZeRO way:

ZeRO-1: Each friend only keeps1/8 of the manual (125 pages) for making changes, but everyone still has a full read-only copy.
ZeRO-2: When sharing suggestions, each friend only writes down the suggestions for their 125-page section.
ZeRO-3: Each friend only has 125 pages, period. When you need to build a tower, you yell "Who has the tower pages?" and that friend reads them out loud to everyone. You build, then forget those instructions.

ZeRO-3 means you only printed 1,000 pages total (8× savings!), but now you spend time asking friends to read pages out loud. That's the trade-off: less paper (memory), more talking (communication).

Practical Considerations

When to Use Each ZeRO Stage

Worked example Decision Tree

Model size $\Phi$ , batch size $B$ , GPU count $N$ , GPU memory $M_{\text{avail}}$

Calculate memory needed:
- Standard DP: $M_{\text{need}} = 16\Phi + M_{\text{act}}(B)$
- ZeRO-1: $M_{\text{need}} = 4\Phi + \frac{12\Phi}{N} + M_{\text{act}}(B)$
- ZeRO-2: $M_{\text{need}} = 2\Phi + \frac{14\Phi}{N} + M_{\text{act}}(B)$
- ZeRO-3: $M_{\text{need}} = \frac{16\Phi}{N} + M_{\text{act}}(B)$
Choose minimum ZeRO stage where $M_{\text{need}}< M_{\text{avail}}$ Example: $\Phi = 20 \times 10^9$ (20B params), $N = 8$ A100 (80 GB each), $M_{\text{act}} \approx 10$ GB

Standard: $16 \times 20 + 10 = 330$ GB → doesn't fit
ZeRO-1: $4 \times 20 + \frac{12 \times 20}{8} + 10 = 80 + 30 + 10 = 120$ GB → doesn't fit
ZeRO-2: $2 \times 20 + \frac{14 \times 20}{8} + 10 = 40 + 35 + 10 = 85$ GB → doesn't fit (close!)
ZeRO-3: $\frac{16 \times 20}{8} + 10 = 40 + 10 = 50$ GB → fits!

Verdict: Use ZeRO-3.

Communication Efficiency Requirements

The efficiency of ZeRO-3 depends on: $\text{Efficiency} = \frac{T_{\text{compute}}}{T_{\text{compute}} + T_{\text{comm}}}$

where $T_{\text{comm}} = \frac{C_{\text{ZeRO-3}}}{\text{Bandwidth}}$

For 75% efficiency (reasonable target): $\frac{T_{\text{compute}}}{T_{\text{compute}} + T_{\text{comm}}} \geq 0.75 \implies T_{\text{comm}} \leq \frac{T_{\text{compute}}}{3}$

Worked example Bandwidth Requirements Check

Given: $\Phi = 175 \times 10^9$ , $N = 64$ , $T_{\text{compute}} = 10$ s per step

$\text{BW}_{\text{min}} = \frac{18 \times 175 \times 10^9 \times 2 \times 63}{64 \times 10} = \frac{396.9 \times 10^{12}}{640} = 620 \text{ GB/s}$

Per-GPU bandwidth: $\frac{620}{64} = 9.7$ GB/s Reality check:

NVLink (intra-node): ~300 GB/s → ✓ sufficient
InfiniBand HDR (inter-node): ~200 Gb/s = 25 GB/s → ✓ sufficient
10 Gbps Ethernet: 1.25 GB/s → ✗ insufficient (would need 78 s communication time!)

Why this step? This shows ZeRO-3 requires high-bandwidth interconnects. Cloud providers with fast networks (AWS p4d instances with EFA, Azure NDv4) are suitable; commodity clusters are not.

Connections

Model Parallelism: Complements ZeRO by splitting model layers across GPUs (vs. replicating)
Pipeline Parallelism: Another form of model parallelism; ZeRO can be combined with pipelines
Mixed Precision Training: ZeRO inherently uses mixed precision (fp16 compute, fp32 optimizer)
Gradient Accumulation: Can combine with ZeRO to handle larger effective batch sizes
Activation Checkpointing: Reduces $M_{\text{act}}$ term, making ZeRO-2 feasible instead of ZeRO-3
AllReduce and Collective Communication: Core primitives; ZeRO replaces AllReduce with ReduceScatter + AllGather
Large Language Models Training: ZeRO is essential for training GPT-3, LaMA, PaLM scale models
Memory-Efficient Optimizers: Adafactor, 8-bit Adam reduce the $12\Phi$ optimizer term

#flashcards/ai-ml

What is data parallelism in distributed training? :: A strategy where each GPU holds a complete copy of the model, processes different data batches in parallel, and synchronizes gradients using AllReduce to update all replicas identically.

What is the memory consumption formula for standard data parallelism with mixed precision and Adam optimizer?

M_{\text{gpu}} = 16\Phi

bytes where

\Phi

is the number of parameters. Breakdown: 2Φ params (fp16) + 2Φ gradients + 12Φ optimizer states (fp32 momentum, variance, master weights).

What does ZeRO stand for and what does it optimize?

Zero Redundancy Optimizer. It eliminates memory redundancy in data parallelism by partitioning model states (optimizer states, gradients, parameters) across GPUs instead of replicating them.

What are the three ZeRO stages and what does each partition?

ZeRO-1: partitions optimizer states only. ZeRO-2: partitions optimizer states + gradients. ZeRO-3: partitions optimizer states + gradients + parameters.

What is the memory per GPU formula for ZeRO-3 with N GPUs?

M_{\text{ZeRO-3}} = \frac{16\Phi}{N}

— all model state components (params, gradients, optimizer states) are partitioned N ways, achieving N× memory reduction.

What communication primitive does ZeRO-2 use instead of AllReduce and why?

ReduceScatter. It combines AllReduce with partitioning—each GPU receives only the averaged gradient for its assigned parameter partition, saving gradient memory while maintaining correctness.

What is the main trade-off when using ZeRO-3 compared to standard data parallelism? :: Memory vs. communication. ZeRO-3 achieves N× memory savings but requires 3× more communication volume (AllGather parameters twice per step + ReduceScatter gradients once, vs. one AllReduce in standard DP).

What is ZeRO-Infinity and when is it used?

An extension of ZeRO-3 that offloads partitioned model states to CPU RAM and NVMe SSD, using bandwidth-aware prefetching. Used for training trillion-parameter models that don't fit even with ZeRO-3 across available GPUs.

Why does ZeRO-3 require high-bandwidth interconnects to be efficient?

Because it must AllGather parameters during forward and backward passes. For 75% compute efficiency, network bandwidth must satisfy

\text{BW}_{\text{min}} = \frac{18\Phi(N-1)}{N \cdot T_{\text{compute}}}

. Low bandwidth causes communication overhead to dominate.

When should you NOT use ZeRO-3 despite having it available?

When your model fits in GPU memory with standard DP or ZeRO-2. ZeRO-3's 3× communication overhead makes training slower if memory isn't the bottleneck. Use the minimum ZeRO stage that fits your model.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Chalo, ek simple picture se samajhte hain. Jab hum ek bada neural network train karte hain, toh data parallelism ka matlab hai ki har GPU ke paas poore model ki ek complete copy hoti hai, aur har GPU alag-alag data batches process karti hai. Phir sab GPUs apne gradients ko aapas mein average kar leti hain (AllReduce), taaki sab ki copy ek jaisi rahe. Yeh kaam isliye karta hai kyunki poore dataset ka gradient basically har chhote subset ke gradients ka average hi hota hai — toh multiple GPUs milke faster training kar leti hain, aur averaging se noise bhi kam ho jaata hai.

Ab problem yeh hai ki har GPU par model ki complete copy rakhne se bahut zyada memory waste hoti hai. Jaise formula batata hai, ek model ke liye har GPU ko lagbhag 16Φ bytes chahiye — parameters, gradients, aur optimizer states (Adam ke momentum, variance, aur fp32 master weights) sab milaake. GPT-3 jaise 175 billion parameter model mein yeh 2800 GB per GPU ban jaata hai, jabki ek A100 GPU mein sirf 80 GB memory hoti hai! Aur agar 8 GPUs hain, toh hum 8 baar wahi same cheez store kar rahe hain — yani 8x redundancy, matlab bilkul zaroorat se zyada memory barbaad.

Yahin par ZeRO (Zero Redundancy Optimizer) kaam aata hai. Iska core idea simple hai: har GPU par saari cheezein duplicate karne ke bajaye, model state (optimizer states, gradients, parameters) ko sabhi GPUs mein baant do. Matlab har GPU sirf apna assigned portion rakhti hai, aur zaroorat padne par smart communication se ek doosre ki madad karti hai. Isse memory redundancy khatam ho jaati hai, aur yehi reason hai ki ZeRO ke saath 100-billion parameter jaise gigantic models ko limited hardware par bhi train kar paana possible ho jaata hai. Yeh matter isliye karta hai kyunki bina ZeRO ke aise bade models train karna practically impossible hota, aur aaj ke sabse powerful AI models isi tarah ki efficient scaling techniques par tike hain.

What is Data Parallelism?

How Standard Data Parallelism Works

Memory Consumption Analysis

The ZeRO Optimization

ZeRO Stages: Progressive Partitioning

ZeRO-1: Optimizer State Partitioning

ZeRO-2: Gradient Partitioning

ZeRO-3: Parameter Partitioning

Advanced: ZeRO-Infinity and ZeRO-Offload

Practical Considerations

When to Use Each ZeRO Stage

Communication Efficiency Requirements

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Scaling & Efficient Architectures

Connections