A visual guide to parallel Transformers

How a transformer’s forward and backward passes are sharded across devices, first one strategy at a time and then in the combinations used to train recent models. The notation follows the JAX scaling book.

Notation #

We use the sharding notation from the scaling book. A mesh such as Mesh({'X': 4}) is a set of four devices along one axis named X. A subscript on a dimension gives the mesh axis that dimension is sharded over. For example, In[B_X, T, D] is sharded along the batch over X, so each device holds an array of shape [B/4, T, D]. A dimension with no subscript is replicated. A trailing {U_X} marks a partial sum: each device holds an array of the full shape, and the true value is the sum of these arrays over X. AllReduce and ReduceScatter compute this sum.

The tooltips assume B=8, T=128, D=1024, F=4096, H=1024 (all attention heads concatenated), E=4 experts with S=256 tokens routed to each, and bf16.

Reading the figures #

The model runs from left to right: In, then the attention and MLP blocks of each layer, then Out. Each device is a horizontal lane, so tensors sent between devices move vertically. Weights are blue, activations amber, keys and values teal, and gradients rose. A dashed outline is the full array and the solid part is the shard held by that device. Hovering over a tensor shows its shape, sharding and memory.

Each segment of the step bar is one operation. Click a segment to see the state after that operation, or press play. For a collective, the panel below the canvas shows how it runs on a ring of devices. The program on the right lists every operation. The figures start with the forward pass; + Backward adds the backward pass, including the activations saved for it, and links each backward operation to the forward operation it comes from.

The axes of transformer sharding #

Each strategy below shards one thing: the batch, the optimizer state, the weights, the feature dimensions, the sequence, the experts, or the layers. Each runs on four devices.

Data parallelism #

The batch is sharded, In[B_X, T, D], and the weights are replicated. The forward pass needs no communication. In the backward pass, each device computes the weight gradients from its part of the batch, so each gradient dW[D, F]{U_X} is a partial sum. It is AllReduced over X before the optimizer step.

ZeRO-1 #

Data parallelism also replicates the optimizer state, which for Adam is two more arrays the size of the weights. ZeRO-1 shards the optimizer state over X and keeps the weights replicated. The forward pass is unchanged. In the backward pass, each gradient dW[D, F]{U_X} is ReduceScattered instead of AllReduced, giving dW[D_X, F]. Each device updates its shard of the weights, W[D_X, F], and an AllGather gives every device the full updated weights. An AllReduce is a ReduceScatter followed by an AllGather, so ZeRO-1 has the same communication cost as data parallelism.

Fully-sharded data parallelism (ZeRO-3) #

FSDP also shards the weights over X: W_in[D_X, F] and W_out[F, D_X]. Each weight is AllGathered just before its matmul and freed just after. The backward pass AllGathers each weight again, and each weight gradient is ReduceScattered so that it is sharded like its weight.

Tensor parallelism #

Tensor parallelism shards the attention heads, W_qkv[D, H_Y] and W_o[H_Y, D], and the MLP hidden dimension, W_in[D, F_Y] and W_out[F_Y, D]. Activations are sharded along D. Each block starts with an AllGather and ends with a ReduceScatter. The first matmul contracts over D, so the activations are AllGathered first. The second matmul contracts over the sharded H or F, so each device holds a partial sum {U_Y}. The ReduceScatter sums it and shards the result along D for the next block. The backward pass does the same in reverse.

Context parallelism #

Context parallelism shards the sequence: In[B, T_X, D]. The MLP acts on each token separately and needs no communication. In attention, each query needs every key and value, so K and V are AllGathered over X while Q stays sharded. The backward pass AllGathers K and V again. Each device computes gradients for all keys and values, so dK and dV are partial sums and are ReduceScattered back over T. The weight gradients are AllReduced over X, as in data parallelism.

Expert parallelism #

