Data parallelism and ZeRO optimization — Hinglish notes

6.1.8 · AI-ML › Scaling & Efficient Architectures

Intuition Core Intuition

Socho ki ek bahut bada neural network train karna ek massive mural paint karne jaisa hai. Data parallelism matlab hai multiple painters hire karna, jisme har ek ke paas design ki apni copy hai (model parameters), aur wo wall ke alag-alag sections paint karte hain (alag data batches), phir apna kaam average karke master design update karte hain. Lekin problem ye hai: agar tumhara mural design 10 GB ka hai aur tumhare paas 8 painters hain, toh tum 80 GB identical designs store kar rahe ho! ZeRO (Zero Redundancy Optimizer) iska solution ye deta hai ki painters design ko smartly share karein—har painter sirf apna assigned portion rakhta hai, lekin phir bhi apna section paint kar sakta hai aur effectively collaborate kar sakta hai. Isi liye ZeRO 100-billion parameter models ko consumer hardware par train karna possible banata hai.

What is Data Parallelism?

How Standard Data Parallelism Works

Har GPU (rank $r$ ) par training loop:

Forward pass: Local model copy $θ_{r}$ se mini-batch $B_{r}$ process karo
Backward pass: Local gradients compute karo $\nabla_{θ_{r}} L (B_{r})$
AllReduce: Sabhi GPUs mein gradients synchronize karo: $\overset{ˉ}{\nabla}_{θ} = \frac{1}{N} \sum_{r = 1}^{N} \nabla_{θ_{r}} L (B_{r})$
Update: Optimizer apply karo: $θ \leftarrow θ - η \overset{ˉ}{\nabla}_{θ}$
Broadcast: Ab sabhi GPUs ke paas identical updated parameters hain

YE KAAM KYU KARTA HAI? Kyunki full dataset ka gradient, har subset ke gradients ka average hota hai (linearity of expectation). Har GPU true gradient ka ek noisy estimate compute karta hai, aur averaging se variance reduce hoti hai.

Memory Consumption Analysis

Formula Memory Per GPU in Standard Data Parallelism

$Φ$ parameters wale model ke liye (fp32 mein har parameter 4 bytes, fp16 mein 2 bytes): $M_{gpu} = params (fp16) 2Φ + gradients 2Φ + optimizer states 12Φ + batch-dependent M_{activations}$

Total model state: $16Φ$ bytes per GPU (mixed precision + Adam optimizer ke liye)

First principles se derivation:

Parameters (fp16): $2Φ$ bytes (compute efficiency ke liye half precision)
Gradients: $2Φ$ bytes (parameters ke same size, fp16)
Optimizer states (Adam):
- Momentum: $4Φ$ bytes (numerical stability ke liye fp32)
- Variance: $4Φ$ bytes (fp32)
- Master weights: $4Φ$ bytes (updates ke liye fp32 copy)
- Total: $12Φ$ bytes OPTIMIZER ke liye fp32 KYU? Mixed precision training speed ke liye fp16 use karta hai, lekin fp32 optimizer states maintain karta hai taaki tiny gradient updates mein underflow na ho. Update $θ_{fp32} \leftarrow θ_{fp32} - η \cdot gradient$ ko chhote $η$ ke liye precision chahiye.

The ZeRO Optimization

ZeRO Stages: Progressive Partitioning

ZeRO-1: Optimizer State Partitioning

Kya hai: Har GPU sirf $\frac{1}{N}$ optimizer states store karta hai.

Kaise:

Parameters ko $N$ groups mein partition karo: $θ = [θ_{1}, θ_{2}, \dots, θ_{N}]$
GPU $r$ sirf partition $θ_{r}$ ke liye optimizer states store karta hai
Gradients ke AllReduce ke baad, GPU $r$ sirf $θ_{r}$ ko apne local optimizer states use karke update karta hai
Sabhi GPUs ke paas full model ho iske liye updated parameters AllGather karo

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

YE KYU KAAM KARTA HAI? Parameter $θ_{i}$ ke optimizer states (momentum, variance) sirf gradient $\nabla θ_{i}$ par depend karte hain, doosre parameters par nahi. Isliye GPU $r$ apna assigned partition independently update kar sakta hai.

ZeRO-2: Gradient Partitioning

Kya hai: Additionally gradients bhi partition karo—har GPU sirf $\frac{1}{N}$ gradients store karta hai.

Kaise:

