Model parallelism (tensor, pipeline) — Hinglish notes

6.1.7 · AI-ML › Scaling & Efficient Architectures

Jab models itne bade ho jaate hain ki ek single GPU ki memory mein fit nahi hote, tab model parallelism model ko hi multiple devices mein split kar deta hai — unlike data parallelism jo poore model ko replicate karta hai.

Connections

Data Parallelism - complementary approach (replicate model, split data)
Gradient Accumulation - batch size ke liye memory kam karta hai
Mixed Precision Training - har parameter ke liye memory kam karta hai
Transformer Architecture - model parallelism ka common target
GPU Memory Hierarchy - hardware constraint jo zaroorat drive karta hai
Activation Checkpointing - compute aur memory ka trade karta hai

Model parallelism kehta hai: "Agar model ek device par fit nahi hota, toh model ko tukdon mein kaat do aur har tukda alag device par rakh do."

Do main strategies hain:

Tensor Parallelism: Individual layers ko split karo (matrices ko horizontally/vertically kaatna)
Pipeline Parallelism: Model ko depth-wise split karo (early layers GPU-0 par, later layers GPU-1 par)

Tensor Parallelism

Derivation: Matrix Multiply Ko Kaise Split Karein

Ek simple linear layer consider karo: $Y = X W$ jahan:

$X \in R^{b \times d_{in}}$ (batch × input dimension)
$W \in R^{d_{in} \times d_{o u t}}$ (weight matrix)
$Y \in R^{b \times d_{o u t}}$ (output)

Hum ise split kyun kar sakte hain? Matrix multiplication concatenation ke saath associative hai.

Column-wise Split (sabse zyada common)

$W$ ko $N$ column slices mein split karo: $W = [W_{1} ∣ W_{2} ∣ \dots ∣ W_{N}]$ jahan har $W_{i} \in R^{d_{in} \times (d_{o u t} / N)}$ .

Derivation: $Y = X W = X [W_{1} ∣ W_{2} ∣ \dots ∣ W_{N}] = [X W_{1} ∣ X W_{2} ∣ \dots ∣ X W_{N}]$

Toh device $i$ compute karta hai $Y_{i} = X W_{i}$ , phir hum concatenate karte hain: $Y = [Y_{1} ∣ Y_{2} ∣ \dots ∣ Y_{N}]$ .

Key insight: Har device ko poora input $X$ chahiye lekin $W$ ka sirf ek fraction. Weights par memory save hoti hai, lekin $X$ ko sab devices par broadcast karna padta hai.

Yeh step kyun? Column split isliye kaam karta hai kyunki matrix-vector multiplication distribute hoti hai: $W$ ka har column block ek independent column block of $Y$ produce karta hai.

Backward Pass Ka Kya?

Gradients flow karte hain: $\frac{\partial L}{\partial W} = X^{T} \frac{\partial L}{\partial Y}$

Agar $Y$ column-wise split hai, toh $\frac{\partial L}{\partial Y}$ bhi split hoga. Device $i$ compute karta hai: $\frac{\partial L}{\partial W _{i}} = X^{T} \frac{\partial L}{\partial Y _{i}}$

Critical: $\frac{\partial L}{\partial X}$ compute karne ke liye (earlier layers ke liye zaroori): $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} W^{T} = \sum_{i = 1}^{N} \frac{\partial L}{\partial Y _{i}} W_{i}^{T}$

Iske liye sab devices se contributions sum karne ke liye ek all-reduce chahiye.

Yeh step kyun? Chain rule demand karta hai ki inputs ke respect mein gradients paane ke liye hum saare output dimensions par sum karein. Kyunki outputs split hain, aggregate karne ke liye communicate karna padta hai.

Ek Transformer MLP mein do linear layers hoti hain: $MLP (x) = GELU (x W_{1}) W_{2}$

jahan $W_{1} : d_{m o d e l} \to 4 d_{m o d e l}$ (expand) aur $W_{2} : 4 d_{m o d e l} \to d_{m o d e l}$ (project).

