Home › Neural network training & inference › Multi-GPU parallelism
algorithm low

Tensor parallelism

Problem A single layer of a very large neural network is too big or too slow for one GPU. Engineers need to divide the layer's work across several GPUs on a fast link.

Split each layer's weight matrices across GPUs (Megatron-style column then row splits) so one layer's math runs on t devices with one all-reduce per block.

per GPU FLOPs and weights ÷ t; 2 all-reduces of b·s·h activations per transformer layer forward, 2 backward
model-parallelmegatronsequence-parallelall-reducetransformer
Overview

When one layer is too big or too slow for one GPU, split its matrices. In an MLP Y = GeLU(X·A)·B, split A by columns: each GPU computes its slice of the hidden units and applies GeLU locally. Then split B by rows: each GPU produces a partial Y, and one all-reduce sums them. Attention splits the same way, by heads.

Key facts

Use when

  • Dense layers too big for one GPU (about 20B+), or latency-bound inference of big models.

Avoid when

  • No fast interconnect (PCIe-only or across nodes).
  • Small hidden sizes: per-GPU GEMMs get too thin.

Deep dive

How it works

In: layer input X (replicated or sequence-sharded) and per-rank weight shards in a group of t GPUs · Out: the layer output, identical to the unsharded layer · State: each rank's 1/t of the weights (and of the optimizer state) · Cost: FLOPs and weights ÷ t; 4 all-reduces of b·s·h elements per transformer layer per step

Megatron's two conjugate operators (t ranks in a TP group):

  • f: identity in forward, all-reduce of the gradient in backward. Placed before a column-parallel layer.
  • g: all-reduce in forward, identity in backward. Placed after a row-parallel layer.
MLP:        X ─f─► [X·A_i] ─► GeLU ─► [·B_i] ─g─► Y          A = [A_1 … A_t] by columns
                   (b·s·4h/t)          partial Y              B = [B_1; …; B_t] by rows
Attention:  X ─f─► Q_i,K_i,V_i = X·W_i (heads h/t per rank) ─► attn ─► ·O_i ─g─► Y
Embedding:  vocab-parallel: each rank holds V/t rows; mask out-of-range ids, all-reduce
Loss:       vocab-parallel cross-entropy: all-reduce the row max and the sum-exp, never gather logits

Sequence parallelism (Korthikanti et al.): between blocks, activations are sharded along the sequence dimension. g becomes a reduce-scatter (each rank gets s/t tokens of Y), LayerNorm and dropout run on the shard, and f becomes an all-gather. Communication is unchanged; activation memory for those regions drops by t.

Code

A Megatron MLP split over t ranks on plain slices. A column-parallel first GEMM and a row-parallel second GEMM need exactly one all-reduce forward ("g") and one backward ("f"). The code checks this against the unsharded layer.

fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {   // (m×k)(k×n)
    let mut c = vec![0.0f32; m * n];
    for i in 0..m { for p in 0..k { let x = a[i * k + p]; for j in 0..n { c[i * n + j] += x * b[p * n + j]; } } }
    c
}
fn transpose(a: &[f32], r: usize, c: usize) -> Vec<f32> {
    let mut t = vec![0.0f32; r * c];
    for i in 0..r { for j in 0..c { t[j * r + i] = a[i * c + j]; } }
    t
}
fn all_reduce(parts: &[Vec<f32>]) -> Vec<f32> {                 // NCCL all-reduce stand-in
    (0..parts[0].len()).map(|i| parts.iter().map(|p| p[i]).sum()).collect()
}
fn relu(x: &[f32]) -> Vec<f32> { x.iter().map(|v| v.max(0.0)).collect() }
fn relu_bwd(g: &[f32], pre: &[f32]) -> Vec<f32> { g.iter().zip(pre).map(|(g, z)| if *z > 0.0 { *g } else { 0.0 }).collect() }
fn close(a: &[f32], b: &[f32]) -> bool { a.iter().zip(b).all(|(x, y)| (x - y).abs() < 1e-3 * (1.0 + y.abs())) }

/// Columns [i·w, (i+1)·w) of a row-major (rows×cols) matrix.
fn col_block(a: &[f32], rows: usize, cols: usize, w: usize, i: usize) -> Vec<f32> {
    (0..rows).flat_map(|r| a[r * cols + i * w..r * cols + (i + 1) * w].iter().copied()).collect()
}

