Distributed training overview

3.3.11 · AI-ML › Deep Learning Frameworks

Core Idea

Ise ek restaurant kitchen ki tarah socho: ek chef (GPU) ek ghante mein sirf itne hi dishes bana sakta hai. 10 chefs milke kaam karein, toh customers ko 10× faster serve karte ho—agar woh sahi se coordinate karein aur ek doosre ke raaste mein na aayein.

The Fundamental Problem

Hum Distribute Kya Kar Rahe Hain?

Jab koi model train karte hain, teen resources dominant hote hain:

1. Computation: Forward pass ($\mathbf{y}=f(\mathbf{x};\theta )$) aur backward pass (${\nabla }_{\theta }\mathcal{L}$)
2. Memory: Parameters $\theta$, activations, gradients, optimizer states store karna
3. Data: Training dataset $\mathcal{D}=\{({\mathbf{x}}_i,y_i){\}}_{i=1}^N$

Main Parallelization Strategies

1. Data Parallelism

Workflow:

1. Mini-batch $B$ of size $b$ ko $k$ sub-batches mein split karo: $B_1,B_2,\dots ,B_k$ (ek per device)
2. Har device apne sub-batch par gradients compute karta hai: ${\mathbf{g}}_i=\frac{1}{|B_i|}{\sum }_{(\mathbf{x},y)\in B_i}{\nabla }_{\theta }\mathcal{L}(f(\mathbf{x};\theta ),y)$
3. Devices ke beech gradients synchronize karo (all-reduce operation): ${\mathbf{g}}_{\text{global}}=\frac{1}{k}{\sum }_{i=1}^k{\mathbf{g}}_i$
4. Har device apni local copy ${\mathbf{g}}_{\text{global}}$ se update karta hai: $\theta \leftarrow \theta -\eta {\mathbf{g}}_{\text{global}}$

Ye kyun kaam karta hai: Mini-batches ke gradients ka average, full batch ke gradient ke barabar hota hai (expectation ki linearity). Mathematically ye ek single device par size $b\times k$ ke batch par training ke equivalent hai.

Ye step kyun? Synchronizing ensure karta hai ki har update ke baad sab devices ke parameters identical hain. Sync ke bina, devices diverge ho jaate aur alag-alag models produce karte.

Result: 4× faster training (ideally), lekin communication overhead actual speedup ko ~3.5× tak reduce kar deta hai.

2. Model Parallelism

Do types:

A) Tensor Parallelism (Intra-layer splitting) Individual layers ko devices mein split karo. Ek badi matrix multiplication $\mathbf{Y}=\mathbf{XW}$ ke liye jahan $\mathbf{W}\in {\mathbb{R}}^{d\times h}$:

$\mathbf{W}=\left[\begin{array}{c}{\mathbf{W}}_1\\ {\mathbf{W}}_2\end{array}\right],\mathbf{Y}=\mathbf{X}\left[\begin{array}{c}{\mathbf{W}}_1\\ {\mathbf{W}}_2\end{array}\right]=\left[\begin{array}{cc}{\mathbf{Y}}_1& {\mathbf{Y}}_2\end{array}\right]$

Device 1 compute karta hai ${\mathbf{Y}}_1=\mathbf{X}{\mathbf{W}}_1$, Device 2 compute karta hai ${\mathbf{Y}}_2=\mathbf{X}{\mathbf{W}}_2$ parallel mein. Phir results concatenate karo.

Ye step kyun? Jab $\mathbf{W}$ ek GPU ki memory mein fit nahi hota (jaise ek 50GB weight matrix), tab use split karna hi ek tarika hai model fit karne ka.

B) Pipeline Parallelism (Inter-layer splitting) Model ko vertically layers ke hisaab se split karo. Device 1 ke paas layers 1-10 hain, Device 2 ke paas layers 11-20, aur aage bhi.

Data ek assembly line ki tarah flow karta hai:

• Batch 1 Device 1 mein enter karta hai (layers 1-10)
• Device 1 ka output → Device 2 (layers 11-20)
• Jab Device 2 Batch 1 process kar raha hota hai, Device 1 Batch 2 shuru kar deta hai

Challenge: Pipeline bubbles—devices idle baithte hain, pichle stage se data ka wait karte hue.

Solution: 8 GPUs mein split karo:

• GPU 0: Layers 0-11
• GPU 1: Layers 12-23
• ...
• GPU 7: Layers 84-95

Ek forward pass mein 7 device-to-device transfers chahiye (expensive!). Lekin model run karne ka yahi ek tarika hai.

3. Hybrid Parallelism

Example: 64 GPUs (8 nodes × 8 GPUs) par ek 100B parameter model train karna:

• Model ko 8 GPUs mein split karo (model parallelism)
• Ise 8 baar replicate karo (nodes ke across data parallelism)
• Total: 8 data-parallel replicas, har ek 8-way model parallelism use karta hai

Key Trade-offs

Communication Overhead

Actual speedup: $S=\frac{T_{\text{compute}}}{T_{\text{compute}}/k+T_{\text{comm}}}$

jahan $T_{\text{comm}}$ gradients/activations synchronize karne mein lagta waqt hai.

Ye kyun matter karta hai: Agar $T_{\text{comm}}\approx T_{\text{compute}}/k$, toh speedup sirf $0.5k$ hai (50% efficiency). High-bandwidth interconnects (NVLink, InfiniBand) $T_{\text{comm}}$ reduce karte hain.

Scaling Efficiency