$W_{1}$ par column-parallel, $W_{2}$ par row-parallel:

Device $i$ of $N$ : $W_{1}^{(i)} \in R^{d_{m o d e l} \times (4 d_{m o d e l} / N)}$ store karo
Locally compute karo: $h_{i} = GELU (x W_{1}^{(i)})$
Abhi koi communication nahi! Har device ke paas hidden state ka ek slice hai
$W_{2}^{(i)} \in R^{(4 d_{m o d e l} / N) \times d_{m o d e l}}$ store karo (row slice)
Compute karo: $y_{i} = h_{i} W_{2}^{(i)}$
All-reduce: $y = \sum_{i = 1}^{N} y_{i}$

Yeh pattern kyun? Column-parallel ke baad row-parallel design communication minimize karta hai: poore MLP block mein sirf ek all-reduce, har layer mein ek nahi.

Memory saved: $W_{1}$ aur $W_{2}$ dono $N$ se split hain, toh roughly $1/ N$ MLP parameters har device par.

Yeh step kyun? $W_{1}$ ke baad activation ek intermediate hai jise gather karne ki zaroorat nahi. Hum sirf $W_{2}$ ke baad final output gather karte hain, communication overhead kam karta hai.

Attention mein projection matrices $W_{Q}, W_{K}, W_{V} \in R^{d_{m o d e l} \times d_{m o d e l}}$ aur output projection $W_{O}$ hoti hai.

Attention heads mein split karo: Agar 8 heads hain, toh 4 devices mein se har ek par 2 heads rakho.

Device $i$ compute karta hai: $head_{2 i}, head_{2 i + 1} = Attention (Q_{i}, K_{i}, V_{i})$

jahan $Q_{i} = X W_{Q}^{(i)}$ , etc., $W_{Q}^{(i)} \in R^{d_{m o d e l} \times (d_{m o d e l} /4)}$ ke saath.

Phir concatenate karo aur $W_{O}$ apply karo (row-parallel), followed by all-reduce.

Yeh step kyun? Attention heads independent hain, isliye naturally parallelize hote hain. Har device ek subset of heads compute karta hai, memory aur compute per device kam hoti hai.

Pipeline Parallelism

Naive Pipeline: Bubble Problem

Setup: 8 layers wala model, 4 GPUs. GPU-0 ke paas layers 0-1 hain, GPU-1 ke paas layers 2-3, etc.

Forward pass:

GPU-0 batch process karta hai → activations GPU-1 ko bhejta hai
GPU-1 process karta hai → GPU-2 ko bhejta hai
etc.

Problem: GPU-1 idle baithta hai jab GPU-0 compute kar raha hota hai. GPU-2 idle rehta hai jab GPU-0 aur GPU-1 kaam kar rahe hote hain. Isse unused compute time ka ek "bubble" banta hai.

Utilization: $D$ devices aur $B$ micro-batches ke saath, bubble time approximately $\frac{D - 1}{B + D - 1}$ total time ka hota hai. $D = 4, B = 1$ ke liye, yeh 75% idle hai!

Yeh kyun hota hai: Sequential dependency. Later stages earlier stages ka wait karte hain. Yeh pipeline parallelism ka fundamental trade-off hai.

GPipe: Micro-Batching Solution

Timeline:

GPU-0: [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]
GPU-1:      [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]
GPU-2:           [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]
GPU-3:                [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]

$F_{i}$ = forward micro-batch $i$ , $B_{i}$ = backward micro-batch $i$

Bubble fraction: $\frac{D - 1}{M + D - 1}$ jahan $M$ = number of micro-batches.

$M = 16, D = 4$ ke saath: bubble = $3/19 \approx 16%$ . Bahut better!

Cost: Backward pass ke liye saare $M$ micro-batches ke activations ek saath store karne padte hain. Memory linearly $M$ ke saath badhti hai.