fn main() {
    let (m, d, f, t) = (8, 32, 128, 4);                        // tokens, hidden, ffn width, TP degree
    let rnd = |n: usize, s: usize| -> Vec<f32> { (0..n).map(|i| ((i * 2654435761 + s) % 1000) as f32 / 1000.0 - 0.5).collect() };
    let (x, a, b) = (rnd(m * d, 1), rnd(d * f, 2), rnd(f * d, 3));  // Y = relu(X·A)·B
    let fs = f / t;
    let a_sh: Vec<Vec<f32>> = (0..t).map(|i| col_block(&a, d, f, fs, i)).collect();     // A by columns
    let b_sh: Vec<Vec<f32>> = (0..t).map(|i| b[i * fs * d..(i + 1) * fs * d].to_vec()).collect(); // B by rows

    // Forward. X is replicated; hidden activations stay sharded (m × f/t per rank).
    let pre: Vec<Vec<f32>> = a_sh.iter().map(|ai| matmul(&x, ai, m, d, fs)).collect();
    let partial: Vec<Vec<f32>> = (0..t).map(|i| matmul(&relu(&pre[i]), &b_sh[i], m, fs, d)).collect();
    let y = all_reduce(&partial);                               // "g": the only forward communication
    let pre_full = matmul(&x, &a, m, d, f);
    assert!(close(&y, &matmul(&relu(&pre_full), &b, m, f, d)));

    // Backward for loss ½‖Y‖², so dY = Y (identical on every rank after the all-reduce).
    let dy = y.clone();
    let mut dx_parts = vec![];
    for i in 0..t {
        let dh = relu_bwd(&matmul(&dy, &transpose(&b_sh[i], fs, d), m, d, fs), &pre[i]);
        let da_i = matmul(&transpose(&x, m, d), &dh, d, m, fs);   // weight grads: local, no comm
        let dh_full = relu_bwd(&matmul(&dy, &transpose(&b, f, d), m, d, f), &pre_full);
        assert!(close(&da_i, &col_block(&matmul(&transpose(&x, m, d), &dh_full, d, m, f), d, f, fs, i)));
        dx_parts.push(matmul(&dh, &transpose(&a_sh[i], d, fs), m, fs, d));
    }
    let dx = all_reduce(&dx_parts);                             // "f": the only backward communication
    let dh_full = relu_bwd(&matmul(&dy, &transpose(&b, f, d), m, d, f), &pre_full);
    assert!(close(&dx, &matmul(&dh_full, &transpose(&a, d, f), m, f, d)));
    println!("TP={t} MLP matches the unsharded layer: 1 all-reduce forward, 1 backward");
}

In PyTorch the same plan is parallelize_module(mlp, tp_mesh, {"fc1": ColwiseParallel(), "fc2": RowwiseParallel()}) on DTensors. Megatron implements f and g as autograd functions around all_reduce.

Complexity and constants

  • Forward FLOPs per layer ≈ 24·b·s·h² + 4·b·s²·h, split t-ways. Communication per layer per step: 4 all-reduces of 2·b·s·h bytes (bf16).
  • Example: b = 1, s = 4096, h = 8192, t = 8. Each all-reduce is 67 MB, about 0.3 ms at about 400 GB/s busbw, so about 1.2 ms per layer per step. Compute: 3 × 24·b·s·h² / t ≈ 2.5 TFLOP per GPU, about 4 ms at about 600 TFLOPS. Exposed communication is about 25% unless overlapped, and roughly 8x worse over InfiniBand.
  • Megatron-LM ran 8-way TP within DGX nodes combined with pipeline and data parallelism (PTD-P) at 502 PFLOP/s on 3072 A100s (52% of peak) for a 1T-parameter model.

Efficiency

  • Cost: per layer 4 all-reduces of 2·b·s·h bytes (bf16). Example from above: b = 1, s = 4096, h = 8192, t = 8 gives 67 MB each, about 0.3 ms at about 400 GB/s, so about 1.2 ms per layer against about 4 ms of compute: about 25% exposed unless overlapped, order of 8x worse over InfiniBand. Megatron ran 8-way TP inside nodes at 52% of A100 peak on 3072 GPUs (Narayanan et al.). Bound: collective latency and bandwidth on the critical path.
  • Levers:
    • keep t inside one NVLink domain (8 GPUs) — avoids the roughly 8x slower inter-node all-reduce (above) — always
    • sequence parallelism — norm and dropout activations ÷ t at the same communication volume (Korthikanti et al. 2022) — long sequences
    • overlap collectives with GEMM chunks (async TP, collective matmul) — hides the exposed communication (Wang et al., ASPLOS 2023) — large layers
    • vocab-parallel embedding and loss — never gathers b·s·V logits — large vocabularies
    • TP × pipeline × data (PTD-P, 2D/3D meshes) — scale beyond one node (Narayanan et al. 2021) — more than about 8 GPUs
  • Naive → standard → state of the art:
    • one GPU, layer does not fit · not runnable
    • column then row split, 4 blocking all-reduces per layer · about 25% exposed communication in the example
    • sequence-parallel TP with overlapped reduce-scatter and all-gather · communication largely hidden

