Tensor and sequence parallelism
Derive local forward and backward matrix shapes, ownership, and collectives for column- and row-parallel linear layers
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Ability to multiply X[M,K] by W[K,N] and partition either K or N
- Chapter 13 all-reduce, all-gather, and reduce-scatter
- Chapter 3 transformer-layer dataflow
One multilayer perceptron projection receives an activation matrix X with shape [8,192 tokens, 8,192 features] and multiplies it by a weight matrix W with shape [8,192, 32,768]. The weight matrix and its matrix multiplication are too large for the desired per-GPU memory and step-time target. Four GPUs must cooperate on this one mathematical layer rather than process four different microbatches.
Partition the 32,768 output columns of W into four 8,192-column shards. Rank t stores W_t with shape [8,192, 8,192] and computes Y_t = XW_t, producing one feature shard with shape [8,192, 8,192]. Every rank needs X, but no reduction is needed because different output columns are mathematically independent. The following layer can consume the Y shards directly when its contracted dimension is partitioned compatibly.
A local shape log makes ownership visible, while PyTorch Profiler or Nsight Systems shows the matrix multiplications and the collectives at layout transitions. Compare tensor degree one, two, four, and eight. Higher degree reduces local dimensions, but the profiler can show shorter matrix kernels followed by a larger fraction of time in latency-sensitive all-reduces. A matrix benchmark can also show when the narrow local shape stops using the GPU efficiently.
Tensor parallelism is the established name for partitioning tensors and the operations that use them across ranks. Column and row partitions describe which weight dimension is divided. Every correct design must state the input layout, local weight shape, local output shape, and forward and backward communication at each boundary; the label alone predicts none of those details.
Derive column-parallel linear algebra
Let X have shape [M,K] and W have shape [K,N]. Split W by output columns across T ranks: W = [W0 W1 ...]. Rank t stores Wt with shape [K,N/T] and computes Yt = XWt with shape [M,N/T]. Concatenating the Yt tensors along output features reconstructs Y.
Every rank needs X. No reduction is required because output features are independent. An all-gather is needed only when the next operation requires the full Y. If the next operation can consume the feature shard, keep it local and defer communication.
In backward, rank t receives the matching output-gradient shard dYt. Rank t computes dWt = X^T dYt locally. Rank t also computes a partial input gradient dXt = dYt Wt^T with full [M,K] shape. If the input gradient is replicated across tensor ranks, sum the dXt contributions with an all-reduce. If the preceding operator expects a sharded input-gradient layout, an equivalent reduce-scatter or layout-specific transition can replace that replicated result.
Derive row-parallel linear algebra
Split W by input rows and X by the same contracted dimension: X = [X0 X1 ...] and W is stacked from Wt blocks of shape [K/T,N]. Rank t computes partial Yt = XtWt with shape [M,N].
The complete output is Y = sum_t Yt. An all-reduce can return full Y to every tensor rank. A reduce-scatter can combine the partial results while leaving Y partitioned along an eligible dimension for the next operation.
Assume the forward output and its gradient dY are replicated. Rank t computes dWt = Xt^T dY and its local input-gradient shard dXt = dY Wt^T. Both match locally owned shards, so this layer needs no additional backward reduction under those layout assumptions. If forward used reduce-scatter or another sequence-sharded output, backward must first provide the dY layout required by the local formulas. The collective follows the declared layouts, not the row-parallel label alone.
Pair adjacent transformer projections
In an MLP, make the D→F expansion column-parallel. Each rank produces F/T intermediate features and applies the activation locally. Make the F→D contraction row-parallel so the same rank consumes its local F/T shard and contributes a partial D-wide output.
This pairing avoids an all-gather of the large F-wide intermediate. Communication occurs at the smaller D-wide output boundary. The same principle applies to attention projections and output projection: choose layouts as a connected operator sequence.
Calculate local shapes and communication
For each operation, record the global matrix, local weight shard, local activation, output ownership, and forward collective. Derive backward in the same way: weight gradients and input gradients can require all-reduce, reduce-scatter, or all-gather according to the selected layout.
Communication frequency is per layer, often several times in forward and backward. Even a moderate message can dominate when it lies on every layer's dependency frontier. Place tensor groups on the fastest links that support the required degree.
Recognize the local-kernel limit
Increasing T reduces the N or K dimension of local GEMMs. A shard can become too narrow to use matrix tiles efficiently or provide enough thread blocks. Launch and collective startup remain while local arithmetic falls.
Measure local GEMM efficiency for the actual M, N/T, and K or M, N, and K/T shapes. Tensor degree should stop increasing when memory need or total step time no longer justifies smaller local work, even if network counters remain below saturation.
Add sequence parallelism
Tensor parallelism can replicate token-position activations for normalization, dropout, and residual work. Sequence parallelism partitions these eligible activations along the token dimension across the same ranks. Each rank performs elementwise or token-local operations on fewer positions.
Operators that need a different layout create transitions. A reduce-scatter can combine tensor-parallel partial outputs and distribute token segments; an all-gather can reconstruct the layout needed by a later projection. Calculate the activation memory saved and the added or transformed collectives.
Distinguish context parallelism
Context parallelism partitions long attention context so no rank stores or computes the complete sequence state. Depending on the algorithm, ranks exchange K and V blocks, partial attention outputs, or normalization statistics while preserving exact attention semantics.
Its communication scales with sequence, head, and datatype dimensions and follows an attention-specific schedule. Do not treat it as a rename of token-local sequence parallelism. Model which context ranges each rank owns and which remote information each query block must receive. Include causal masking, backward communication during exact training, and load imbalance when context blocks have unequal valid lengths.
Worked example: Partitioning an MLP pair
A feed-forward block receives X with shape [8,192, 4,096], expands to F = 16,384, applies a gated activation, and contracts back to D = 4,096 across four tensor-parallel ranks.
- Column-partition each expansion or gate weight from [4,096,16,384] into local [4,096,4,096] shards. Every rank receives the common [8,192,4,096] input and produces [8,192,4,096] intermediate features per projection.
- Apply the gated activation locally. The global intermediate would contain 8,192×16,384 elements, but no rank materializes or gathers that full tensor.
- Row-partition the contraction weight into four [4,096,4,096] shards along F. Each rank multiplies its local [8,192,4,096] activation by its local weight and produces a partial [8,192,4,096] output.
- All-reduce the four partial outputs to obtain the full residual-width result on each tensor rank, or reduce-scatter it along tokens if the next section uses sequence-parallel ownership.
- Calculate one BF16 partial-output payload: 8,192×4,096×2 bytes, about 67.1 MB. Apply the collective model and count this communication at every such layer boundary.
- Compare tensor degrees two, four, and eight. Degree eight halves local F shard width to 2,048, reduces per-rank memory, and changes GEMM efficiency. Measure whether the memory need and faster local compute exceed the higher collective and smaller-matrix costs.
Conclusion: Complementary partitions keep the expanded activation local and pay communication at the smaller layer boundary. Operator pairing is the source of efficiency.
Common errors
- Sharding each operation independently. Adjacent input and output ownership determines whether an avoidable gather appears between them.
- Applying a row-parallel split without partitioning X along the contracted dimension. Every rank would compute redundant or incompatible partial results.
- Increasing tensor degree until memory fits without measuring local GEMM shapes. Narrow shards can underuse matrix hardware.
- Placing per-layer tensor collectives across slow inter-node links when a sufficient faster local group is available.
Reference summary
Tensor parallelism divides one layer’s matrices across ranks. For Y = XW, a column-parallel projection partitions the output-feature columns of W, so each rank produces a slice of Y from the same X. A row-parallel projection partitions the contracted dimension, so each rank produces a partial contribution to the full Y and the contributions must be reduced.
Backward communication follows the same ownership rules. A column-parallel layer forms local weight gradients but produces partial input gradients that must be summed when the input gradient is replicated. A row-parallel layer with replicated output gradients forms its local input-gradient shard and local weight gradient without another reduction. State the activation layout at each boundary because sequence sharding or another adjacent layout can change these collectives.
- Pair a column-parallel expansion with a row-parallel contraction so the large intermediate activation remains sharded between the two operations.
- Tensor collectives occur inside many layers and often lie on the critical path. Keep their groups in the fastest suitable topology domain.
- Increasing tensor degree reduces per-rank matrix dimensions. Kernel efficiency can fall before the links reach saturation.
- Sequence parallelism partitions eligible activations along the token dimension and introduces gather or reduce-scatter transitions around operators that require another layout.
- Context parallelism partitions long attention context and communicates keys, values, or partial attention results according to the selected algorithm.
- Reduction order and format can change numerical results. Validate model quality and convergence with the actual collective path.
Knowledge check
Why can excessive tensor-parallel degree reduce throughput even when communication links remain unsaturated?
Show the answer guide
- A larger tensor degree shrinks per-rank M, N, or K dimensions and can produce narrow or poorly aligned matrix shapes that use the GPU's matrix pipelines inefficiently.
- Each layer still pays launch, synchronization, layout, and tensor-collective costs; these fixed costs become larger relative to smaller local arithmetic even before links saturate.
- The stopping point comes from measured local kernel efficiency, communication on the critical path, memory need, and total step time, not network utilization alone.
Derive and benchmark tensor-parallel shapes
For one X[M,K] × W[K,N] projection, derive column- and row-parallel local shapes for T = 1, 2, 4, and 8. Benchmark or model the local GEMMs and required collectives.
Evidence to keep: Shape and ownership table, exact forward and backward collective contracts, local kernel timings or roofline estimates, communication bytes, and the first degree whose complete predicted or measured time worsens.
Tensor parallelism partitions within a layer. Pipeline and expert parallelism partition layers or conditional subnetworks, introducing bubbles and routing imbalance.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.