Tensor, Pipeline, and Expert Parallelism

Using more GPUs can make a model fit, reduce the work each GPU performs, or serve more independent requests. Those goals lead to different layouts. Splitting a model also adds communication, so a layout that saves memory can make a single request slower.

The memory budget gives the first constraint. A rounded 70B model has about 140 GB of raw 16-bit weights or 35 GB of raw four-bit weights. The latter can leave substantial capacity on an 80 GB device, although metadata, workspace, and KV cache still need room. Parameter count alone does not tell you how many GPUs to use.

Tensor parallelism: distinguish columns from rows

Write a linear layer using row-vector convention, $y=xW$.

If $W$ is partitioned by output columns, each GPU can receive the same $x$ and compute a different slice of $y$. Those slices concatenate to form the complete output. They are different coordinates, so summing them would be wrong.

If $W$ is partitioned by input rows, each GPU multiplies the matching slice of $x$ by its rows of $W$. These results are contributions to the same output coordinates, so they must be summed, often with an all-reduce.

A common Megatron-style feedforward layout pairs a column-parallel expansion with a row-parallel contraction. The intermediate activation remains sharded, avoiding a gather between them. There is a reduction after the contraction. A similarly arranged attention block adds another reduction after its output projection: typically two forward collectives per full transformer block in this layout, not one for every matrix.

Implementations can use different collective combinations, sequence parallelism, or fused communication. Some tensors remain replicated. KV heads can also be replicated when the tensor-parallel degree exceeds the model's KV-head count.

Pipeline parallelism: different layers, overlapping work

Pipeline parallelism places consecutive layer groups on different devices. Activations move from one stage to the next, while each stage keeps the KV cache for its own attention layers.

One request traversing four stages leaves three stages idle while the current stage works. Independent microbatches can overlap: stage 0 handles B while stage 1 handles A. A later generated token from the same sequence still depends on sampling the previous token after the final stage; it cannot be treated as an independent microbatch.

If there are $P$ equally fast stages and $M$ microbatches, each taking one time unit per stage, a forward pass takes $M+P-1$ units. The idle fraction across that window is

With four stages and four microbatches this is $3/7$, about 43%. Sixteen microbatches reduce it to $3/19$, about 16%, under the same assumptions. Communication, unequal layers, changing batch shapes, and autoregressive feedback complicate a real inference schedule.

Expert parallelism: send activations to the chosen experts

In a mixture-of-experts model, a router chooses a subset of expert networks for each token. Expert parallelism places those networks on different GPUs. Tokens' hidden states are dispatched to their selected experts, then the expert outputs return for combination.

Dispatch and return commonly use all-to-all communication. If routing is uneven, one GPU may receive much more work than another. Expert replication and placement can help balance execution while preserving the model's calculation; changing routing decisions or dropping tokens changes the computation and needs separate justification.

An expert need not occupy an entire GPU. A GPU can hold several experts, and a large expert can itself use tensor parallelism.

Combine axes around the hardware

A TP=4, PP=2 layout uses eight GPUs: two pipeline stages, each containing a four-GPU tensor-parallel group. Tensor-parallel communication stays within each group; pipeline activations cross between stages.

For 140 GB of raw 16-bit weights, perfect eight-way sharding would average 17.5 GB per GPU. Real placement also includes replicated tensors, stage imbalance, communication buffers, and KV entries. An eight-GPU capacity sum alone cannot prove that every individual GPU fits.

Fast intra-node links often make tensor parallelism attractive within a node. Pipeline parallelism communicates at fewer boundaries and can be useful across nodes, depending on message size, latency, topology, batch size, and overlap. H100's advertised 900 GB/s NVLink figure is aggregate bidirectional bandwidth, not an effective one-way rate for every collective.

For a ring all-reduce over $P$ ranks and a payload of $D$ bytes per rank, the standard bandwidth model sends roughly $2(P-1)D/P$ bytes per rank. Expert dispatch volume also depends on top-k and how many assignments stay local.

Serving replicas are another option

If a model fits on one GPU and one-request latency is acceptable, several independent replicas can increase aggregate throughput without model-parallel communication between them. This is serving data parallelism: replicas process different requests and do not synchronize gradients.

A larger tensor-parallel group is not equivalent to more replicas. It spends more devices on each model execution. Benchmark both layouts under the intended request distribution before choosing.

A concrete launch shape

The following are alternative vLLM server launches for the same checkpoint. They require compatible GPUs with enough memory, access to the model, and an installed vLLM version supporting the selected architecture and backend.

# Four tensor-parallel ranks in one pipeline stage.
vllm serve meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 4 --pipeline-parallel-size 1
 
# Alternative: two stages, with two tensor-parallel ranks per stage.
# Stop the first server before launching this configuration.
vllm serve meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 2 --pipeline-parallel-size 2

A multi-node deployment also needs a configured distributed executor, networking, and model availability on its workers.

The next post connects those choices to the throughput and latency tradeoff.

Additional reading

• Megatron-LM, Shoeybi et al. Column/row tensor partitions and communication placement.
• GPipe, Huang et al. Pipeline microbatching; its training schedule should not be confused with the forward-only illustration here.
• vLLM distributed serving. Tensor and pipeline launch configuration.
• NVIDIA NCCL collective operations. All-reduce, all-gather, and related operations.
• NVIDIA H100 specifications. Hardware bandwidth definitions.