Implementation notes

  • Make the hidden size, heads and FFN width divisible by t. Keep h/t a multiple of 64-128 for tensor-core tiles.
  • GQA with fewer KV heads than t: replicate KV heads (each rank holds ≥1), or TP would split a head.
  • Overlap TP communication with GEMMs: decompose the GEMM into chunks and pipeline reduce-scatter/all-gather chunks against compute (async TP, collective matmul in XLA, Transformer Engine userbuffers).
  • Initialize shards consistently. Either slice one full matrix (as above, from a shared seed) or use a per-shard RNG offset, or TP results differ from the unsharded model.
  • Use TP with ZeRO/FSDP across data-parallel replicas (a 2D mesh) and pipeline stages across nodes (3D parallelism).

Pitfalls

  • Adding the row-parallel bias on every rank before the all-reduce adds it t times.
  • Dropout in TP regions needs a per-rank RNG stream for sharded activations and an identical stream for replicated activations (Megatron's model_parallel_cuda_manual_seed).
  • Gradient clipping must count sharded parameters once and replicated parameters once, not t times.
  • Calling .contiguous() before collectives is mandatory: NCCL needs dense buffers, and transposed views silently copy or fail.
  • Inference: TP across PCIe GPUs is often slower than one GPU plus quantization (weight quantization).

Variants and alternatives

Sequence parallelismMegatron SP inside TP regions; context parallelism or ring attention to shard long sequences across GPUs.
2D/2.5D/3D tensor parallelismSUMMA-style grid sharding, lower communication asymptotically. Rarely worth it below hundreds of GPUs.
Expert parallelismMoE experts on different GPUs with all-to-all token routing.
ZeRO-3/FSDPshards storage, not math. Simpler, and works across nodes.Pipeline parallelismsplit by layers; cheaper communication, but adds a bubble.

In production

Megatron-LM and Megatron-Core, NVIDIA NeMo, DeepSpeed (with Megatron), PyTorch DTensor parallelize_module (TorchTitan), JAX/XLA GSPMD sharding, and the vLLM, SGLang and TensorRT-LLM --tensor-parallel-size inference paths.

Implementations

namelangrepoinstallverdictlicencehealthas_ofnotes
Megatron-LM / Megatron CorePython / CUDANVIDIA/Megatron-LMpip install megatron-core————origin of column/row TP and sequence parallelism
PyTorch DTensorPython / C++pytorch/pytorchpip install torch————parallelize_module with ColwiseParallel / RowwiseParallel
TorchTitanPythonpytorch/torchtitanpip install torchtitan————TP, PP and FSDP on one device mesh
vLLMPython / C++ / CUDAvllm-project/vllmpip install vllm————--tensor-parallel-size for inference
DeepSpeedPython / CUDAdeepspeedai/DeepSpeedpip install deepspeed————tensor parallelism for training and inference

Sources

  • Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019. paper · code
  • Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021. paper
  • Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022. paper
  • Xu et al., GSPMD: General and Scalable Parallelization for ML Computation Graphs, 2021. paper
  • Wang et al., Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models, ASPLOS 2023.

Builds on

Collective communicationMany machines or GPUs must exchange or combine large buffers in lockstep, for example sharding weights, routing items to owners or sharing one result. Cost grows with link speed, latency and machine count.Ring all-reduceData-parallel training on many GPUs must add up large gradient buffers from every worker each step. Without a good scheme, communication time grows with worker count and becomes the bottleneck.Tiled matrix multiplicationDevelopers writing GPU or CPU kernels for matrix multiplication, the main workload of neural networks, find naive code runs at a tiny fraction of hardware speed because memory cannot keep up with the arithmetic.

Used by

Systems (11):

training fabricrack scale systemscale up fabricllm pretrainingllm inferencellm api platformautodiff frameworkvideo generationml weather modelmulti datacenter trainingreasoning rl

Alternatives and related

ZeRO and FSDP shardingTraining a model whose weights, gradients and optimizer state do not fit in one GPU's memory, without rewriting the model to split its layers across devices, is the problem for data-parallel training.Pipeline parallelismA neural network is too large to fit on one GPU, and the machines are linked by slow connections. Engineers need to spread consecutive layers across devices without leaving most of them idle.Gradient checkpointingTraining a deep network or using long sequences runs out of GPU memory because stored intermediate values from the forward pass grow with depth and batch size. Engineers need to fit the run in memory.Mixture of experts (MoE)Training and serving a very large neural network is limited by how much computation each input item needs. Designers want far more model capacity without a proportional rise in per-token compute.Transformer blockAnyone building or sizing a language model needs the standard structure of its layers, and the exact sizes, compute and memory cost that follow from that structure.LLM pretrainingHow is a large language model trained from scratch in practice? This guide covers deciding data mix, scale, schedules and parallel layout so a months-long run on thousands of GPUs stays stable and efficient, and whether to pretrain at all.