Tensor Parallelism › Communication Patterns

The communication overhead in tensor parallelism is dominated by two collectives: all-reduce and all-gather.

In the Megatron-style column-row chaining, the forward pass uses one all-reduce per transformer sub-layer to combine partial sums from row-parallel layers. This is the dominant communication pattern during training. The backward pass mirrors this with a symmetric but inverted set of collectives, as discussed below.

The communication volume per transformer layer scales as $O(bsd)$ where $b$ is batch size, $s$ is sequence length, and $d$ is hidden dimension. This is the volume of one activation tensor, and each all-reduce transmits this volume. With modern NVLink interconnects providing 600 GB/s of bidirectional bandwidth on an A100 DGX node, the communication overhead is manageable for moderate tensor parallelism degrees (typically 8 within a node), but becomes expensive for larger degrees that span nodes and are limited to slower InfiniBand links at roughly 25-50 GB/s.

The practical implication is that tensor parallelism is almost always limited to a single node. Within a node, NVLink provides sufficient bandwidth to keep communication overhead at a few percent of compute time. Across nodes, the bandwidth drop of 10-20x makes the communication overhead prohibitive. This is why production training systems for the largest models use tensor parallelism within a node (for intra-layer splitting) and pipeline parallelism across nodes (for inter-layer splitting), a combination known as 3D parallelism when combined with data parallelism across pipeline groups.

A subtle but important point about tensor parallelism is that the communication pattern in the backward pass is the dual of the forward pass. This symmetry is what makes the approach efficient and predictable.

In the forward pass:

• Column-parallel layers require no communication (full input available, partial output produced)
• Row-parallel layers require an all-reduce (partial outputs summed to full output)

In the backward pass, the gradients flow in the opposite direction through the same graph. For row-parallel layers, the gradient of the loss with respect to the full output (which every GPU has after the forward all-reduce) must be distributed back to the input shards. This requires a reduce-scatter: each GPU receives the gradient slice corresponding to its input shard. For column-parallel layers, the gradient with respect to the weight shard can be computed locally (since the full input was available on every GPU), and the gradient with respect to the full input requires an all-reduce of the partial gradients computed from each weight shard.

This duality means that the backward pass has exactly as many collective operations as the forward pass, and each one transmits roughly the same data volume. The total communication volume per training step (forward plus backward) is approximately $4$ times the per-layer all-reduce volume times the number of layers. For a model like GPT-3 with 96 layers, an activation size of roughly 50 MB per layer, and a tensor-parallel degree of 8, this amounts to several gigabytes of data transmitted per training step, well within the capacity of a DGX node's NVLink fabric at reasonable batch sizes.

An important optimization for memory efficiency comes from observing that an all-reduce is equivalent to a reduce-scatter followed by an all-gather. This decomposition matters for sequence parallelism, which is discussed in the next section, but it is worth understanding in the context of standard tensor parallelism as well.

In a naive all-reduce, every GPU accumulates the full-size sum and every GPU holds the full result, using $O(S)$ memory on each device. In a reduce-scatter plus all-gather, after the reduce-scatter each GPU holds only $O(S/N)$ of the sum, and the all-gather distributes it. During the time between reduce-scatter and all-gather, memory usage is reduced by a factor of $N$. By exploiting this gap (inserting useful computation like layer normalization between the two primitives), sequence parallelism reduces peak activation memory while keeping the same total communication volume.