Parallel and Distributed Inference: Megatron-LM / GPipe / ZeRO

When model parameters exceed a single GPU's memory (a 70B model's FP16 weights alone are ~140GB; an A100 has 80GB), there's only one answer: split the model and let multiple GPUs carry it together. Different splitting directions gave rise to the three papers here — Megatron-LM (within-layer split: tensor parallelism), GPipe (between-layer split: pipeline parallelism), and DeepSpeed ZeRO (state split: parameter/gradient/optimizer partitioning). They began as training techniques, but today they are standard equipment for large-model inference.

Before reading, skim GPUs and Hardware Selection (NVLink, the memory hierarchy) and LLM Inference Optimization.

Parallelism Strategies at a Glance

Strategy	What gets split	Problem it solves	Best for
Data parallelism (DP)	Data spread across GPUs, each holding a full model	Training throughput (parallelism)	Mostly training; not used for inference (duplicated weights per GPU are waste)
Tensor parallelism (TP)	Within a layer, matrices split by row/column across GPUs	A single layer too large for one GPU	Inference + training (vLLM's -tp flag)
Pipeline parallelism (PP)	Between layers, layers assigned to GPUs forming a pipeline	The whole doesn't fit; split by layers to share the load	Inference + training (vLLM's -pp flag)
ZeRO	Training state (optimizer/gradients/parameters) partitioned	Training memory redundancy; on inference = parameter partitioning/offload	Mostly training; inference uses ZeRO-Inference variants

The parallelism formula

What fits on one card = per-card memory × number of cards ÷ parallelism overhead. Common deployment combos: DP × PP × TP (three-dimensional parallelism). The Megatron paper calls this "1D/2D/3D parallelism" — TP alone is 1D, TP+PP is 2D, TP+PP+DP is 3D.

1. Megatron-LM: Within-Layer Tensor Parallelism (Shoeybi et al., 2019)

The contribution in one sentence

Within-layer tensor parallelism with a handful of communication primitives: split the Transformer layer's matrix multiplications by row/column across GPUs, adding only two all-reduces per layer — enough to train/run models whose single layers don't fit on one card; 8.3B parameters across 512 GPUs reached 15.1 PFLOPS at 76% scaling efficiency.

Background and motivation

In 2019, parallel training meant data parallelism — but once a model gets large enough, a single layer's weights plus activations won't fit on one card, and data parallelism is helpless (every card needs the full model). Megatron chose to "split within the layer": given the matrix multiply Y = X·W, cut W into two pieces by column, or by row, let different GPUs compute each half, then combine — no new compiler required, just a few communication calls inserted into PyTorch.

The method in brief

Using the Transformer's MLP and attention as examples:

Tensor parallelism: split weight matrix W into two pieces, X unchanged

Y = X · W        ┌──────────┬──────────┐
                 │ W1 (col) │ W2 (col) │      → Y = [Y1, Y2]
                 └──────────┴──────────┘
 GPU0: Y1 = X·W1   GPU1: Y2 = X·W2   (f split: compute, then communicate to join)

Z = W · X        ┌──────────┐
                 │ W1 (row) │             → Z = Z1 + Z2 (each computes, then all-reduce the sum)
                 ├──────────┤
                 │ W2 (row) │
                 └──────────┘
 GPU0: Z1 = W1·X   GPU1: Z2 = W2·X  (g split: communicate first, then compute)

• The f / g splits alternate: the column split (f) joins outputs by concatenation; the row split (g) sums them via all-reduce — guaranteeing only two all-reduces per Transformer layer, with communication that can overlap computation.
• 1D → 2D → 3D: tensor parallelism (within layer), pipeline parallelism (between layers), and data parallelism (data) compose orthogonally; the paper trained with configurations like 8-GPU TP × 64-GPU DP.
• Supporting engineering: fused operators (bias+gelu fusion), scaled masked softmax, communication/computation overlap.

Key results

Metric	Number
Scale	8.3B-parameter GPT-2-style model, trained on 512 GPUs
Throughput	15.1 PFLOPS (end-to-end); single-GPU baseline 39 TFLOPS (~30% of peak)
Scaling efficiency	76% (versus a strong single-GPU baseline, at 512 GPUs)
Quality	WikiText-103 perplexity 10.8, LAMBADA 66.5% — SOTA at the time

Limitations

