Parallelism for serving: tensor, pipeline and expert parallelism

Parallelism for serving is how one model is spread over several GPUs at inference time. There are four basic moves. Tensor parallelism (TP) splits every weight matrix across GPUs so that each layer runs on all of them at once. Pipeline parallelism (PP) gives each GPU a contiguous slice of layers and passes activations from one stage to the next. Expert parallelism (EP) places the experts of a mixture-of-experts (MoE) layer on different GPUs and routes each token to wherever its experts live. Data parallelism (DP) runs independent replicas. Training uses the same vocabulary, but serving optimises for different things: there are no gradients or optimiser states, the batch per step is set by live traffic, and the user waits for every decode step. A post-trained model inherits its base model's shapes, so its parallel layout is mostly decided by the base; what post-training changes is which layouts are safe (LoRA sharding, quantization group sizes, MoE router balance after fine-tuning) and whether outputs are reproducible across layouts.

The constraints: memory first, then the decode step#

A layout must first fit. With $n$ GPUs of memory $M$, weights $W_{\text{bytes}}$, per-token KV-cache bytes $\kappa$ and $T$ tokens of context held across all running requests,

for TP or PP over $n$ GPUs (each GPU holds $1/n$ of the weights and of each token's KV), with $L$ layers, $h_{\text{kv}}$ KV heads of dimension $d_{\text{head}}$ and $b$ bytes per value. Then it must be fast. For decode, which is memory-bound (see the roofline), the step time under TP of degree $n$ is roughly

The first term shrinks as $1/n$, because every GPU reads only its shard of the weights. The second is the price: two all-reduces per layer of the activations $B\times d$, modelled as a ring with per-hop latency $\alpha$ and link bandwidth ${\beta }_{\text{link}}$. For decode, $S=Bdb$ is tiny (64 tokens × 8192 × 2 bytes = 1 MB), so the latency term $2(n-1)\alpha$ dominates, and it is paid $2L$ times per step. TP therefore helps latency until the communication term catches up with the shrinking weight term, and it is only viable over fast intra-node links.

Where the techniques came from#

Shoeybi et al.'s Megatron-LM (2019) introduced the intra-layer split that defines TP for transformers: the first MLP matrix split by columns, the second by rows, attention split by heads, so that each block needs one all-reduce in the forward pass. Huang et al.'s GPipe (2018) defined pipeline parallelism with micro-batches and characterised its bubble. Lepikhin et al.'s GShard (2020) and Fedus, Zoph and Shazeer's Switch Transformers (2021) made MoE layers with experts sharded across devices practical, with all-to-all communication for routing and capacity factors for load. For inference specifically, Rajbhandari et al.'s DeepSpeed-MoE (2022) combined expert, tensor and data parallelism for MoE serving, and Pope et al. (2022), "Efficiently Scaling Transformer Inference", analysed partitioning layouts for latency versus throughput on TPU v4 and showed that the best layout differs between prefill and decode. The DeepSeek-V3 report (2024) is the most detailed public description of a large MoE serving layout: prefill on units of 32 GPUs with TP4 and DP8 for attention and EP32 for experts, and decode on units of 320 GPUs with TP4 and DP80 for attention, EP320 for the MoE layers with one expert per GPU, and 64 GPUs hosting redundant copies of high-load experts and the shared expert, rebalanced from live statistics about every ten minutes.

Tensor parallelism, step by step#

For an MLP $Y=\varphi (XA)B$, split $A=[A_1,\dots ,A_n]$ by columns and $B=[B_1;\dots ;B_n]$ by rows. GPU $i$ computes $Y_i=\varphi (X{A}_i)B_i$ with no communication, because $\varphi$ is applied elementwise to its own columns, and $Y={\sum }_i{Y}_i$ is one all-reduce. Attention splits by heads: each GPU owns $h/n$ query heads and their K and V heads, computes attention locally, and the output projection is row-parallel. With grouped-query attention there are only $h_{\text{kv}}$ KV heads (8 in most current dense models), so TP above $h_{\text{kv}}$ must replicate KV heads, which stops shrinking the KV cache per GPU. Embeddings and the LM head are split along the vocabulary.

What the model predicts

For a 72B dense model (Qwen2.5-72B shapes: 80 layers, hidden size 8192, 72.7B parameters, 145 GB in bf16) and assumed link parameters, the weight term at batch 1 is 21.7 ms at TP2, 10.9 ms at TP4 and 5.4 ms at TP8. With NVLink-class links (450 GB/s, 2 µs per hop) the all-reduce term is 0.65, 1.9 and 4.5 ms, so the total falls from 22.4 to 12.8 to 9.9 ms: each doubling helps less. With PCIe-class links (25 GB/s, 5 µs) the communication at TP8 (11.4 ms at batch 1, 22.9 ms at batch 64) exceeds the weight time it saves, and TP8 is no faster than TP4. These are model outputs from assumed constants; the shape, diminishing returns over NVLink and a reversal over slow links, is what to expect on real hardware, and why TP almost never crosses node boundaries.

Pipeline parallelism#

With $p$ stages each holding $L/p$ layers, a token's forward pass visits every stage in sequence, so PP does not reduce per-token latency at all; it only adds the point-to-point transfer of a $B\times d$ activation between stages, which is small and tolerates slow inter-node links. What PP buys is capacity: memory per GPU falls as $1/p$, like TP, without a collective on every layer. To keep all stages busy the scheduler must have $m\ge p$ micro-batches in flight; with a synchronous schedule the idle fraction is the GPipe bubble

so 4 stages with 4 micro-batches idle 43% of the time and with 16 micro-batches 16%. Serving engines avoid the synchronous schedule by keeping different requests' decode steps in different stages at once, which works when there is enough concurrent traffic and fails at low load, where a single request's step must traverse all stages while the others sit idle. The usual pattern is TP within a node and PP across nodes, for models that do not fit in one node.

Expert parallelism and MoE serving#

An MoE layer replaces the dense MLP with $E$ experts and a router that sends each token to its top $k$. Under EP over $G$ GPUs each holds $E/G$ experts, and every MoE layer does two all-to-alls: dispatch (send each token's hidden state to the GPUs owning its experts) and combine (send expert outputs back and sum them with the router weights). Per layer, each token moves about $kdb$ bytes each way, so communication scales with tokens and $k$, not with expert size.

Two properties make MoE serving different from dense serving. The first is that the bytes a decode step reads depend on the batch. With $B$ tokens routed uniformly, the expected fraction of experts touched is

Using the expert counts in the models' own config files: for Mixtral 8x7B ($E=8,k=2$) a batch of 16 already touches 99% of experts; for DeepSeek-V2-Lite ($E=64,k=6$) batch 32 touches 95.7%; for DeepSeek-V3 ($E=256,k=8$) batch 64 touches 86.9% and batch 256 essentially all. At batch 1 an MoE model reads only its active parameters and decodes like a much smaller dense model; at serving batch sizes it reads nearly everything, so per-step memory traffic approaches that of a dense model of the same total size while compute stays at the active size. That is why MoE models are cheap per token only at high batch, where the large weight read is shared by many tokens, and why their serving layouts aim to give each expert a large enough batch (DeepSeek-V3's stated reason for EP32 in prefill).

The second property is load imbalance. Routers are not uniform, and the slowest GPU sets the step time. We simulated 64 experts with top-6 routing over 8 GPUs (8 experts each), with expert popularities drawn from a lognormal distribution ($\sigma =0.7$, seed 0; an assumption, not measured router statistics), sampling top-6 per token by Gumbel-perturbed popularity.

The imbalance does not average away with larger batches: above a few hundred tokens it converges to the imbalance of the popularity distribution itself (1.43 here), meaning the busiest GPU does 43% more work than average and the others wait. Replicating the 8 hottest experts onto a second GPU each and splitting their traffic cut it to 1.23. This is the mechanism behind DeepSeek-V3's redundant experts. Fine-tuning can change which experts are hot, so the statistics must come from the post-trained model's live traffic.

Production MoE layouts usually combine axes: attention (dense, with a per-request KV cache) runs data-parallel or with small TP, while expert layers run with large EP, so each GPU has its own requests' KV cache and a share of the experts. Prefill and decode often get different layouts, which pairs naturally with prefill-decode disaggregation.

Worked example: a 72B model on 8 GPUs#

With 80 GB GPUs and 72 GB usable, Qwen2.5-72B-shaped weights take 145 GB. Its KV cache costs $\kappa =2\times 80\times 8\times 128\times 2=327,680$ bytes per token in bf16. Three layouts on one 8-GPU node:

TP8 gives the lowest latency and the largest single KV pool, which matters for very long contexts. Two TP4 replicas give about 58% more throughput at the same batch per replica, because they spend less of each step in all-reduces, at the cost of 30% higher per-token latency and two smaller KV pools. If the latency SLO allows 13 ms per output token, the replicas win; if it demands 10 ms, only TP8 meets it (see latency SLO tuning). This trade, a bigger TP group for latency or more replicas for throughput, is the central decision in serving parallelism.

Measured: tensor parallelism changes the outputs#

TP is often described as exact. In floating point it is not, because the row-parallel layers sum partial results in a different order and with extra roundings. We emulated this on an Apple-silicon Mac (PyTorch 2.10, MPS) with Qwen2.5-0.5B-Instruct in bf16: the o_proj and down_proj layers were split into $n$ input shards, each shard's partial product rounded to bf16 and the partials summed in bf16, as a bf16 all-reduce does. Column-parallel layers were left alone since they do not change arithmetic. We compared against the unsplit model on 60 held-out chat prompts (teacher-forced on the model's own answers) and on 96-token greedy generations. Two unsplit runs were bitwise identical, so every difference below comes from the split.