Backward pass ke dauran, har GPU full local gradients compute karta hai
AllReduce ki jagah ReduceScatter use karo: GPU $r$ ko sirf partition $θ_{r}$ ke liye averaged gradient milta hai
GPU $r$ $θ_{r}$ update karta hai, phir AllGather

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

ReduceScatter KYU? Ye AllReduce + partition ek hi operation mein hai. Iske bajaye ki har GPU ko full averaged gradient mile (AllReduce), har GPU ko averaged gradient ka sirf apna assigned slice milta hai.

ZeRO-3: Parameter Partitioning

Kya hai: Ultimate step—model parameters bhi partition karo. Har GPU sirf $\frac{1}{N}$ parameters store karta hai.

Kaise:

Forward pass: GPU $r$ ko layer $l$ ke parameters chahiye
- Agar $θ_{l}$ GPU $r$ par hai: directly use karo
- Warna: owner GPU se $θ_{l}$ AllGather karo, compute karo, phir discard karo
Backward pass: Gradient computation ke liye jab zaroorat ho tab parameters phir AllGather karo
Update: ZeRO-2 jaisa hi (ReduceScatter gradients, local update, AllGather updated params)

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

Ye hai $N \times$ reduction! Sabhi model state components partition hain.

Advanced: ZeRO-Infinity and ZeRO-Offload

Memory hierarchy: $GPU HBM AllGather 5 TB/s GPU HBM or CPU RAM PCIe 30 GB/s GPU or NVMe PCIe 7 GB/s CPU$

Strategy: GPU par sirf currently-computed layer ke parameters rakho. Current layer compute karte waqt next layer ke parameters CPU se prefetch karo.

Common mistake Common Misconception: "ZeRO training ko faster banata hai"

Galat soch: "Main training speed up karne ke liye ZeRO-3 use karunga!" Ye sahi kyun lagta hai: Zyada memory savings = better, hai na?

Sach: ZeRO memory ko communication ke badle trade karta hai. Ye training possible banata hai jab memory bottleneck ho, lekin generally 3× communication overhead ki wajah se training slow ho jaati hai. ZeRO-3 tab use karo jab tumhara model ZeRO-2 ya standard DP ke saath GPU memory mein fit na ho. Agar model fit hota hai, toh standard DP faster hai. Example: 1.3B parameter model ko 8 GPUs par train karna:

Standard DP memory: $16 \times 1.3 = 20.8$ GB → 80 GB A100 mein fit hota hai
ZeRO-3 memory: $\frac{20.8}{8} = 2.6$ GB → fit bhi hota hai, lekin 3× slow communication
Verdict: Standard DP use karo!

Recall Ek 12-saal ke bache ko explain karo

Socho tum aur 7 dost identical LEGO castles bana rahe ho (wo hai tumhara neural network). Tumhe sab ko instruction manual chahiye (model parameters), aur tum alag-alag rooms bana rahe ho (alag data batches process kar rahe ho).

Standard tarika: Har kisi ke paas 1000-page manual ki apni full copy hai. Jab tum sab ek room finish karte ho, notes compare karte ho: "Mujhe lagta hai tower 10 blocks oonchi honi chahiye," "Mujhe 11 lagti hai," etc. Tum suggestions average karte ho aur sabhi apna manual update karte hain "10.5 blocks." Ye kaam karta hai, lekin tumne total 8,000 pages print ki hain!

ZeRO tarika:

ZeRO-1: Har dost changes karne ke liye sirf manual ka 1/8 (125 pages) rakhta hai, lekin sabhi ke paas abhi bhi ek full read-only copy hai.
ZeRO-2: Suggestions share karte waqt, har dost sirf apne 125-page section ke liye suggestions likhta hai.
ZeRO-3: Har dost ke paas sirf 125 pages hain, bas. Jab tower banana ho, tum chillate ho "Kiske paas tower pages hain?" aur woh dost sabko unhe parhta hai. Tum banate ho, phir woh instructions bhool jaate ho.

ZeRO-3 ka matlab hai tumne total sirf 1,000 pages print ki (8× savings!), lekin ab dosto se pages padhwaane mein time lagta hai. Yahi trade-off hai: kam paper (memory), zyada baat (communication).

Practical Considerations

When to Use Each ZeRO Stage

Communication Efficiency Requirements

ZeRO-3 ki efficiency depend karti hai: $Efficiency = \frac{T _{compute}}{T _{compute} + T _{comm}}$