• Communication grows with TP degree: two all-reduces per layer — push TP too high and communication bites back. So TP is generally confined within a single node (where NVLink bandwidth is high), with PP spanning nodes.
• Requires each layer to divide cleanly into integer pieces (head counts, hidden dimensions must align).

What it means today

• The first choice on the inference side: vLLM, TensorRT-LLM, and DeepSpeed's tensor-parallel-size are exactly this; a 70B model runs on 2–4 A100s with TP alone — see the vLLM case study.
• The decision rule: a single layer doesn't fit → you need TP; too many layers overall → TP+PP.

2. GPipe: Between-Layer Pipeline Parallelism (Huang et al., 2018)

The contribution in one sentence

Cut the model by layers into K segments, one per GPU, and fill the pipeline with micro-batches — by "splitting into micro-batches," parallelism decouples from the model's batch size; a 6B-parameter, 128-layer Transformer approaches linear speedup.

Background and motivation

Splitting by layers (layers 1–32 on GPU0, 33–64 on GPU1...) looks simple, but the naive "send the whole batch through the pipeline once" produces a severe bubble: each GPU spends most of its time waiting for upstream, and utilization is dismal. GPipe's key improvement: split the batch into multiple micro-batches flowing through the pipeline, keeping every GPU busy simultaneously.

The method in brief

K=4 GPUs, M=8 micro-batches, pipelined execution (starting from layer 0):

