Robots Atlas>ROBOTS ATLAS
Inference

Distributed Inference

2022ActivePublished
Key innovation
Enables inference of models too large to fit on a single accelerator by partitioning the model and its computation across multiple GPU/TPU devices and nodes.
Category
Inference
Abstraction level
Pattern
Operation level
InferenceServingDeployment
Use cases
Serving large LLMs (tens to hundreds of billions of parameters)Inference of models that do not fit on a single GPUHigh-throughput text-generation APIsInference of Mixture-of-Experts modelsLong-context inference with a large KV cacheMulti-node GPU/TPU clusters

How it works

The model and computation are partitioned along one or more dimensions. In tensor parallelism, a layer's weight matrices are sharded across GPUs and results are combined with an AllReduce/AllGather within the layer. In pipeline parallelism, consecutive layers are assigned to separate stages (devices) and activations are sent point-to-point between stages; micro-batching reduces pipeline bubbles. Data/replica parallelism replicates the model and distributes requests for higher throughput. In MoE models the analogue is expert parallelism, routing tokens via an All-to-All operation. When memory is insufficient, weight/state sharding and offloading to CPU/NVMe (ZeRO-Inference) are used. A distributed KV cache and fast interconnects (NVLink/InfiniBand) minimize the communication cost that dominates the decoding phase.

Problem solved

A single accelerator has limited memory and compute, so the largest models (tens to hundreds of billions of parameters) do not fit on one GPU or run too slowly. Distributed inference addresses the memory and throughput limits by splitting the model and its computation across multiple devices.

Components

Tensor ParallelismSplits individual layers/matrices across devices (intra-layer).

Weight matrices (e.g., attention and MLP projections) are split column-/row-wise; results are combined via AllReduce/AllGather within the layer.

Pipeline ParallelismAssigns consecutive layers to separate stages (inter-layer).

The model is split into stages assigned to different devices; activations are passed point-to-point; micro-batching reduces pipeline bubbles.

Data / Replica ParallelismReplicates the model and distributes requests for throughput.

Official

Expert ParallelismDistributes MoE experts across devices with All-to-All routing.

Official

Collective CommunicationAllReduce/AllGather/All-to-All operations synchronizing devices.

Implemented by libraries such as NCCL; the dominant communication cost of distributed inference.

Weight & State Sharding / OffloadShards weights/state and offloads to CPU/NVMe (ZeRO-Inference).

Official

Implementation

Reference implementations
DeepSpeed
Python
Official
Megatron-LM
Python
Official
vLLM
Python
Official
Alpa
Python
Official
TensorRT-LLM
C++/Python
Official
Implementation pitfalls
Communication bottleneckHigh

Collective operations (AllReduce/All-to-All) dominate time when inter-device/inter-node links are slow.

Fix:Fast interconnects (NVLink/InfiniBand), overlap communication with compute, keep TP within a single node.
Pipeline bubblesMedium

Device idle time at pipeline fill/drain reduces efficiency.

Fix:Micro-batching and 1F1B-style schedules.
Load imbalance / uneven shardingMedium

Uneven weight or token (MoE) partitioning stalls faster devices.

Fix:Balance shards and routing; tune TP/PP to the architecture.

Evolution

Original paper · 2022 · Reza Yazdani Aminabadi
DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale
Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, Yuxiong He
2018
GPipe – pipeline parallelism
Inflection point

GPipe introduces pipeline parallelism with micro-batching for giant networks.

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (paper)
2019
Megatron-LM – tensor parallelism
Inflection point

Megatron-LM popularizes tensor parallelism for Transformer layers.

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (paper)
2022
DeepSpeed Inference
Inflection point

A dedicated distributed Transformer inference system combining TP, PP and offload (ZeRO-Inference).

DeepSpeed Inference (paper)
2022
Alpa – automated parallelism

Alpa automates inter- and intra-operator parallelism for distributed training and inference.

Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning (paper)
2023
vLLM / PagedAttention
Inflection point

vLLM introduces PagedAttention and distributed serving with KV-cache management, an industry standard for LLM inference.

Efficient Memory Management for Large Language Model Serving with PagedAttention (paper)

Hyperparameters (configurable axes)

Tensor parallel size (TP)High

Number of devices splitting a single layer.

Pipeline parallel size (PP)High

Number of pipeline stages (layer groups) on separate devices.

Data / replica parallel sizeMedium

Number of model replicas serving requests in parallel.

Micro-batch sizeMedium

Micro-batch size that reduces pipeline bubbles.

Computational complexity

Space complexity: O(P / N) na urządzenie.

Parallelism

Parallelism level
Partially parallel

Tensor parallelism is fully parallel within a layer, but pipelining introduces sequential dependencies between stages; the decoding phase is communication-bound.

Scope
InferenceAcross devicesAcross layers

Hardware requirements

Primary

Requires multiple GPUs with fast interconnects (NVLink/InfiniBand) for collective operations.

Good fit

TPU pods with ICI interconnects support model and pipeline parallelism well.