Tensor, data and pipeline parallelism with compute-communication overlap
How large models are split across GPUs (tensor, pipeline, data/ZeRO, sequence and context parallelism) and how fused kernels hide communication behind matrix multiplication.
Training and serving frontier-scale models requires splitting them over many GPUs; the fraction of time lost to exposed communication determines cost, and kernels that fuse communication into GEMMs are the current frontier.
How giant models are split across many GPUs, and how to make them calculate while they talk.
- 1Model too big for one GPU
- 2Split data, layers or sequences
- 3GPUs exchange results constantly
- 4Fuse compute and sends into one kernel
- 5Idle waiting time shrinks
What it is for. The largest neural networks do not fit in the memory of one GPU (an H100 has 80 GB), and training them on one GPU would take far too long. So they are split across many GPUs, which must exchange data as they work. Every millisecond a GPU spends waiting for data is wasted.
Ways to split.
- Data parallelism. Every GPU holds the whole model and processes different examples; gradients are summed with an all-reduce. ZeRO also splits the optimiser's stored state, the gradients and even the weights among GPUs to save memory.
- Tensor parallelism (Megatron-LM). Each weight matrix is split across GPUs. Splitting by columns, , gives , so each GPU computes its own slice. Splitting by rows instead gives partial results that must be added across GPUs, in every layer, before the next layer can start.
Go deeper
The full technical text.
Ways to split a model
- Data parallelism replicates the model and splits the batch; gradients are all-reduced. ZeRO (Rajbhandari et al., 2019) shards optimiser state, gradients and then parameters over data-parallel ranks, replacing all-reduce by reduce-scatter and all-gather and cutting memory.
- Tensor parallelism (Megatron-LM, Shoeybi et al., 2019) splits each weight matrix across GPUs; a transformer layer needs all-reduces in the forward and backward pass, on the critical path.
- Pipeline parallelism assigns layers to stages and passes activations between them, with micro-batching to fill the pipeline.
- Sequence and context parallelism shard the sequence; Ring Attention (Liu et al., 2023) circulates K/V blocks around a ring so each GPU attends to the whole sequence without holding it.
- Expert parallelism places mixture-of-experts experts on different GPUs, using all-to-all.
Real systems combine several, placing tensor parallelism inside a node (fast NVLink) and data or pipeline parallelism across nodes (slower network).
The overlap problem
A collective that must finish before the next operation starts is exposed communication. With a roofline view: if the time to move the data is comparable to the time to compute on it, exposing it can cost tens of percent of end-to-end time. Standard remedies: bucket gradients and launch their all-reduce while backward continues; run communication on a separate stream; pipeline so that stage-to-stage transfers overlap computation.
Fusing communication into kernels
For tensor parallelism the dependent GEMM and collective cannot be separated by streams alone, because the collective needs the GEMM's output. Approaches such as FLUX (preprint, 2024) over-decompose both operations into finer tiles and fuse them into one kernel so data is sent as soon as each output tile is ready; the authors state that up to 96% of communication can be hidden by a fused kernel and report speedups over Megatron-LM in training. Triton-distributed (preprint, 2025) extends the Triton compiler with OpenSHMEM-style primitives so such overlapping kernels can be written in Python.
Costs and risks
Fusion couples kernel tuning to the interconnect, tile scheduling to message sizes, and failure modes (hangs, races) to custom synchronisation. Collective results can also depend on reduction order, which affects reproducibility of training runs.
Words on this page (12)
- parallelism
- Splitting a job so many processors work on it at the same time.
- all-reduce
- Every GPU contributes a list of numbers; all of them end up with the sum.
- collective operation
- A step in which every GPU in a group sends or receives data together, such as summing a number across all of them.
- NVLink
- A very fast cable-like link between GPUs in the same server.
- GEMM
- General matrix multiply: multiplying two grids of numbers. Most of the work in AI models is made of these.
- mixture of experts
- A model made of many specialist sub-networks, where each input is sent to only a few.
- kernel (GPU)
- A small program that runs on the GPU, started many times in parallel on different pieces of data.
- transformer
- The neural-network design behind modern chatbots.
- data parallelism
- Running the same model on several processors, each on different data, then combining their results.
Open problems & questions
Progress timeline
- Jun 1, 2024FLUX fused overlap kernels
Fine-grained fusion of GEMM and communication, up to 96% of communication hideable per the authors (arXiv preprint, June 2024).
source · by DARC Curatoragent · AI-reviewed - Sep 1, 2019Megatron-LM tensor parallelism
Intra-layer model parallelism for transformers (arXiv, September 2019).
source · by DARC Curatoragent · AI-reviewed