Tensor Parallelism › Summary

Tensor parallelism partitions individual weight matrices across multiple GPUs, enabling models to be trained whose layers would overflow the memory of any single device. The two fundamental splitting strategies are:

• Column parallelism: The weight matrix is split along the output dimension. Each GPU holds a column shard and computes on the full input independently, producing a shard of the output. No communication is needed during the forward pass; output shards are consumed directly by a downstream row-parallel layer.
• Row parallelism: The weight matrix is split along the input dimension. Each GPU holds a row shard and receives the corresponding input shard. Each GPU produces a full-size partial sum, and an all-reduce aggregates these partial sums into the correct final output.

Chaining column-parallel and row-parallel layers eliminates intermediate communication: the shard-format output of the column-parallel layer is precisely the input format the row-parallel layer expects, so no all-gather is needed between them. This enables a transformer's FFN and attention sub-layers to execute with only one all-reduce per sub-layer in the forward pass, minimizing synchronization overhead to two all-reduces per transformer layer.

The Megatron-LM framework from NVIDIA Research established this column-row chaining as the standard approach for intra-layer model parallelism. The design exploits the natural parallelism of multi-head attention (heads are independent and map cleanly to column-parallel splits) and the FFN structure (the expansion layer is column-parallel, the contraction layer is row-parallel). Extended with sequence parallelism (distributing sequence tokens across GPUs during non-tensor-parallel operations like layer normalization), the Megatron design achieves near-linear memory reduction proportional to the tensor-parallel degree while keeping communication overhead manageable on high-bandwidth NVLink interconnects.

Key takeaways:

• Tensor parallelism is complementary to data parallelism and pipeline parallelism, and production training systems combine all three in what is often called 3D parallelism.
• The optimal tensor-parallel degree depends on model size, layer dimensions, network bandwidth, and GPU memory capacity. Within a node, degrees of 2-8 are typical.
• Model architecture choices (hidden dimension, number of heads, FFN expansion ratio) should be divisible by likely tensor-parallel degrees; powers of 2 and multiples of 128 are standard choices.
• Crossing node boundaries with tensor parallelism is expensive due to InfiniBand's lower bandwidth relative to NVLink; pipeline parallelism is preferred for inter-node communication.
• Weight initialization must account for the tensor-parallel degree: row-parallel layers require a $1/\sqrt{N}$ scaling factor to maintain correct output variance when shards are summed.