Yeh step kyun? Zyada micro-batches later stages ko busy rakhte hain jab early stages abhi bhi process kar rahe hote hain. Pipeline fill ho jaati hai, idle time kam hoti hai, activation memory ki cost par.

Model: 48 layers, 4 GPUs (12 layers each), micro-batch size 1, total batch 32.

Scenario 1: $M = 8$ micro-batches

Bubble time: $(4 - 1) / (8 + 4 - 1) = 27%$
Activation memory: 8× single micro-batch
Throughput: ~5.5 samples/sec

Scenario 2: $M = 32$ micro-batches

Bubble time: $(4 - 1) / (32 + 4 - 1) = 8.6%$
Activation memory: 32× single micro-batch
Throughput: ~6.8 samples/sec

Trade-off: Zyada $M$ efficiency improve karta hai lekin memory badhta hai. GPipe activation checkpointing use karta hai (backward ke dauraan activations recompute karna) memory kam karne ke liye, 33% zyada compute ki cost par.

Yeh example kyun? Pipeline depth ka quantitative impact dikhata hai. Real systems ko available memory aur desired throughput ke based par $M$ balance karna hota hai.

PipeDream: 1F1B Schedule

GPipe ki problem: Pehle saare forward passes, phir saare backward passes. Activations ke liye memory spike.

PipeDream solution: Forward aur backward ko interleave karo: "1 Forward, 1 Backward" schedule.

Timeline:

GPU-0: [F₀] [F₁] [F₂] [F₃] [B₀] [F₄] [B₁] [F₅] [B₂] [F₆] [B₃] ...
GPU-1:      [F₀] [F₁] [F₂] [B₀] [F₃] [B₁] [F₄] [B₂] [F₅] [B₃] ...

Yeh kyun help karta hai: Early micro-batches ke backward passes later forward passes se pehle hote hain. Activations jaldi free ho jaate hain. Steady-state memory: sirf $D$ micro-batches' activations (pipeline stages ki sankhya), $M$ nahi (total micro-batches).

Complication: Earlier layers older micro-batches ke gradients dekhte hain (version inconsistency). PipeDream weight stashing use karta hai: training ke dauraan multiple weight versions rakhna taaki har micro-batch ke forward aur backward same weights use karein.

Yeh step kyun? Forward aur backward passes overlap karke, peak memory $O (M)$ se $O (D)$ ho jaati hai, jo deeper pipelines ko memory khatam kiye bina feasible banata hai.

Yeh sahi kyun lagta hai: Pipeline simple lagti hai — bas layers split karo, koi complex matrix slicing nahi.

Yeh galat kyun hai: Pipeline mein unavoidable bubble time hai aur har layer partition ke beech inter-device communication hai. Tensor parallelism mein har layer ke andar communication hoti hai lekin saare devices par simultaneously batches process kar sakta hai (steady state mein koi bubble nahi).

Kab kaun jeetat hai:

Tensor parallelism: Wide models (badi hidden dimensions) ke liye better, high inter-device bandwidth (NVLink) ke saath. Communication cost hidden size ke proportional hai, depth ke nahi.
Pipeline parallelism: Deep models ke liye better jab bandwidth limited ho. Communication sirf partition boundaries par hoti hai (forward pass mein kam baar).

Reality: Modern systems (Megatron-LM, DeepSpeed) dono combine karte hain. Tensor parallelism ek node ke andar (fast NVLink) aur pipeline parallelism nodes ke across (slower InfiniBand).

Fix: Apna model aur hardware profile karo. Agar communication time < bubble time, toh tensor prefer karo. Agar memory per layer bottleneck hai, toh pipeline prefer karo.

Yeh sahi kyun lagta hai: Chote batches → kam activation memory → model ek GPU par fit ho jaata hai.

Yeh galat kyun hai: Model parameters aur optimizer states batch size se independent hain. Ek 175B parameter model ko sirf weights ke liye 700GB chahiye (FP32). Batch size 1 karne se yeh nahi badalta.

