Mixed precision training — Hinglish notes

3.3.9 · AI-ML › Deep Learning Frameworks

What Problem Does This Solve?

Large neural networks train karne mein teen bottlenecks hain:

Memory: Modern models (GPT, Vision Transformers) mein billions of parameters hote hain. Har FP32 parameter 4 bytes leta hai. 175B parameter model ko sirf weights store karne ke liye 700GB chahiye!
Compute: Matrix multiplications training time ko dominate karti hain. GPUs mein specialized hardware (Tensor Cores) hote hain jo FP16 ops ko FP32 se 2-8× faster run karte hain.
Communication: Distributed training mein, gradients ko GPUs ke across sync karna padta hai. Precision aadha karne se bandwidth requirements bhi aadhi ho jaati hain.

Lekin naively FP16 par switch karne se gradient underflow hota hai: tiny gradients zero ho jaate hain, aur training diverge kar jaati hai.

Key insight yeh hai: gradients ka dynamic range bahut bada hota hai. Early layers mein gradients $1 0^{- 6}$ ke around ho sakte hain, jabki later layers mein $1 0^{- 2}$ . FP16 ki limited range dono ko represent nahi kar sakti.

How Mixed Precision Training Works

Is algorithm ke teen core components hain:

1. FP32 Master Weights

Kyun: Weight updates cumulative hote hain. Agar learning rate $1 0^{- 4}$ hai aur gradient $1 0^{- 3}$ hai, to update $1 0^{- 7}$ hoga—jo FP16 ki precision ( $1 0^{- 3}$ ) se bhi chhota hai. Update zero ho jaayega, aur weights kabhi change nahi honge.

Kaise: Weights ki do copies maintain karo:

FP16 copy forward/backward passes ke liye (fast compute)
FP32 master copy updates accumulate karne ke liye (precision preserve karo)

Derivation from First Principles:

Ek single parameter $w$ ko gradient descent se update karte hain: $w_{t + 1} = w_{t} - η \cdot g_{t}$

FP16 mein, agar $η g_{t} < ϵ_{F P 16} \approx 1 0^{- 3}$ hai, to update ban jaata hai: $w_{t + 1}^{F P 16} = round_{F P 16} (w_{t} - η g_{t}) = w_{t}$ (chhota update zero ho jaata hai)

$N$ steps ke baad jab koi update nahi hua, humne $N η g_{t}$ ka accumulated change kho diya—training ruk jaati hai.

Solution: Master weights ko FP32 mein store karo: $w_{t + 1}^{F P 32} = w_{t}^{F P 32} - η \cdot g_{t}^{F P 32}$ Chahe individual updates tiny hon, higher-precision master mein accumulate hote rehte hain. Phir next forward pass ke liye FP16 mein cast karo: $w_{t + 1}^{F P 16} = cast (w_{t + 1}^{F P 32})$

2. Loss Scaling

Kyun: Early layers ke gradients often FP16 ki minimum normal value ( $6 \times 1 0^{- 8}$ ) se neeche hote hain. Yeh underflow hokar zero ho jaate hain, jisse "dead neurons" ban jaate hain jo kabhi learn nahi karte.

Kya: Loss ko backprop se pehle ek bade scalar $S$ se multiply karo (typically $512$ to $6536$ ). Isse gradients FP16 ke representable range mein shift ho jaate hain. Gradients compute karne ke baad, $S$ se divide karo taaki master weights update karne se pehle true values restore ho jayein.

Derivation:

Gradients ke liye chain rule deta hai: $\frac{\partial L}{\partial w} = \frac{\partial L}{\partial a} \cdot \frac{\partial a}{\partial w}$

Agar $\frac{\partial L}{\partial a} \approx 1 0^{- 7}$ hai (deep networks mein common), to FP16 ise represent nahi kar sakta—underflow hokar zero ho jaata hai.