time →
GPU0: [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
GPU1:      [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
GPU2:           [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
GPU3:                [mb1][mb2][mb3][mb4][mb5][mb6][mb7][mb8]
                      └──────── tail bubble ────────┘

The more micro-batches, the fuller the pipeline, the smaller the idle (bubble) share

• Bubble overhead = (K-1)/(M+K-1): fewer GPUs K and more micro-batches M shrink the bubble share (approaching 0 as M→∞).
• Micro-batches decouple "parallelism" from "batch size": previously, parallelism 8 demanded batch=8; now batch=1 split into 8 micro-batches can still occupy 8 GPUs.
• Activation re-materialization (recomputation) trims peak training memory.

Key results

Metric	Number
Models	6B-parameter, 128-layer Transformer (multilingual translation); 557M AmoebaNet (84.4% top-1 on ImageNet)
Speedup	Near-linear after micro-batching (bubbles filled)
Applicability	Works for any "sequence of layers" model — architecture-agnostic

Limitations

• Bubbles can't be fully eliminated (the head and tail never fill).
• On the training side, activation recomputation costs some speed; on the inference side, PP adds first-token latency (the request must "flow through" the whole pipeline).
• Uneven per-layer compute (embedding/output layers vs. middle layers) causes load imbalance.

What it means today

• On inference, PP combines with TP: TP solves "one layer doesn't fit," PP solves "too many layers"; vLLM supports PP.
• Deployment caveats: PP's added first-token latency and cross-node communication bottlenecks mean small-scale setups should prefer TP.
• The bubble analysis ((K-1)/(M+K-1)) is the first back-of-envelope calculation for evaluating any pipeline scheme.

3. DeepSpeed ZeRO: State Partitioning (Rajbhandari et al., 2019)

The contribution in one sentence

Removing data parallelism's "every GPU stores everything" redundancy: partitioning optimizer states, gradients, and parameters in three successive stages (Zero Redundancy) so memory scales linearly with GPU count — 400 GPUs training 100B+ parameters at 15 PFLOPS with superlinear speedup.

Background and motivation

Data parallelism looks simple, but every GPU stores the full model plus full optimizer states — for a 7.5B model, pure data parallelism needs ~120GB per card, of which model weights are only 30GB; the redundancy dominates (75GB optimizer states, 30GB gradients). ZeRO's insight: that redundancy isn't necessary — it can be partitioned like model parallelism while keeping data parallelism's communication pattern (low volume, large granularity).

The method in brief: three stages

Stage 1: partition optimizer states → each GPU stores 1/N of optimizer state (memory ÷ N)
Stage 2: + partition gradients     → each GPU stores 1/N of gradients (memory ÷ N)
Stage 3: + partition parameters    → each GPU stores 1/N of parameters (memory ÷ N)
         (all-gather a layer's parameters only when computing it; discard after use)

Stage 3 makes per-GPU memory ÷ N: 100 cards = 100× the trainable scale
The inference-side counterpart: ZeRO-Inference / offloading parameters to CPU

• Keeps data parallelism's communication pattern: not model parallelism's per-layer all-reduce, but "all-gather on demand, discard when done" — low communication volume, simple to implement.
• ZeRO-Offload: pushes parameters/optimizer states further onto CPU/memory, letting even a single GPU train huge models (on inference: "weights on CPU, compute on GPU" offloading).

Key results

Metric	Number
Scale	400 GPUs training 100B+ parameters, throughput 15 PFLOPS
Speedup	Superlinear scaling (batch-size gains from freed memory outweigh communication costs)
Versus SOTA	8× larger models, 10× higher performance
Usability	13B parameters trainable without model parallelism (ZeRO alone)

Limitations

• Stage 3's on-demand all-gather is communication-heavy in low-bandwidth environments (superlinear gains depend on fast interconnects).
• Primarily a training technique; applying ZeRO directly to inference has limited value (inference stores no gradients/optimizer states) — use ZeRO-Inference (weight partitioning + offload) instead.

What it means today

• On inference: ZeRO-Inference can partition weights across GPUs or offload them to CPU memory, letting 70B-class models run on 2 cards or even 1 card + CPU memory — trading speed for "it fits."
• The decision rule: need throughput → TP/PP (weights resident in VRAM, low latency); need "it fits" → ZeRO partitioning/offload (slower, but it runs).

4. Comparison Table

Method	Communication cost	Memory effect	Implementation complexity	Best for
Tensor parallelism (Megatron)	2 all-reduces per layer, growing with TP degree	Single-layer weights/activations spread across GPUs	Medium (needs operator-level splitting support)	First choice for inference: single layer doesn't fit; vLLM -tp
Pipeline parallelism (GPipe)	Only activations passed between stages; low communication	Each GPU stores only some layers	Medium (needs layer splitting + micro-batching)	Training + large-model inference (stacked with TP)
ZeRO (DeepSpeed)	Stage 3 on-demand all-gather, high	Parameters/gradients/optimizer split linearly across GPUs	High	Mostly training; inference via offload variants

5. Parallelism in Inference Practice

The common combo for large-model inference is TP + PP (supported natively by vLLM):

• Single node, layers fit: pure TP (e.g. 2–4 cards within NVLink range) — fast communication, low latency.
• Model too large / multiple nodes: TP (within node) × PP (between nodes); only activations cross nodes, sidestepping the inter-node communication bottleneck.
• Severely memory-starved: ZeRO-Inference / offload — weights live in CPU memory, fetched by the GPU on demand; low throughput but "it runs."
• Deployment details and flags (--tensor-parallel-size, --pipeline-parallel-size): the vLLM case study.

More parallelism isn't automatically better

Doubling TP cards doubles communication; doubling PP stages doubles bubbles and first-token latency. Add parallelism only until "the model fits and latency is acceptable" — extra cards are better spent running additional replicas (DP) to take traffic, which often yields higher throughput. That's the horizontal-scaling idea from Deployment Architecture Patterns.

6. Shared Limitations

1. Communication bottlenecks: inter-node bandwidth (25/50/100 Gbps NICs) is far below NVLink (~600GB/s class); TP across nodes slows markedly — keep TP within the node is the iron rule.
2. Load imbalance: uneven per-layer compute in PP and indivisible TP splits both create a "weakest plank" effect; splits must be tuned finely by compute/memory.
3. Complexity: scheduling, fault tolerance, and memory planning for 3D parallelism are a high bar for in-house engines — which is why mature frameworks like vLLM, TensorRT-LLM, and DeepSpeed come first. See Choosing Frameworks and Platforms.

Further Reading

• LLM Inference Optimization — the complete picture of KV cache, scheduling, and parallelism in LLM inference
• The vLLM case study — real configurations and selection for -tp/-pp
• PagedAttention: The vLLM Paper — how memory management stacks with parallelism
• GPUs and Hardware Selection — how NVLink and memory bandwidth determine the parallelism plan
• Serving Systems: Clipper / Orca / Nexus — the scheduling perspective on distributed serving
• Deployment Architecture Patterns — why horizontal scaling comes first when one GPU suffices

References

• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (arXiv 1909.08053)
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (arXiv 1811.06965)
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arXiv 1910.02054)
• DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale (arXiv 2207.00032)