Load balancing in MoE — Hinglish notes

6.1.6 · AI-ML › Scaling & Efficient Architectures

MoE Mein Load Balancing Kya Hai?

Iski zaroorat kyun hai?

Balancing ke bina: Gating network (router) ek degenerate solution par converge kar sakta hai jahan wo zyaadatar tokens ek chote subset of experts ko bhejta hai. Baaki experts ko almost koi gradient signal nahi milta aur wo seekh nahi paate.
Balancing ke saath: Saare experts diverse data par train hote hain, specialized skills develop karte hain, aur model apna poora parameter budget efficiently use karta hai.

Load Balancing Problem: First Principles Se

Step 1: Routing Kaise Kaam Karta Hai

MoE gating function yaad karo. Token $x$ ke liye, hum $N$ experts ke liye gating scores compute karte hain:

g_{i} (x) = softmax (W_{g} x)_{i}

jahan $g_{i} (x)$ expert $i$ ko route karne ki probability hai. Top-K routing mein (typically $K = 2$ ), hum $K$ experts choose karte hain jinki $g_{i}$ sabse zyaada ho aur normalize karte hain:

\tilde{g}_{i} (x) = {\frac{g _{i} ( x )}{\sum _{j \in TopK} g _{j} ( x )} 0 if i \in TopK otherwise

Output yeh hai:

y = i \in TopK \sum \tilde{g}_{i} (x) E_{i} (x)

Step 2: Collapse Kyun Hota Hai

Gating network $W_{g}$ ko main task loss par gradient descent se seekha jaata hai. Training ke shuruaat mein, random initialization ki wajah se expert 1 kisi common pattern mein thoda behtar ho sakta hai. Router zyaada tokens expert 1 ko bhejta hai, usse zyaada gradient milta hai, wo aur tezi se improve hota hai, aur aur bhi tokens attract karta hai — ek feedback loop. Experts 2-8 bhookhe reh jaate hain.

Step 3: Load Imbalance Measure Karna

$B$ tokens ke ek batch mein expert $i$ ko route hone wale tokens ka fraction define karo:

f_{i} = \frac{1}{B} b = 1 \sum B 1 [expert i \in TopK for token b]

Perfectly balanced: $f_{i} = K / N$ sabhi $i$ ke liye (agar 8 experts mein se Top-2 hai, toh har expert ko average par $2/8 = 25%$ tokens milne chahiye).

Collapsed: $f_{1} \approx 0.9$ , $f_{2} \approx 0.1$ , $f_{3} \dots f_{8} \approx 0$ .

Load Balancing Loss: Scratch Se Derivation

Standard approach Switch Transformer (aur related GShard) paper se aata hai. Key insight: hum ek saath do cheezein balanced chahte hain, aur hume ek aisa term chahiye jo (a) uniform distribution par minimize ho aur (b) router weights ke through differentiable ho.

Wo Do Quantities Jo Hum Track Karte Hain

$B$ tokens aur $N$ experts ke ek batch par define karo:

$f_{i}$ : expert $i$ ko actually dispatch hone wale tokens ka fraction (yeh ek hard count hai, non-differentiable):

f_{i} = \frac{1}{B} b = 1 \sum B 1 [expert i \in TopK for token b]

$P_{i}$ : router dwara expert $i$ ko assign ki gayi average routing probability (yeh $W_{g}$ ke through differentiable hai):

P_{i} = \frac{1}{B} b = 1 \sum B g_{i} (x_{b})

Switch Transformer Auxiliary Loss

L_{aux} = α \cdot N i = 1 \sum N f_{i} \cdot P_{i}

jahan $α$ ek hyperparameter hai (typically $0.01$ ).

Formula Yeh EXACT Form Kyun? (Corrected)

Yeh ek carefully engineered loss hai, aur iski correctness subtle facts par depend karti hai. Yeh simply "ek dot product minimize karo" nahi hai:

Fact 1 — Constraints matter karte hain. Construction se $\sum_{i} f_{i} = K$ (har token $K$ experts choose karta hai) aur $\sum_{i} P_{i} = 1$ (softmax ka sum ek hota hai). Bare dot product $\sum_{i} f_{i} P_{i}$ ko sirf in sum constraints ke saath minimize karne par sabse choti- $P$ expert par saara $f$ -mass daalna padega (ek corner solution) — uniformity se nahi. Toh raw dot product ke liye "uniform par minimize hota hai" wali baat galat hai.

Fact 2 — Sirf $P_{i}$ mein gradients hain. $f_{i}$ ek hard indicator count hai → uska gradient zero hai (backprop ke dauran ise constant maana jaata hai). Toh effective differentiable loss yeh hai: $L_{aux} = α N \sum_{i} const f_{i} \cdot P_{i} \Rightarrow \frac{\partial L _{aux}}{\partial P _{i}} = α N f_{i} .$ Har $P_{i}$ par gradient current measured load $f_{i}$ ke proportional hai.

Fact 3 — Yahi kaam karta hai isliye. Router ko push kiya jaata hai ki wo exactly un experts ki probability $P_{i}$ kam kare jo currently overloaded hain (bada $f_{i}$ ). Kyunki $\sum_{i} P_{i} = 1$ (softmax), overloaded experts ka $P_{i}$ kam karne se under-used ones ki probability automatically badh jaati hai. Yeh ek feedback controller hai, symmetric anti-correlation objective nahi. Product $f_{i} P_{i}$ isliye use kiya jaata hai taaki $P_{i}$ par corrective push ki magnitude us expert ke overload hone ke saath scale kare.

Fact 4 — $f_{i} \cdot P_{i}$ kyun, aur akela $\sum f_{i}^{2}$ ya $\sum P_{i}^{2}$ kyun nahi? Akela $\sum_{i} f_{i}^{2}$ non-differentiable hai (koi gradient nahi). Akela $\sum_{i} P_{i}^{2}$ sirf router ki confidence regularize karta hai lekin actual dispatched load (jo hardware/compute imbalance cause karta hai) ko ignore karta hai. Cross term $f_{i} P_{i}$ measured hard load (ek scaling constant ke roop mein) ko differentiable probability se couple karta hai, ek aisa gradient deta hai jo observed imbalance ko directly correct karta hai. Yeh (fixed $f$ ke saath) minimize hota hai jab router heavily loaded experts se probability mass door shift karta hai.

Key takeaways (corrected mental model):

Loss magically $f_{i} = P_{i} = 1/ N$ par minimal nahi hota jaise ek unconstrained dot product mein. Iski purpose dynamically achieve hoti hai: gradients load $f_{i}$ ke saath scale karte hain aur busy experts ke liye $P_{i}$ niche push karte hain.
Training ke fixed point par, agar router probabilities equalize kare ( $P_{i} \approx 1/ N$ ) aur dispatch equalize ho ( $f_{i} \approx K / N$ ), toh loss $α N \cdot N \cdot (K / N) (1/ N) = α K$ par settle hota hai, ek chota constant. Imbalance ise isse upar uthata hai.

Total Training Loss

L_{total} = L_{task} + L_{aux}

Router ab ek trade-off optimize karta hai: tokens ko task ke liye best expert tak route karo ( $L_{task}$ minimize karo) jabki load balanced rakho ( $L_{aux}$ minimize karo by already-busy experts par probability over-concentrate na karo).

Alternative: Capacity Factor aur Token Dropping

Kuch MoE implementations (jaise Switch Transformer) ek capacity factor $C$ bhi use karte hain:

Expert Capacity = C \cdot \frac{B}{N}

Yeh kaise kaam karta hai:

Agar $C = 1.0$ hai, toh har expert exactly $B / N$ tokens handle kar sakta hai (perfect balance).
Agar expert ki assignments capacity se zyaada ho jaayein, toh excess tokens drop ho jaate hain (unki representation residual ke roop mein pass hoti hai, expert ko bypass karke).
Dropping back-pressure create karta hai: agar expert 1 overloaded hai, toh tokens drop hote hain, model un tokens par worse perform karta hai, aur gradients expert 1 ko over-routing discourage karte hain.

Yeh kyun help karta hai:

Hard constraint: physically ek expert ko saare tokens lene se rokta hai.
Router ko under-utilized experts explore karne ke liye encourage karta hai.
Trade-off: tokens drop karna performance hurt kar sakta hai agar $C$ bahut low ho. Typical values: $C \in [1.0, 1.5]$ .

Worked example Example 1: Switch Auxiliary Loss Compute Karna

Setup: 4 experts, 8 tokens ka batch, Top-1 routing. Routing outcomes (hard dispatch):

Tokens 1,2,3,4 → Expert 1
Tokens 5,6 → Expert 2
Token 7 → Expert 3
Token 8 → Expert 4

$f_{i}$ compute karo (dispatched fraction):

f_{1} = 4/8 = 0.5, f_{2} = 2/8 = 0.25, f_{3} = 1/8 = 0.125, f_{4} = 1/8 = 0.125

(Note: $\sum_{i} f_{i} = 1 = K$ kyunki $K = 1$ hai.)

Average gating probabilities (softmax, toh $\sum_{i} P_{i} = 1$ ):

P_{1} = 0.6, P_{2} = 0.2, P_{3} = 0.1, P_{4} = 0.1

Auxiliary loss ( $α = 0.01$ , $N = 4$ ):

L_{aux} = 0.01 \cdot 4 \cdot (0.5 \cdot 0.6 + 0.25 \cdot 0.2 + 0.125 \cdot 0.1 + 0.125 \cdot 0.1)

= 0.04 \cdot (0.3 + 0.05 + 0.0125 + 0.0125) = 0.04 \cdot 0.375 = 0.015

Yeh step kyun? $P_{i}$ ke w.r.t. gradient $α N f_{i}$ hai. Expert 1 ke liye: $0.01 \cdot 4 \cdot 0.5 = 0.02$ ; expert 3 ke liye: $0.01 \cdot 4 \cdot 0.125 = 0.005$ . Router ko push kiya jaata hai ki $P_{1}$ ko $P_{3}$ se chaar guna zyaada strongly lower kare — exactly us hisaab se jo har expert kitna overloaded hai. Softmax phir wo probability under-used experts ki taraf redistribute karta hai.

Worked example Example 2: Balancing Ka Expert Utilization Par Effect

Auxiliary loss ke bina:

1000 steps ke baad: Expert 1 ko 80% tokens milte hain, experts 2-8 ko 2-3% each.
Experts 2-8 ke paas near-random weights hain (koi training signal nahi).
Effective model capacity: total parameters ka 1/8.

Auxiliary loss ke saath ( $α = 0.01$ ):

1000 steps ke baad: Har expert ko 10-15% tokens milte hain (roughly uniform).
Saare experts distinct specializations develop karte hain (e.g., syntax, math, dialogue, ...).
Effective capacity: ~8x zyaada, same parameter count par better performance.

Yeh step kyun? Gradient $\propto f_{i}$ feedback loop ko rokta hai: chahe expert 1 initially better ho, uska bada $f_{1}$ ek bada downward push $P_{1}$ par produce karta hai, probability mass free karta hai taaki dusre experts dispatch ho sakein aur finally learning signal pa sakein.

Common Hyperparameters

Hyperparameter	Typical Value	Purpose
$α$ (aux loss weight)	0.01	Task loss aur balance ke beech trade-off
Capacity factor $C$	1.0 - 1.5	Load variation ke liye buffer
Top-K	1 ya 2	Har token ke liye experts ki sankhya

$α = 0.01$ kyun? Empirically tuned. Bahut zyaada: router probabilities ko flat kar deta hai balance satisfy karne ke liye aur ignore karta hai ki kaunsa expert best hai (accuracy hurt hoti hai). Bahut kam: load collapse wapas aata hai. $0.01$ ek robust sweet spot hai.

Common mistake Mistake 1:

$α$ Bahut Zyaada Set Karna Galat idea: "Zyaada load balancing hamesha better hoti hai, toh main $α = 1.0$ set karunga." Kyun sahi lagta hai: Aap sach mein expert collapse avoid karna chahte ho, toh ek strong penalty safe lagta hai.

Reality: $α = 1.0$ ke saath, auxiliary loss dominate karta hai. Kyunki gradient $α N f_{i}$ bahut bada hota hai, router $P_{i}$ ko uniform ki taraf flatten karta hai chahe har token ke liye kaunsa expert actually best ho. Result: saare experts mediocre generalists ban jaate hain, performance ek dense baseline se bhi niche aa jaati hai. Aapne expert collapse ki jagah expert homogeneity le li.

Fix: $α \in [0.001, 0.01]$ use karo. Task loss dominate karna chahiye, auxiliary loss sirf ek gentle regularizer ke roop mein. Dono $L_{task}$ aur expert utilization monitor karo; agar utilization abhi bhi imbalanced hai, toh $α$ gradually badhao.

Common mistake Mistake 2:

$\sum_{i} f_{i} P_{i}$ Ko Ek Symmetric "Anti-Correlation" Objective Samajhna Galat idea: "Loss $f_{i}$ aur $P_{i}$ ko anti-correlated banata hai, toh yeh uniform point $f_{i} = P_{i} = 1/ N$ par minimize hota hai." Kyun sahi lagta hai: Do "balanced" quantities ka product aisa lagta hai ki jab sab equal ho toh sabse chota hona chahiye.

Reality: True constraints $\sum_{i} f_{i} = K$ aur $\sum_{i} P_{i} = 1$ ke under, raw dot product $\sum_{i} f_{i} P_{i}$ actually ek corner par minimize hota hai — sabse chote- $f$ expert par saara $P$ -mass dump karke — uniformity par nahi. Iske alawa, $f_{i}$ non-differentiable hai; gradients sirf $P_{i}$ ke through flow karte hain. Toh yeh loss bilkul bhi symmetric anti-correlation objective nahi hai. Yeh ek feedback controller hai: yeh $f_{i}$ ko $P_{i}$ par ek constant weight maanta hai aur currently-overloaded experts se probability door push karta hai. Balance training ke dauran dynamically emerge karta hai, static minimization se nahi.

Fix: Gradient $\partial L_{aux} / \partial P_{i} = α N f_{i}$ ke baare mein socho, na ki "function kahan minimal hai" ke baare mein. Har expert par corrective push uske measured load ke proportional hai.

Jab Load Balancing Kaafi Nahi Hota

Bahut chote batch sizes: $B = 16$ aur $N = 8$ experts (Top-2) ke saath, stochastic imbalance unavoidable hai. Solution: micro-batches par stats accumulate karo, ya $B$ badhao.
Highly skewed data: Agar 90% English hai aur 10% code, toh ek early-specialized expert dominate kar sakta hai. Solution: data-aware balancing ya mild imbalance accept karo agar task performance improve hoti hai.
Sparse expert gradients: Kam tokens se feed hone wale experts ko noisy updates milte hain. Solution: capacity factor $C$ badhao ya gradient accumulation use karo.

Recall Ek 12-Saal Ke Bacche Ko Explain Karo

Ek group project imagine karo jisme 8 dost hain aur ek team leader (router) hai jo tasks baanta hai. Agar akela chhod do, toh leader SAARI drawing Alex ko de sakta hai kyunki Alex ne achha start kiya. Alex overloaded ho jaata hai aur baaki 7 kabhi drawing seekh hi nahi paate.

Load balancing aisa hai jaise teacher ek tally sheet dekh raha ho ki kisne kitne tasks liye. Agar Alex ka tally bahut bada hai, toh teacher leader ko ek nudge deta hai Alex ke tally ke proportional: "Alex ko itna prefer karna band karo." Nudge us insaan ke liye bada hota hai jo sabse zyaada overloaded hai. Kyunki leader ki preferences 100% mein add honi chahiye, Alex se preference kheenchne par quiet dosto ko automatically zyaada milta hai — toh unhe finally tasks milte hain aur wo seekhte hain.

Clever baat yeh hai: tally khud (kisne kitne tasks liye) sirf ek number hai jo teacher padhta hai — teacher sirf leader ki preferences badalta hai, aur unhe sabse busy insaan ke liye zyaada strongly badalta hai. Yahi exactly MoE auxiliary loss kaam karta hai.

Connections

Mixture of Experts (MoE) Architecture: Base architecture; load balancing MoE ke liye scale par essential hai.
Top-K Routing in MoE: Routing mechanism jo load imbalance problem create karta hai.
Switch Transformer: $α N \sum f_{i} P_{i}$ auxiliary loss aur capacity factor ka source.
GShard: Large sharded MoE ke liye related importance/load auxiliary losses.
Expert Capacity and Token Dropping: Soft auxiliary loss ka hard-constraint alternative.
MoE Training Dynamics: MoE models kaise converge karte hain, balancing effects including.
Gradient Routing in MoE: Gradients router ke through kaise flow karte hain (sirf $P_{i}$ ke via).

#flashcards/ai-ml

MoE mein load balancing kaunsi core problem solve karta hai?

Load balancing ke bina, kuch experts ko zyaadatar tokens milte hain aur wo training dominate karte hain, jabki dusre experts useless mein collapse ho jaate hain, parameters aur compute waste karte hain. Load balancing ensure karta hai ki saare experts ko roughly equal training signal mile aur wo diverse skills develop karen.

Switch Transformer auxiliary load-balancing loss state karo.

L_{aux} = α \cdot N \sum_{i = 1}^{N} f_{i} \cdot P_{i}

, jahan

f_{i}

expert

i

ko dispatch hone wale tokens ka fraction hai (non-differentiable count),

P_{i}

expert

i

ke liye average router softmax probability hai (differentiable),

N

experts ki sankhya hai, aur

α \approx 0.01

. Note karo ki

α

pehle se andar hai, toh total loss

L_{task} + L_{aux}

hai (na ki

L_{task} + α L_{aux}

Auxiliary loss mein kaun sa term gradients carry karta hai, aur wo gradient kya hai?

Sirf

P_{i}

gradients carry karta hai (softmax differentiable hai).

f_{i}

ek hard dispatch count hai aur ise constant maana jaata hai. Gradient hai

\partial L_{aux} / \partial P_{i} = α N f_{i}

, yaani expert

i

ke measured load ke proportional.

Yeh kehna kyun galat hai ki

\sum_{i} f_{i} P_{i}

uniform point

f_{i} = P_{i} = 1/ N

par minimize hota hai?

True constraints

\sum_{i} f_{i} = K

aur

\sum_{i} P_{i} = 1

ke under, dot product

\sum_{i} f_{i} P_{i}

ek corner par minimize hota hai — sabse chote-

f

expert par saara

P

-mass concentrate karke — uniformity par nahi. Loss is baat se kaam nahi karta ki uniform par minimal ho, balki ek feedback controller ki tarah: uska gradient

α N f_{i}

probability ko overloaded experts se door push karta hai, aur softmax redistribute karta hai.

Cross term

f_{i} P_{i}

akele

\sum f_{i}^{2}

\sum P_{i}^{2}

ki jagah kyun use karte hain?

Akela

\sum f_{i}^{2}

non-differentiable hai (hard count se koi gradient nahi). Akela

\sum P_{i}^{2}

sirf router confidence regularize karta hai aur actual dispatched load jo real compute imbalance cause karta hai use ignore karta hai. Cross term measured hard load

f_{i}

(ek constant scaling ke roop mein) ko differentiable

P_{i}

se couple karta hai, toh har

P_{i}

par corrective gradient observed overload ke saath directly scale karta hai.

Top-K routing mein "perfectly balanced" ka kya matlab hai?

Har expert ko average par

f_{i} = K / N

tokens milte hain. 8 experts mein se Top-2 ke liye, har expert ko

2/8 = 25%

milna chahiye.

Capacity factor

C

kya hai aur yeh balance kaise enforce karta hai?

Expert capacity

= C \cdot B / N

. Capacity se zyaada assign hone wale tokens drop ho jaate hain (residual ke roop mein pass hote hain). Overloaded experts tokens drop karte hain, performance hurt hoti hai, toh gradients over-routing discourage karte hain. Typical

C \in [1.0, 1.5]

Kya hota hai agar

α

bahut zyaada set kar dein (e.g., 1.0)?

Auxiliary gradient

α N f_{i}

bahut bada ho jaata hai, router probabilities ko uniform ki taraf flatten kar deta hai chahe har token ke liye kaunsa expert actually best ho. Experts mediocre generalists ban jaate hain aur performance ek dense base se bhi niche aa sakti hai.