The per-token distributions barely moved (KL about 0.001), yet only a fifth to a quarter of greedy generations matched the single-GPU output token for token, because one near-tie flipped early and the continuation diverged from there. For post-trained models this matters in three places: evaluation numbers measured under one layout will not reproduce exactly under another; RL pipelines that sample rollouts with one TP degree and compute log-probabilities with another have a small but real sampler-learner mismatch; and regression tests that compare greedy outputs byte for byte will fail after a re-sharding without any real change. Pin the serving layout in eval configs, or compare distributions rather than strings.

Interaction with post-trained models#

LoRA under TP. Adapters must be sharded consistently with their base layer. For a column-parallel base, $A$ is replicated and $B$ is split by rows; for a row-parallel base, $A$ is split by columns and its partial products need a small extra reduction before $B$. S-LoRA describes such a scheme to keep the added communication small; it is why multi-LoRA engines support TP but sometimes not every adapter target under every layout (see multi-LoRA serving). Quantized checkpoints. Group-wise INT4 weights need each shard's input dimension to be a multiple of the group size; a 128-group checkpoint with a 14,336-wide MLP splits cleanly at TP8 (1,792 per shard, 14 groups), but Qwen2.5-72B's 29,568-wide MLP gives 3,696 per shard at TP8, which is not a multiple of 128, so such a checkpoint needs padding, a smaller group, or a different TP degree (see GPTQ and AWQ). Extended vocabularies. Added tokens change the embedding and LM-head shapes, which must stay divisible by the TP degree; pad the vocabulary when you extend the tokenizer. MoE fine-tunes. Fine-tuning shifts router preferences, so redundancy and placement decisions computed on the base model are stale; collect expert-load statistics from the fine-tune. Fine-tunes that drift toward a few experts hurt EP throughput even when quality improves. Reproducibility. As measured above, use the same layout for evaluation and production, and for RL sampling and scoring where possible.