In a mixture-of-experts layer, the MLP is replaced by E experts, one per device, W_in[E_Z, D, F], and a router assigns each token to one expert. Tokens are drawn as squares colored by their expert. An AllToAll sends each token to the device that holds its expert, each expert runs its MLP, and a second AllToAll sends the tokens back. In the backward pass, the token gradients go through the same two AllToAlls. The expert weight gradients need no communication, since each device holds its own expert. The attention weight gradients are AllReduced over Z, as in data parallelism.

Pipeline parallelism #

Pipeline parallelism shards the layers: stage 0 holds layer 1 and stage 1 holds layer 2. The activations are sent from one stage to the next with a point-to-point send. With a single batch, one stage is always idle, so the batch is split into four microbatches that go through the stages in turn. The grid below the lanes shows which microbatch each stage runs at each time step. The empty cells are idle time, called the bubble. The backward pass takes about twice the compute of the forward pass, so its cells are twice as wide.

Examples from frontier models #

Large training runs combine several of these strategies. The technical reports for DeepSeek-V3, Kimi K2, GLM-4.5V, Nemotron-4, Nemotron-H, Nemotron 3, Qwen3-VL and Llama 3 use three combinations. The figures use two devices per axis; the sizes used in the reports are given in the text.

In these figures, each pipeline stage is a grid of devices. The row and column labels give each device’s coordinate along each mesh axis. Each collective runs over one axis and is drawn in that axis’s color, and the label under the grid names the axis and the strategy. The row under the step bar counts the collectives on each axis.

Dense models: Llama 3 #

FSDP × TP × CP × PP. Llama 3 405B used TP 8, PP 16 and FSDP 64 for 8K-token sequences, and TP 8, CP 16, PP 16 and FSDP 8 for 128K-token sequences. TP is placed within a server, where bandwidth is highest, and data parallelism across servers. Nemotron-4 340B used TP 8, PP 12 and data parallelism; Nemotron-H used TP 8 and 768-way data parallelism. The figure uses Mesh({'X': 2, 'Y': 2, 'C': 2, 'stage': 2}) with In[B_X, T_C, D_Y]. Each matmul has an FSDP AllGather of the weight and a TP AllGather of the activations. CP adds the K and V AllGathers in attention, and PP adds one send in each direction.

Open MoE models: DeepSeek and Kimi #

EP × CP × PP × ZeRO-1, without TP. DeepSeek-V3 used PP 16, EP 64 across 8 nodes and ZeRO-1, with no tensor parallelism. Kimi K2 used PP 16, EP 16 and ZeRO-1. Neither used context parallelism; DeepSeek-V3 extends its context to 128K with YaRN. Their successors add it: DeepSeek-V4 and Kimi K3 both use context parallelism for long-context training, and GLM-4.5V adds CP 4 for its long-context stage. The figure uses Mesh({'X': 2, 'C': 2, 'Z': 2, 'stage': 2}) with In[B_XZ, T_C, D]. The batch is sharded over both X and Z, since outside the MoE layers the expert axis is another data axis. The forward pass has no weight AllGathers. It has the K and V AllGathers over C, two AllToAlls over Z per MoE layer, and the send between stages. In the backward pass, each attention weight gradient is ReduceScattered over X and AllReduced over Z and C. Each expert weight gradient is ReduceScattered over X and AllReduced over C only, since each expert is held at a single Z coordinate.

All five: Nemotron 3 and Qwen3-VL #

FSDP × TP × CP × EP × PP. The long-context stage of Nemotron 3 Nano used CP 8, TP 8, EP 8 and PP 4, and Nemotron 3 Ultra used EP 128. Qwen3-VL uses TP, PP, CP, EP and ZeRO-1 on up to 10,000 GPUs. The figure uses 32 devices, Mesh({'X': 2, 'Y': 2, 'C': 2, 'Z': 2, 'stage': 2}), with In[B_XZ, T_C, D_Y]. Each stage is a 4 × 4 grid with rows X and C and columns Z and Y. Most collectives are the FSDP and TP AllGathers and ReduceScatters. CP communicates only in attention, EP adds two AllToAlls per MoE layer, and PP adds one send in each direction. In the backward pass, each attention weight gradient is summed over X, Z and C, the three axes that shard the tokens, which takes three collectives.