Confusion yeh hai: Activation memory (forward pass intermediate results) batch size ke saath scale karti hai. Lekin parameter memory bade models ke liye dominate karti hai.

Numbers:

Parameters: 175B × 4 bytes = 700GB (fixed)
Optimizer states (Adam): 175B × 12 bytes = 2.1TB (fixed)
Activations: ~batch_size × seq_len × hidden_dim × num_layers (variable)

GPT-3 ke liye, batch size 1 ke saath bhi, parameter + optimizer memory kisi bhi single GPU se zyada hai.

Fix: Model parallelism tab zaroori hai jab parameter count device memory se zyada ho. Activation memory reduction (gradient accumulation, checkpointing) complementary hai.

Communication Patterns

Tensor Parallelism (per layer): $Comm_{tensor} = 2 b s h$

jahan $b$ = batch size, $s$ = sequence length, $h$ = hidden dimension.

Factor of 2: forward ke baad all-gather (activations broadcast karne ke liye) + backward ke baad reduce-scatter (gradients aggregate karne ke liye).

Pipeline Parallelism (per stage boundary): $Comm_{pipeline} = b s h$

Sirf forward activation next stage ko bheji jaati hai aur backward gradient previous stage ko. Pipeline ke andar koi all-reduce nahi.

Yeh kyun matter karta hai: Tensor parallelism mein, communication hidden size ke saath badhti hai. Pipeline ke liye, yeh har partition boundary par constant hai. Yeh hybrid strategy guide karta hai.

Goal: 530B parameter model train karna (Megatron-Turing NLG).

Hardware:

4480 GPUs (A100 80GB)
8 GPUs wale nodes NVLink se connected (900 GB/s)
Nodes InfiniBand se connected (200Gb/s)

Configuration:

Tensor parallelism: $N = 8$ har node ke andar (fast NVLink all-reduce sasta banata hai)
Pipeline parallelism: $P = 35$ stages nodes ke across (cross-node communication kam karta hai)
Data parallelism: 16-way pipeline replicas ke across

Yeh kyun kaam karta hai:

Tensor parallelism fast intra-node bandwidth exploit karta hai
Pipeline parallelism slow inter-node bandwidth use limit karta hai
Data parallelism poore cluster mein training scale karta hai

Memory per GPU:

Model: 530B / (8 × 35) = 1.9B parameters per GPU ≈ 7.6GB (FP32)
Optimizer states ke saath: ~23GB per GPU
Activations aur workspace ke liye jagah bachti hai

Yeh step kyun? Yeh ek real-world example hai ki parallelism strategies kaise compose hoti hain. Parallelism ka har dimension ek alag bottleneck address karta hai.

Recall Ek 12 saal ke bachche ko explain karo

Socho tum duniya ka sabse bada LEGO castle banana chahte ho, lekin yeh itna bada hai ki saare pieces tumhari table par fit nahi honge. Tum kya karoge?

Option 1 (Tensor Parallelism): Tum aur tumhara dost same wall ke pieces aadhe-aadhe le lo. Tum dono apna aadha banao, phir snap karke jodo. Agli wall ke liye, pieces phir se split karo. Tum dono har part par saath kaam karte ho, lekin load share karte ho.

Option 2 (Pipeline Parallelism): Tum pehli manzil banao, phir apne dost ko do jo doosri manzil banaye, jabki tum ek naye castle ki pehli manzil shuru karo. Jaise ek assembly line — sab ka apna stage hai, aur castles flow karte hain.

Problem? Option 1 mein, tum dono ko instruction book dekhni padti hai (communication). Option 2 mein, tumhara dost wait karta hai jab tum finish karo (bubble time).

Real AI systems dono karte hain! Fast splits (tensor) apni table par dosto ke saath, assembly line (pipeline) kamre mein tables ke beech.

TENSOR = Together Each Node Slices One Row → Saare devices same layer par kaam karte hain, horizontally split

PIPELINE = Pass Inputs Progressively Each Layer In Next Element → Data layers ke through pipe ki tarah flow karta hai