Tuning knobs#

When each kind of parallelism helps#

Failure modes and how to detect them#

Communication-bound decode. Raising TP does not lower inter-token latency; a profiler shows all-reduce kernels taking a large share of each step. Fix with a smaller TP and more replicas, faster links, or custom all-reduce. Pipeline bubbles at low load. PP stages idle when there are fewer in-flight requests than stages; GPU utilisation drops at off-peak hours while latency stays high. Expert hot spots. One EP rank is consistently the last to finish its MoE layer; log per-rank token counts per layer. KV replication beyond $h_{\text{kv}}$. TP above the number of KV heads stops increasing KV capacity per token, visible as a smaller-than-expected gain in maximum concurrent context. Layout-dependent outputs that break regression tests or RL importance ratios, as in our measurement. Uneven memory under PP when the first and last stages also hold the embedding and LM head; balance layer counts per stage.

Implementation sketch#

In vLLM each axis is a flag (--tensor-parallel-size, --pipeline-parallel-size, --data-parallel-size, --enable-expert-parallel); SGLang and TensorRT-LLM expose the same axes. A layout search is cheap to script before any benchmarking. In POFE, the serving layout is recorded alongside each evaluation so results can be reproduced under the same sharding.

def plan(params_b, layers, kv_heads, head_dim, hidden, gpu_gb=80, usable=0.9,
         hbm=3.35e12, alpha=2e-6, link=450e9, bytes_w=2, bytes_kv=2):
    kappa = 2 * layers * kv_heads * head_dim * bytes_kv              # KV bytes per token
    out = []
    for tp in (1, 2, 4, 8):
        w = params_b * 1e9 * bytes_w / tp
        free = gpu_gb * 1e9 * usable - w
        if free <= 0: continue                                     # does not fit
        kv_tokens = tp * free / (kappa * max(1, tp // kv_heads))     # KV replicated beyond kv_heads
        for B in (1, 64):
            ar = 0 if tp == 1 else 2*(tp-1)*alpha + 2*(tp-1)/tp * B*hidden*2/link
            t = w / hbm + 2 * layers * ar
            out.append(dict(tp=tp, batch=B, ms=1e3*t, tok_s_per_gpu=B/t/tp, kv_tokens=int(kv_tokens)))
    return out

for row in plan(72.7, 80, 8, 128, 8192): print(row)
# vllm serve MODEL --tensor-parallel-size 4 --data-parallel-size 2          # two TP4 replicas
# vllm serve MOE_MODEL --tensor-parallel-size 1 --data-parallel-size 8 --enable-expert-parallel

Further reading#

Shoeybi et al. (2019) for the Megatron tensor-parallel split; Huang et al. (2018) for pipeline micro-batching and bubbles; Lepikhin et al. (2020) and Fedus et al. (2021) for sharded MoE layers; Rajbhandari et al. (2022) for MoE inference with combined parallelism; Pope et al. (2022) for an analytical treatment of inference partitioning; the DeepSeek-V3 report (2024) for a production MoE layout with redundant experts; Jiang et al. (2024) for Mixtral's architecture. For the economics of these layouts, see utilization and batching economics and cost per million tokens.

Frequently asked questions#

References#

1. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arxiv.org/abs/1909.08053
2. Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, et al. (2018). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. arxiv.org/abs/1811.06965
3. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. arxiv.org/abs/2006.16668
4. William Fedus, Barret Zoph, Noam Shazeer (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arxiv.org/abs/2101.03961
5. Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, et al. (2022). DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale. arxiv.org/abs/2201.05596
6. Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, et al. (2022). Efficiently Scaling Transformer Inference. arxiv.org/abs/2211.05102
7. DeepSeek-AI (2024). DeepSeek-V3 Technical Report. arxiv.org/abs/2412.19437
8. Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, et al. (2024). Mixtral of Experts. arxiv.org/abs/2401.04088
9. Ying Sheng, Shiyi Cao, Dacheng Li, Coleman Hooper, Nicholas Lee, et al. (2023). S-LoRA: Serving Thousands of Concurrent LoRA Adapters. arxiv.org/abs/2311.03285

Related reading