Loss scaling trick: Ek scaled loss define karo $L^{'} = S \cdot L$ jahan $S ≫ 1$ . Tab: $\frac{\partial L ^{'}}{\partial w} = S \cdot \frac{\partial L}{\partial w}$

Ab agar $\frac{\partial L}{\partial w} = 1 0^{- 7}$ hai aur $S = 1024$ hai, to $\frac{\partial L ^{'}}{\partial w} = 1 0^{- 4}$ hai, jise FP16 represent kar sakta hai!

Backprop ke baad, unscale karo: $\frac{\partial L}{\partial w} = \frac{1}{S} \cdot \frac{\partial L ^{'}}{\partial w}$

Yeh step kyun?: Scaling loss par hoti hai (computational graph ka top), isliye automatically poore backward pass mein propagate ho jaati hai. Yeh har layer ke gradients ko individually scale karne se zyaada efficient hai.

Formula Loss Scaling Algorithm

# Forward pass (FP16)
loss = model(x, y)  # Computed in FP16
 
# Scale loss BEFORE backward
scaled_loss = loss * scale_factor  # scale_factor ~ 1024
 
# Backward pass (gradients are now scaled, stay in FP16 range)
scaled_gradients = backward(scaled_loss)
 
# Unscale gradients before optimizer step
true_gradients = scaled_gradients / scale_factor
 
# Update FP32 master weights
optimizer.step(true_gradients)

Dynamic Loss Scaling: Bade $S$ se shuru karo (jaise $6536$ ). Agar gradients overflow karein (inf ya nan produce karein), $S$ ko aadha karo aur retry karo. Agar training $N$ steps (jaise $2000$ ) tak stable rahe, $S$ ko double karo. Yeh har model ki gradient distribution ke hisaab se adapt karta hai.

3. FP16 Arithmetic for Most Operations

Kyun: Matrix multiplications (GEMM) training time dominate karti hain—often 90%+ compute. Modern GPUs mein Tensor Cores hain jo FP16 ops ko FP32 se 2-16× zyaada throughput par execute karte hain.

Kya: Forward/backward passes ke dauran activations, gradients, aur weights FP16 mein store karo. Sirf optimizer state aur master weights FP32 mein rehte hain.

Kaise:

Har forward pass ke start mein FP32 master se weights ko FP16 mein cast karo
Activations FP16 mein compute karo: $a^{F P 16} = σ (W^{F P 16} x F P 16 + b^{F P 16})$
Loss compute karo aur scale karo
Backward pass FP16 mein (gradients FP16 mein accumulate hote hain)
Optimizer ke liye gradients ko FP32 mein cast karo

Worked example Example 1: Single Weight Update

Setup: Ek weight $w_{0}^{F P 32} = 0.5$ se start karta hai, learning rate $η = 0.001$ , gradient $g = 0.003$ . Step 1 (Forward): FP16 mein cast karo $w_{0}^{F P 16} = cast (0.5) = 0.5$ (Koi precision loss nahi kyunki $0.5$ FP16 mein exactly representable hai) Step 2 (Backward): FP16 mein gradient compute karo, $g^{F P 16} = 0.003$ Step 3 (Optimizer): FP32 mein update karo

g^{FP32} &= \text{cast}(0.003) = 0.003 \\ w_1^{FP32} &= 0.5 - 0.0001 \times 0.003 = 0.4997 \end{aligned}$$ **Yeh step kyun?**: Update ka magnitude $3 \times 10^{-7}$ hai. FP16 precision $\sim 10^{-3}$ hai, isliye: $$w_1^{FP16} = \text{round}_{FP16}(0.5 - 0.0003) = 0.5$$ Update kho jaata! FP32 ise preserve karta hai: $0.4997 \neq 0.5$. **Step 4**: Next forward pass $w_1^{FP16} = \text{cast}(0.4997) \approx 0.4997$ use karta hai.

Common Pitfalls and Fixes

Common mistake Mistake 1: Forgetting to Unscale Gradients Before Clipping

Kyun sahi lagta hai: "Main gradients ko norm 1.0 par clip karunga taaki exploding gradients na hon." Kyun galat hai: Agar loss scale $S = 1024$ hai, to scaled gradients ke norms $1024 \times$ bade hain. 1.0 par clipping essentially unscaled gradients ko $1/1024 \approx 0.001$ par clip karti hai—bahut zyaada aggressive! Training barely progress karti hai.

Fix: Clipping se pehle unscale karo:

scaled_grad = backward(scaled_loss)
true_grad = scaled_grad / scale_factor
clipped_grad = clip_by_norm(true_grad, max_norm=1.0)
optimizer.step(clipped_grad)

Steel-man: Yeh mistake isliye hoti hai kyunki gradient clipping aksar alag line mein add ki jaati hai, aur unscaling order bhoolna aasan hai. Fix: unscaling ko apne gradient preparation pipeline ka hissa banao, hamesha clipping se pehle.

Common mistake Mistake 2: Using Static Loss Scale for All Models

Kyun sahi lagta hai: "ResNet ke liye loss scale 1024 kaam kiya, to main ise har jagah use karunga."

Kyun galat hai: Alag architectures mein alag gradient distributions hoti hain. Vision Transformers mein aksar CNNs se bade gradients hote hain. Fixed scale bahut chhota (underflow) ya bahut bada (overflow) ho sakta hai.

Evidence: Scale $1024$ ke saath ViT train karne par 100 steps ke baad nan produce hota hai. Dynamic scaling ke saath ( $6536$ se shuru hokar, $512$ tak adapt hokar), training stable rehti hai.

Fix: Dynamic loss scaling use karo ya per-model tune karo. Gradient histograms monitor karo—agar bahut saare zeros (underflow) ya inf (overflow) dikh rahe hain, to scale adjust karo.

Steel-man: Static scaling reproducibility ke liye appealing lagti hai, lekin gradient magnitudes architectures across 1000× vary karte hain. Dynamic scaling robust default hai.

Implementation Pseudocode

# Initialization
model_fp16 = model.half()  # FP16 copy for compute
master_weights = model.parameters()  # FP32 master copy
optimizer = Adam(master_weights)
scale = 2**16  # Initial loss scale (6536)
 
for batch in dataloader:
    # Forward pass in FP16
    with autocast(dtype=torch.float16):
        output = model_fp16(batch.x)
        loss = criterion(output, batch.y)
    
    # Scale loss to prevent gradient underflow
    scaled_loss = loss * scale
    # Backward pass (gradients computed in FP16, but scaled)
    scaled_loss.backward()
    
    # Unscale gradients
    for param in master_weights:
        if param.grad is not None:
            param.grad = param.grad / scale
    # Clip gradients (on unscaled values!)
    torch.nn.utils.clip_grad_norm_(master_weights, max_norm=1.0)
    
    # Optimizer step on FP32 master weights
    optimizer.step()
    optimizer.zero_grad()
    # Copy updated master weights back to FP16 model
    for param_fp16, param_fp32 in zip(model_fp16.parameters(), master_weights):
        param_fp16.data.copy_(param_fp32.data)
    # Dynamic loss scaling: check for overflow
    if any(torch.isnan(p.grad).any() for p in master_weights if p.grad is not None):
        scale = scale / 2  # Halve scale on overflow
    elif step % 2000 == 0:
        scale = min(scale * 2, 2**16)  # Double scale if stable

Theoretical Guarantees

Convergence: Mixed precision training provably FP32 training ke same optimum par converge karta hai mild conditions ke andar:

Assumption: Gradients $g$ satisfy $∥ g ∥_{\infty} < K$ (bounded) aur loss scale $S$ ensure karta hai \cdot g_{\min} > \epsilon_{FP16}$ (no underflow).
Result: Agar learning rate $η < \frac{1}{L}$ (standard SGD condition), to expected distance to optimum FP32 ke same rate se decrease hoti hai: $E [∥ w_{t} - w^{*} ∥^{2}] \leq \frac{C}{t}$ jahan $C$ $L$ (loss smoothness) par depend karta hai lekin precision par nahi.

Proof sketch: Master weights updates exactly accumulate karte hain (koi precision loss nahi). FP16 forward/backward har operation par rounding errors $δ$ introduce karte hain. Lekin modern GPUs round-to-nearest mode use karte hain, isliye $E [δ] = 0$ (unbiased). Error variance $O (ϵ_{F P 16}^{2})$ hai, jo averaging ki wajah se $t \to \infty$ par vanish ho jaata hai. Loss scaling ensure karti hai ki $δ$ scaled gradients ke relative measure hota hai, $δ / g$ ko chhota rakhte hue.

When to Use Mixed Precision

Tab use karo jab:

10M parameters wale models train kar rahe ho (memory savings matter karti hain)
Tensor Cores wale GPUs use kar rahe ho (V100, A100, RTX 20/30 series)
Distributed training mein (gradient communication bottleneck hai)
Inference memory-bound ho (FP16 models directly deploy kar sakte hain)

Tab use mat karo jab:

Numerical issues debug kar rahe ho (mushkil hai batana ki bug tumhara hai ya precision-related)
Aisi datasets ke saath kaam kar rahe ho jahan values >5 orders of magnitude span karti hain (jaise scientific computing with huge ranges)
Custom CUDA kernels use kar rahe ho jo FP16 ke saath tested nahi hain

Real-world impact: GPT-3 training ne mixed precision use ki aur iske bina infeasible hoti (175B parameters → ~700GB FP32, lekin mixed precision aur activation checkpointing ke saath ~350GB mein fit ho jaata hai).

Recall 12 Saal ke Bachche Ko Samjhao

Pretend karo ki tum ek video game khel rahe ho, aur tumhare paas apne gold coins store karne ke liye ek piggy bank hai. Har coin ko do tareekon se store kiya ja sakta hai:

Chhote pouches (FP16): Coins jaldi store karte hain, kam jagah lete hain, lekin sirf 1 se 100 tak count kar sakte hain.
Bade chests (FP32): Slow, heavy, lekin 0.001 se 1,000,000 tak count kar sakte hain.

Zyaadatar time tum coins collect aur spend karte hue bhaag rahe ho, isliye chhote pouches use karte ho—yeh quick hai! Lekin har raat, tum apne saare pouches ek bade chest mein daalte ho apni total wealth jodne ke liye. Is tarah, chhote coins bhi (jaise 0.01 gold) kho nahi jaate.

Ab kabhi kabhi tumhe bahut chhote treasure milte hain (0.001 gold). Agar seedha chhote pouch mein daalo, to yeh round hokar zero ho jaata hai aur gayab! Isliye ek trick use karte ho: collect karte waqt saare treasure ko 1000 se multiply karo. Ab 0.0001, 0.1 ban jaata hai, jo pouch hold kar sakta hai. Baad mein, jab bade chest mein daalo, 1000 se divide karo real value waapas paane ke liye.

Mixed precision training exactly yahi hai: zyaadatar kaam ke liye fast, chhota storage (FP16) use karo, lekin ek careful ledger (FP32 master weights) rakho taaki tiny updates vanish na hon. Gradients ko scale up karo taaki woh disappear na hon. Isse huge AI models train karna 2× faster aur 30% kam memory use karta hai—jaise video game boss ko speed potion ke saath harna!

Connections

3.3.08-Gradient-accumulation – Mixed precision ko aksar gradient accumulation ke saath combine kiya jaata hai bade batches fit karne ke liye
3.3.05-Learning-rate-schedules – Loss scaling LR ke saath interact karti hai: dono gradient magnitude affect karte hain
4.2.03-Tensor-cores – Mixed precision ke liye hardware motivation
3.3.10-Gradient-checkpointing – Ek aur memory optimization, aksar saath use hoti hai
2.4.06-Batch-normalization – Special handling chahiye (FP32 running stats)
5.1.02-Model-parallelism – Mixed precision distributed training mein communication reduce karti hai
1.3.07-Numerical-stability – Speed aur precision ke beech fundamental tradeoffs

#flashcards/ai-ml

Mixed precision training ke teen core components kya hain? :: (1) Accurate updates ke liye FP32 master weights, (2) Gradient underflow rokne ke liye Loss scaling, (3) Compute speed up karne ke liye forward/backward passes mein FP16 arithmetic.

Hum poori tarah FP16 mein train kyun nahi kar sakte?

FP16 ki limited range (

1 0^{- 8}

6504

) se (1) gradient underflow (tiny gradients → zero) aur (2) weight update underflow (chhote learning rates × gradients → koi change nahi) hota hai. FP32 mein master weights update underflow ko prevent karte hain.

Derive karo ki loss scaling gradient underflow kyun rokti hai :: Agar gradient $g = 1 0^{- 7} < 6 \times 1 0^{- 8}$ (FP16 min) hai, to yeh 0 par underflow hota hai. Loss ko $S = 1024$ se scale karo: backward deta hai $g^{'} = S g = 1 0^{- 4}$ (representable). Unscale karo: $g = g^{'} / S$ optimizer step se pehle FP32 mein true value recover karta hai.

Mixed precision training ka memory formula kya hai?

Memory

\approx

(FP16 weights + FP32 master weights + FP16 gradients + FP16 activations + FP32 optimizer state). Typically pure FP32 se 25-35% savings, mainly activations se jo memory ka ~50% hoti hain.

Gradients unscale kab karna chahiye—gradient clipping se pehle ya baad mein?

Clipping se pehle. Scaled gradients ke norms

S \times

bade hote hain. Agar scaled grads ko norm 1.0 par clip karo, to actually true grads ko

1/ S \approx 0.001

par clip kar rahe ho (with

S = 1024

), jo bahut zyaada aggressive hai.

Batch normalization running statistics ko FP32 mein kyun rehna chahiye?

Running mean/var EMA se update hote hain:

μ_{new} = 0.9 μ_{old} + 0.1 μ_{batch}

0.1 \times

chhota value update often

< 1 0^{- 3}

(FP16 precision limit) hota hai, underflow cause karta hai. Accumulated error inference tod deta hai.

Dynamic loss scaling kya hai aur kyun zaroori hai?

Bade scale se shuru karo (jaise

2^{16}

). Agar gradients overflow karein (inf/nan produce karein), scale aadha karo. Agar training

N

steps tak stable rahe, scale double karo. Zaroori hai kyunki alag models mein gradient distributions magnitude mein 1000× span karti hain—koi ek static scale universally kaam nahi karta.

FP16 vs BF16 compare karo :: FP16: 5 exp bits, 10 mantissa bits. Range $\sim 1 0^{- 8}$ to $6504$ , precision $\sim 1 0^{- 3}$ . BF16: 8 exp bits, 7 mantissa bits. FP32 jaisi range ( $\sim 1 0^{- 38}$ to $1 0^{38}$ ), kam precision. BF16 mein overflow issues kam hote hain lekin accuracy kam hai; TPUs aur newer GPUs par preferred hai.

Modern GPUs par mixed precision training se speedup kitna hota hai?

V100/A100 par Tensor Cores ke saath typically 2-3× end-to-end training speedup. Raw matrix multiply FP16 mein 8-16× faster ho sakta hai, lekin overhead (data loading, optimizer CPU work) total speedup limit karta hai. Memory savings (25-35%) se bade batch sizes allow hote hain, jo aur speedup deta hai.