jahan $T_{comm} = \frac{C _{ZeRO-3}}{Bandwidth}$

75% efficiency ke liye (reasonable target): $\frac{T _{compute}}{T _{compute} + T _{comm}} \geq 0.75 ⟹ T_{comm} \leq \frac{T _{compute}}{3}$

Connections

Model Parallelism: ZeRO ka complement hai—model layers ko GPUs ke across split karta hai (replicate karne ki jagah)
Pipeline Parallelism: Model parallelism ka ek aur form; ZeRO ko pipelines ke saath combine kiya ja sakta hai
Mixed Precision Training: ZeRO inherently mixed precision use karta hai (fp16 compute, fp32 optimizer)
Gradient Accumulation: ZeRO ke saath combine karke larger effective batch sizes handle kar sakte hain
Activation Checkpointing: $M_{act}$ term reduce karta hai, jisse ZeRO-3 ki jagah ZeRO-2 feasible ban jaata hai
AllReduce and Collective Communication: Core primitives; ZeRO AllReduce ko ReduceScatter + AllGather se replace karta hai
Large Language Models Training: ZeRO GPT-3, LLaMA, PaLM scale models train karne ke liye essential hai
Memory-Efficient Optimizers: Adafactor, 8-bit Adam $12Φ$ optimizer term reduce karte hain

#flashcards/ai-ml

Distributed training mein data parallelism kya hota hai? :: Ek strategy jisme har GPU ke paas model ki complete copy hoti hai, wo parallel mein alag data batches process karta hai, aur sabhi replicas ko identically update karne ke liye AllReduce se gradients synchronize karta hai.

Standard data parallelism mein mixed precision aur Adam optimizer ke saath memory consumption ka formula kya hai?

M_{gpu} = 16Φ

bytes jahan

Φ

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

ZeRO ka full form kya hai aur ye kya optimize karta hai?

Zero Redundancy Optimizer. Ye data parallelism mein memory redundancy eliminate karta hai—model states (optimizer states, gradients, parameters) ko replicate karne ki jagah GPUs ke across partition karta hai.

Teen ZeRO stages kya hain aur har ek kya partition karta hai?

ZeRO-1: sirf optimizer states partition karta hai. ZeRO-2: optimizer states + gradients partition karta hai. ZeRO-3: optimizer states + gradients + parameters partition karta hai.

N GPUs ke saath ZeRO-3 mein memory per GPU ka formula kya hai?

M_{ZeRO-3} = \frac{16Φ}{N}

— sabhi model state components (params, gradients, optimizer states) N ways partition hote hain, N× memory reduction achieve hota hai.

ZeRO-2 AllReduce ki jagah kaun sa communication primitive use karta hai aur kyun?

ReduceScatter. Ye AllReduce aur partitioning ko ek hi operation mein combine karta hai—har GPU ko sirf apne assigned parameter partition ke liye averaged gradient milta hai, gradient memory bachata hai aur correctness maintain karta hai.

Standard data parallelism se compare karte hue ZeRO-3 use karne mein main trade-off kya hai? :: Memory vs. communication. ZeRO-3 N× memory savings achieve karta hai lekin 3× zyada communication volume chahiye (AllGather parameters do baar per step + ReduceScatter gradients ek baar, vs. standard DP mein ek AllReduce).

ZeRO-Infinity kya hai aur kab use hota hai?

ZeRO-3 ka ek extension jo partitioned model states ko CPU RAM aur NVMe SSD par offload karta hai, bandwidth-aware prefetching use karta hai. Tab use hota hai jab trillion-parameter models available GPUs ke saath ZeRO-3 mein bhi fit nahi hote.

ZeRO-3 ko efficient rehne ke liye high-bandwidth interconnects kyun chahiye?

Kyunki isse forward aur backward passes ke dauran parameters AllGather karne padte hain. 75% compute efficiency ke liye, network bandwidth satisfy karna chahiye

BW_{min} = \frac{18Φ ( N - 1 )}{N \cdot T _{compute}}

. Low bandwidth se communication overhead dominant ho jaata hai.

ZeRO-3 available hone ke bawajood kab use NAHI karna chahiye?

Jab tumhara model standard DP ya ZeRO-2 ke saath GPU memory mein fit ho jaata ho. ZeRO-3 ka 3× communication overhead training slow kar deta hai agar memory bottleneck nahi hai. Minimum ZeRO stage use karo jo tumhara model fit kare.