Kab use karein: "Tight bandwidth → Tensor; Poor links → Pipeline"

Flashcards

Data parallelism aur model parallelism mein fundamental difference kya hai? :: Data parallelism poore model ko har device par replicate karta hai aur data batch split karta hai. Model parallelism model ko hi devices mein split karta hai kyunki woh ek device par fit nahi hota.

Tensor parallelism kya hai?

Individual weight matrices ko ek layer ke andar multiple devices mein split karna, jahan har device matrix ke ek slice par compute karta hai (intra-layer parallelism).

Pipeline parallelism kya hai? :: Model ko depth-wise devices mein split karna, jahan har device alag layers rakhta hai aur data sequentially flow karta hai jaise ek pipeline (inter-layer parallelism).

Tensor parallelism mein

N

devices mein matrix

W

ke column-wise split ke liye forward pass mein kya communication chahiye?

Ek all-gather operation jo output pieces

[Y_{1} ∣ Y_{2} ∣ \dots ∣ Y_{N}]

ko concatenate kare after har device

Y_{i} = X W_{i}

compute kare.

Tensor parallelism mein

W

ke column-wise split ke liye backward pass mein kya communication chahiye?

Ek all-reduce operation jo gradient contributions

\sum_{i} \frac{\partial L}{\partial Y _{i}} W_{i}^{T}

sum kare

\frac{\partial L}{\partial X}

compute karne ke liye.

Naive pipeline parallelism mein "bubble problem" kya hai?

Later stages idle baithte hain jab earlier stages process karte hain, sequential dependencies ki wajah se wasted compute time hota hai. Bubble fraction approximately

(D - 1) / (B + D - 1)

hoti hai

D

devices aur

B

micro-batches ke liye.

GPipe pipeline bubbles kaise kam karta hai?

Har mini-batch ko bahut saare micro-batches mein split karke, later stages ko busy rakhta hai jab early stages subsequent micro-batches process karte hain. Bubble time

O (1)

(D - 1) / (M + D - 1)

ho jaata hai

M

micro-batches ke liye.

GPipe ke micro-batching ki memory cost kya hai?

Backward pass ke liye saare

M

micro-batches ke activations simultaneously store karne padte hain, memory linearly

M

ke saath badhti hai.

PipeDream mein 1F1B schedule kya hai?

"1 Forward 1 Backward" — forward aur backward passes interleave karna taaki activations jaldi free hon. Steady-state memory

O (D)

(stages ki sankhya) hoti hai

O (M)

(total micro-batches) ki jagah.

Tensor parallelism ko pipeline parallelism se zyada communication kyun chahiye?

Tensor parallelism ko har layer ke andar all-gather aur all-reduce operations chahiye (communication volume

2 b s h

per layer). Pipeline sirf stage boundaries ke beech activations bhejta hai (volume

b s h

per boundary).

Tensor parallelism pipeline parallelism se kab prefer kiya jaata hai?

Jab model wide ho (badi hidden dimension), inter-device bandwidth high ho (NVLink), aur communication cost pipeline bubble time se kam ho.

Pipeline parallelism tensor parallelism se kab prefer kiya jaata hai?

Jab model deep ho, inter-device bandwidth limited ho, aur memory per layer bottleneck ho. Pipeline ke communication points kam hote hain (sirf partition boundaries par).

Megatron-LM jaise modern systems hybrid parallelism kyun use karte hain?

Interconnect bandwidth ke alag levels exploit karne ke liye: tensor parallelism nodes ke andar (fast NVLink) aur pipeline parallelism nodes ke across (slower InfiniBand), plus data parallelism poore cluster mein scaling ke liye.

Kya batch size kam karna model parallelism ki zaroorat khatam kar sakta hai?

Nahi. Model parameters aur optimizer states batch size se independent hain. Ek 175B model ko parameters aur optimizer states ke liye ~2TB chahiye batch size se regardless. Model parallelism tab zaroori hai jab parameters device memory se zyada hon.