Distributed Scaling and Quantization: MoE, Megatron-LM, Pipeline 1F1B, and ZeRO
How does Sparse MoE scale total parameter capacity without increasing active FLOPs per token?
How this lesson unfolds.
When model parameters exceed the memory of a single accelerator, distributed computing becomes essential. In this chapter, we master modern distributed training architectures and post-training compression.
How does Sparse MoE scale total parameter capacity without increasing active FLOPs per token?
Instead of forcing one general doctor to treat every symptom, a triage nurse routes each patient to the two best medical specialists, keeping compute time constant while vastly expanding collective hospital knowledge.
- $$y = \sum_{i \in \text{TopK}} \frac{e^{h_i(x)}}{\sum_{j \in \text{TopK}} e^{h_j(x)}} E_i(x)$$
- Learning goalDerive the top-k sparse routing equation and explain why auxiliary load-balancing loss prevents expert collapse.
- Helpful foundationsFeed-forward network architecture · Softmax gating · Shazeer et al. (2017) Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer · Jiang et al. (2024) Mixtral of Experts
- Common trapBelieving an 8x7B MoE model requires 56B FLOPs per token; only 2 experts (13B FLOPs) execute per token.
- Next connectionNext, explore how large weights are compressed into 4-bit integers with zero perplexity degradation.
More context & mechanics
In dense transformers, all tokens activate all parameters in the Feed-Forward Network (FFN). Sparse Mixture of Experts (MoE) introduces E expert FFNs:
y = ∑_(i=1)^(E) G(x)ᵢ Eᵢ(x)
where gating network G(x) = Softmax(TopK(H(x), K)) selects only the K highest-scoring experts (K ≪ E, typically K=2).
The Routing Collapse Hazard:
Without regularization, gating networks suffer rich-get-richer feedback loops where 1 or 2 experts receive all tokens while others starve. To ensure uniform hardware utilization across GPUs, modern MoE models enforce an auxiliary load balancing loss:
L_(aux) = α · E ∑_(i=1)^(E) fᵢ Pᵢ
where fᵢ is the fraction of tokens routed to expert i, and Pᵢ is the average gating probability assigned to expert i.
Consider Mixtral 8x7B (E = 8 experts, K = 2 active):
- Total parameters: 8 × 7B ≈ 46.7B parameters (due to shared attention layers).
- Active parameters per token: Only 12.9B parameters!
- For T = 4 tokens and router logits H(x), top-2 gating normalizes probabilities:
- Token "quantum": Expert 2 (p=0.68), Expert 5 (p=0.32). Output = 0.68 E₂(x) + 0.32 E₅(x).
- Token "court": Expert 1 (p=0.55), Expert 7 (p=0.45).
Each token computes only 2/8 = 25\% of total FFN FLOPs, achieving the reasoning power of a 47B model at the speed of a 13B model.
Core architecture of state-of-the-art models including Mixtral 8x7B, Mixtral 8x22B, DBRX, and DeepSeek-V2/V3. MoE enables frontier models to scale parameter counts into hundreds of billions while keeping inference latency within real-time interactive thresholds.
The mixture of experts concept was introduced by Robert Jacobs, Michael I. Jordan, and Geoffrey Hinton in 1991. It was adapted to modern deep transformers by Noam Shazeer et al. (Google Brain 2017) and elevated to modern frontier efficiency by Mistral AI in 2024 with Mixtral.
"Sparse experts route compute where needed: build a giant's brain, spend a child's FLOPs." Top-2 gating keeps 75% of parameters dormant on any single token.
How do second-order Hessian corrections in GPTQ and activation awareness in AWQ quantize weights to 4 bits with near-zero loss?
If you must compress a high-resolution painting into a 16-color palette, rounding every pixel independently makes the image grainy. Adjusting surrounding pixels to compensate for earlier color rounding preserves crisp sharp edges.
- $$\Delta w = -\frac{w_q - \hat{w}_q}{[H^{-1}]_{qq}} H^{-1}_{:,q}$$
- Learning goalDerive the optimal brain surgeon weight update formula and contrast GPTQ with activation-aware AWQ.
- Helpful foundationsMatrix multiplication · Taylor series approximation · Frantar et al. (2022) GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers · Lin et al. (2023) AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
- Common trapBelieving 4-bit quantization degrades language reasoning; second-order error compensation maintains near-lossless output quality.
- Next connectionNext, discover how multi-GPU clusters split individual matrix multiplications across devices with Megatron-LM.
More context & mechanics
Uniform affine quantization maps continuous floating-point weights W to discrete b-bit integers q:
q = clamp(⌊ W/S ⌉ + Z, -2^(b-1), 2^(b-1)-1), Ŵ = S(q - Z)
where S = (W_(max) - W_(min)) / (2^b - 1) is the scale factor and Z is the zero-point.
Second-Order Optimal Brain Surgeon (GPTQ):
When rounding weight w_q, the optimal compensation Δ w for unquantized weights minimizes squared error E = 1/2(W - Ŵ)^T H (W - Ŵ):
Δ w = -(w_q - ŵ_q) / ([H^(-1)]_(qq)) H^(-1)_(:,q)
where H = 2 X X^T is the Hessian matrix of layer inputs X.
Activation-Aware Weight Quantization (AWQ):
AWQ observes that only 0.1\% - 1\% of channels contain salient activation outliers. Instead of quantizing all channels equally, AWQ applies a per-channel scale factor s > 1 to protect salient weights: W' = W · diag(s), X' = diag(s)^(-1) X.
Quantizing LLaMA-3 70B:
- FP16 baseline size: 70B × 2 bytes = 140 GB VRAM (requires two 80GB A100s).
- INT4 quantized size: 70B × 0.5 bytes = 35 GB VRAM (fits comfortably on a single consumer GPU or desktop workstation!).
- Perplexity impact: FP16 baseline perplexity = 5.68; GPTQ INT4 perplexity = 5.74 (<1.1\% deviation!).
Powers modern high-efficiency inference frameworks including llama.cpp, Ollama, vLLM, TensorRT-LLM, and ExLlamaV2, allowing developers to run 70B parameter models locally on a single workstation.
Quantization mathematics originates in Claude Shannon's information theory and PCM (1948). In deep learning, Optimal Brain Damage (LeCun et al., 1989) introduced Hessian-based weight pruning. Frantar et al. (2022) created GPTQ using second-order Taylor approximations, while Lin et al. (MIT 2023) developed AWQ by protecting salient activation channels.
"Quantization trades precision bits for memory bandwidth: scale the range, zero the offset, preserve the salient 1%." High-bit outliers preserve reasoning; low-bit matrices deliver speed.
How does Megatron-LM partition Multi-Head Attention and MLP layers across GPUs with minimal inter-GPU communication?
To multiply giant matrices that cannot fit in one brain, two workers each take half of the blueprint. One computes his half, the other computes her half, and they only speak once at the end to add their numbers together.
- $$Z = \text{All-Reduce}\left( \text{GELU}(X A_1) B_1 + \text{GELU}(X A_2) B_2 \right)$$
- Learning goalIllustrate how column and row matrix slicing fuses element-wise activations without intermediate synchronization.
- Helpful foundationsDistributed matrix operations · Shoeybi et al. (2019) Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
- Common trapBelieving tensor parallelism requires inter-GPU communication after every matrix operation; Megatron-LM pairs column and row splits to require only 2 All-Reduces per block.
- Next connectionNext, see how pipeline parallelism stitches layers across multiple servers when a model exceeds 8 GPUs.
More context & mechanics
In a Transformer layer, Feed-Forward Networks execute Y = GELU(X A) B. Splitting A and B naively requires an expensive communication step after every GEMM. Megatron-LM (Shoeybi et al., 2019) eliminates intermediate communication via conjugate column and row slicing:
- Column Parallel Layer (A): Weight matrix A = [A₁, A₂] is split along columns across GPU 1 and GPU 2:
Y₁ = GELU(X A₁), Y₂ = GELU(X A₂)
Because GELU is element-wise: [Y₁, Y₂] = GELU([X A₁, X A₂]). Zero communication is required!
- Row Parallel Layer (B): Weight matrix B = B₁ \ B₂ is split along rows:
Z = Y₁ B₁ + Y₂ B₂
GPU 1 computes Y₁ B₁ and GPU 2 computes Y₂ B₂. An All-Reduce Sum communication primitive combines the two halves: Z = AllReduce-Sum(Yᵢ Bᵢ).
Similarly, attention projections are split across heads:
- W_Q, W_K, W_V are column-sliced (each GPU hosts H/N heads).
- Attention scores and values are computed independently inside each GPU.
- Output projection W_O is row-sliced, followed by a single All-Reduce.
Total communication per transformer block: exactly 2 All-Reduce operations in the forward pass!
Forms the bedrock of distributed multi-GPU training and inference across NVIDIA Megatron-LM, vLLM, TensorRT-LLM, and NeMo Megatron, driving linear throughput scaling over high-speed NVLink interconnects.
Introduced by Mohammad Shoeybi, Mostofa Patwary, and the NVIDIA Applied Deep Learning Research team in 2019 (Megatron-LM). It broke the single-GPU memory wall by exploiting the duality between column-split and row-split linear operators.
"Column slice the expansion, row slice the projection, let All-Reduce reunite the sum." Element-wise activations need zero cross-talk.
How does One-Forward-One-Backward (1F1B) scheduling eliminate pipeline bubbles and bound activation memory in multi-node training?
An automobile assembly line does not wait for one entire car to finish painting before welding the next chassis. Workers pass parts down the line continuously so nobody stands idle waiting for work.
- $$F_{\text{bubble}} = \frac{P - 1}{M + P - 1}$$
- Learning goalDerive the pipeline bubble fraction formula and demonstrate how 1F1B bounds memory consumption to P stages.
- Helpful foundationsPipeline scheduling · Narayanan et al. (2021) Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
- Common trapAssuming pipeline parallelism requires infinite GPU memory for large batches; 1F1B bounds peak activation memory strictly to P micro-batches.
- Next connectionNext, see how ZeRO partitions optimizer states and gradients across data-parallel ranks to unlock 100B+ models.
More context & mechanics
When training models across P pipeline stages, placing layers 1 ṡ L/P on GPU 0, layers L/P+1 ṡ 2L/P on GPU 1, etc., naive batch processing causes idle time (the pipeline bubble F_(bubble)):
Bubble Fraction (Naive) = (P - 1) / (M)
where M is the number of micro-batches (M ≫ P).
The 1F1B (One-Forward-One-Backward) Schedule:
Under the 1F1B schedule (Narayanan et al., 2021):
- Warmup Phase: Stage p executes P - p forward micro-batches to prime the pipeline.
- Steady State: Each GPU alternates exactly one backward step for every forward step (1F1B).
- Cooldown Phase: Finishes remaining backward passes.
The Memory Breakthrough:
In naive pipelining, all M forward activations must be stored until backward execution begins (O(M) memory). 1F1B releases activation memory immediately after each backward step, strictly bounding peak memory to O(P) micro-batches regardless of total batch size M!
Standard execution schedule for training frontier models with hundreds of layers across multi-node clusters in Megatron-Core, DeepSpeed Pipeline, and PyTorch PiPPy.
GPipe (Huang et al., Google Brain 2019) introduced micro-batch pipelining but suffered huge memory footprint. Narayanan et al. (Stanford and NVIDIA 2021) designed the 1F1B schedule in PipeDream and Megatron-LM, capping memory overhead to the pipeline depth P.
"Prime with forward steps, cruise on 1F1B, cool down with backward passes: memory caps at depth, throughput flows like clockwork."
How does the Zero Redundancy Optimizer (ZeRO) eliminate duplicated optimizer states, gradients, and parameters without increasing communication?
If four chefs are preparing soup, they do not each need to buy their own separate 50-pound sack of salt. One chef holds the salt, another holds the pepper, and they share spoons as needed, saving 75% of kitchen counter space.
- $$\text{Memory}_{\text{ZeRO-3}} = \frac{2\Phi + 2\Phi + 12\Phi}{N_d} = \frac{16\Phi}{N_d}$$
- Learning goalFormulate the memory footprint across ZeRO Stages 1, 2, and 3, and explain the dynamic parameter gather-and-release mechanism.
- Helpful foundationsOptimizer states & AdamW memory · Rajbhandari et al. (2020) ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
- Common trapBelieving ZeRO-3 increases total training communication; communication volume matches standard data parallelism within a factor of 1.5x.
- Next connectionProceed to Chapter 48 to master frontier FP8 training, mixture of depths, and long-context needle retrieval.
More context & mechanics
In mixed-precision training (FP16/BF16 weights and FP32 Adam optimizer states), training a model with Φ parameters requires 16 Φ bytes of baseline memory:
- FP16 Parameters (2 Φ) + FP16 Gradients (2 Φ)
- FP32 Master Weights (4 Φ) + FP32 Momentum (4 Φ) + FP32 Variance (4 Φ) = 12 Φ bytes.
In standard Data Parallelism, all 16 Φ bytes are duplicated across every GPU.
The Three ZeRO Stages (Rajbhandari et al., 2020):
- ZeRO-Stage 1 (P_(os)): Optimizer State Partitioning. Adam states are partitioned across N_d GPUs (4Φ + 12Φ / N_d). Memory reduction: 4×.
- ZeRO-Stage 2 (P_(os+g)): Gradient Partitioning. Each GPU updates only its local gradient slice (2Φ + 14Φ / N_d). Memory reduction: 8×.
- ZeRO-Stage 3 (P_(os+g+p)): Parameter Partitioning. Model parameters are partitioned; layers gather weights dynamically via All-Gather and discard them immediately (16Φ / N_d). Memory reduction: Proportional to N_d!
Training a 100B parameter model across 64 GPUs:
- Standard Data Parallelism: 16 × 100B = 1,600 GB per GPU (Impossible on an 80GB GPU).
- ZeRO-Stage 3 on 64 GPUs: 1600 / 64 = 25 GB per GPU (Fits easily within a single 80GB A100/H100 GPU!).
Powering training and fine-tuning workloads worldwide through DeepSpeed ZeRO and PyTorch FSDP (Fully Sharded Data Parallel), enabling open-source research labs to train 70B+ models on commodity GPU clusters.
Invented by Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He at Microsoft Research in 2020 (ZeRO: Memory Optimizations Toward Training Trillion Parameter Models). It proved that memory redundancy could be eliminated without increasing communication volume.
"Stage 1 shards optimizers, Stage 2 shards gradients, Stage 3 shards parameters: zero waste, maximum scale." Divide state by rank to conquer trillion-parameter frontiers.
How does Distributed Scaling: Pipeline, Tensor, and Data Parallelism (3D Parallelism) operate mathematically at scale?
A 3D cube representing an 8x8x8 GPU cluster (512 GPUs):
- Learning goalMaster the algebraic and geometric derivation of Distributed Scaling: Pipeline, Tensor, and Data Parallelism (3D Parallelism).
- Helpful foundationsFoundations of Build Llm From Scratch · Coordinate Invariance & Geometric Projections · Radford et al. (2019), Language Models are Unsupervised Multitask Learners, OpenAI
- Common trapAssuming Distributed Scaling: Pipeline, Tensor, and Data Parallelism (3D Parallelism) is an arbitrary computational trick rather than an invariant spatial law.
- Next connectionConnects upstream foundational lemmas to modern high-dimensional computational systems.
More context & mechanics
A 3D distributed systems cluster architecture diagram in square 1:1 format on dark slate (#0f172a) with a 3D isometric cube.
Top section: Title in white sans-serif: "3D Parallelism: Scaling to Trillions of Parameters". Subtitle in muted cyan: "Coordinating Data Parallelism (DP), Tensor Parallelism (TP), and Pipeline Parallelism (PP)."
Middle section: A 3D cube representing an 8x8x8 GPU cluster (512 GPUs):
- Tensor Parallelism (TP - Intracore, Fast NVLink):
Splits individual weight matrices across GPUs within the same server node (Column-parallel QKV, Row-parallel output). Requires 900 GB/s NVLink.
- Pipeline Parallelism (PP - Inter-node, Medium Bandwidth):
Splits layers sequentially across server racks (Layers 1-8 on Rack 1, Layers 9-16 on Rack 2). Uses 1F1B micro-batch scheduling to minimize pipeline bubble.
- Data Parallelism (DP / ZeRO-3 - Inter-cluster, High Scale):
Replicates the entire model pipeline across multiple clusters, sharding batches and averaging gradients via All-Reduce.
Bottom section: Dimension equation: "Total GPUs = TP × PP × DP. Enables training models exceeding the memory capacity of any single GPU by 100x."
Stunning 3D isometric cluster rendering, clean color-coded axes, modern systems styling.
Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.
Formal derivation and peer-reviewed analysis:
Radford et al. (2019), Language Models are Unsupervised Multitask Learners, OpenAI
"Distributed Scaling: Pipeline, Tensor, and Data Parallelism (3D Parallelism): Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."