• Linear scaling: Efficiency = 1 (rare, sirf minimal communication ke saath)
• Sub-linear scaling: Efficiency = 0.7-0.9 (well-optimized systems ke liye typical)
• Super-linear scaling: Efficiency > 1 (possible jab bada batch cache mein fit ho)

Efficiency kyun degrade hoti hai: Jaise-jaise $k$ badhta hai, har device kam kaam karta hai ($T_{\text{compute}}/k$ shrink hota hai), lekin $T_{\text{comm}}$ constant rehta hai ya badhta hai (zyada devices sync karne hain).

Common Implementation Patterns

Synchronous vs. Asynchronous Training

Pros: Mathematically single-device training ke equivalent (deterministic) Cons: Speed slowest device se limit hoti hai (stragglers performance hurt karte hain)

Pros: Stragglers ka wait nahi, better hardware utilization Cons: Stale gradients (worker ${\theta }_t$ use karta hai lekin ${\theta }_{t+5}$ update karta hai), convergence hurt ho sakti hai

Gradient Accumulation

1. $k=B/b$ micro-batches ke upar gradients accumulate karo: ${\mathbf{g}}_{\text{accum}}=\frac{1}{k}{\sum }_{i=1}^k{\mathbf{g}}_i$
2. $k$ forward-backward passes ke baad update karo: $\theta \leftarrow \theta -\eta {\mathbf{g}}_{\text{accum}}$

Ye kyun kaam karta hai: Gradients linear hain—micro-batches ka average, full batch ke gradient ke barabar hota hai. Ye size $B$ ke batch process karne ke equivalent hai, bas slower hai.

Mistakes & Misconceptions

Reality: Communication overhead device count ke saath badhta hai. Ek certain point ke baad ("scaling wall"), devices add karne se diminishing returns milte hain ya cheezein slow bhi ho jaati hain.

Example: Ek chhote CNN ko 100 GPUs par train karna slower ho sakta hai 8 GPUs ke comparison mein kyunki gradient sync time computation ko dominate kar leta hai.

Fix: Parallelism ko model size ke hisaab se match karo. Chhote models: data parallelism ke saath 1-8 GPUs. Bade models: hybrid strategies ke saath 100+ GPUs.

Difference:

• Large batch on 1 GPU: 256 samples process karo → update
• Data parallel on 4 GPUs (batch 64 each): Har GPU par 64 samples process karo → sync update (effective batch 256)

Math mein dono equivalent hain lekin implementation mein alag. Data parallelism communication cost add karta hai; ek GPU par large batches memory pressure add karte hain.

Nuance: Model parallelism har layer mein cross-device communication add karta hai (high latency). Data parallelism sirf backward pass ke end mein sync karta hai.

Better strategy: Jitna ho sake data parallelism use karo, model parallelism sirf tab jab zaroor ho. For example, ek 10B parameter model shayad activation checkpointing ke saath ek 40GB GPU par fit ho jaaye—model parallelism se bachte raho.

Practical Considerations

Strategy Kaise Chunein

Framework Support

• PyTorch: DistributedDataParallel (data), Tensor Parallelism (via Megatron-LM), PipeDream (pipeline)
• TensorFlow: tf.distribute.MirroredStrategy (data), tf.distribute.TPUStrategy (TPUs par data)
• JAX: pmap (data), xmap (arbitrary sharding)
• DeepSpeed: ZeRO stages (optimizer sharding), pipeline parallelism
• Horovod: MPI-based data parallelism (framework-agnostic)

Active Recall Practice

Distributed training AI models ke liye aisa hi hai. Ek bade AI brain (jaise ChatGPT) ko ek computer par train karna bahut time leta hai. Toh hum kaam ko bahut saare computers (GPUs) mein split karte hain. Har computer alag pictures ya sentences par train karta hai, kuch seekhta hai, aur phir sab computers share karte hain jo unhone seekha. Sab ke answers combine karke, AI faster seekhta hai.

Tricky part? Computers ko answers share karne ke liye ek doosre se baat karni padti hai (jaise tum homework answers compare karte ho), aur us baat karne mein waqt lagta hai. Agar woh zyada waqt baat karne mein lagaayein, toh akele kaam karne se bhi slow ho jaata hai!

• Model parallelism: Model ko devices mein split karo

• Overhead: Communication cost devices ke saath badhti hai

• Devices: Zyada hamesha better nahi (scaling wall)

• Efficiency: Actual speedup / ideal speedup

• Layers: Pipeline layers ke hisaab se split karta hai

• Hybrid: Data + model parallelism combine karta hai

• Yield: 100B+ parameter models ke liye best

• Batch: Gradient accumulation large batches simulate karta hai

• Ring all-reduce: Efficient gradient sync topology

• Interconnect: NVLink/InfiniBand model parallelism ke liye critical

• Deterministic: Synchronous SGD single-GPU se match karta hai

Connections

• 3.1-Neural-network-basics - forward/backward passes jo parallelize ho rahe hain
• 3.3.5-Batch-normalization - data parallelism mein batch statistics
• 3.9-Gradient-descent-optimizers - optimizer state sharding (ZeRO)
• 3.4.1-Transformers-overview - GPT-3 ko model parallelism kyun chahiye
• 4.2-GPU-acceleration - hardware jo distribution enable karta hai
• 5.1-Training-at-scale - distributed systems ki practical engineering

#flashcards/ai-ml

What is the main goal of distributed training? :: Training speed up karna aur single-device memory se bade models enable karna, computation ko multiple GPUs/TPUs mein parallelize karke.

Concept Map