Tensor Parallelism › Megatron-LM Parallelism

Megatron-LM, developed by NVIDIA Research and described in the 2019 paper "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism" by Shoeybi et al., is the reference implementation of tensor parallelism for transformer models. It introduced the column-parallel and row-parallel chaining strategy described above and demonstrated that it could train models with billions of parameters efficiently on NVIDIA DGX clusters.

The paper was significant because it paired the technique with careful engineering validation. Shoeybi et al. trained models up to 8.3 billion parameters (at the time, the largest dense language model trained on a single cluster) and showed that their approach scaled without degradation in efficiency as the degree of parallelism increased from 1 to 8 GPUs per tensor-parallel group. This empirical validation was important because theoretical analyses of distributed training often miss hardware-specific bottlenecks that only appear at scale.

The key contributions of Megatron-LM are:

• Intra-layer model parallelism: Splitting individual weight matrices across devices, as opposed to inter-layer model parallelism where different layers go to different devices. Intra-layer splitting allows simultaneous computation across all GPUs on every layer, whereas inter-layer splitting creates idle GPUs whenever a layer is not the bottleneck.
• Minimal communication design: The column-row pairing ensures only one all-reduce per sub-layer in the forward pass and one dual collective in the backward pass, minimizing synchronization overhead.
• Fusion with data parallelism: Megatron-LM trains with both tensor parallelism (splitting layers within a node) and data parallelism (replicating groups across nodes), allowing scaling to thousands of GPUs with straightforward composition.

A Megatron-style transformer block processes input through both the self-attention and FFN sub-layers, each using the column-row chaining pattern. From the perspective of a single GPU in a tensor-parallel group of size $t$:

For the self-attention sub-layer, the input $X \in \mathbb{R}^{b \times s \times d}$ is available on all GPUs in the group (after the residual connection from the previous layer, which is identical on all devices after the previous all-reduce). The QKV weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ are each column-partitioned: this GPU holds columns $[d \cdot i/t, d \cdot (i+1)/t)$ for its rank $i$ within the tensor-parallel group. The output projection $W_O \in \mathbb{R}^{d \times d}$ is row-partitioned along its input dimension.

The computation within a single transformer layer flows as follows. Each GPU computes its column-parallel QKV projections independently, using only the locally stored weight shards. It then performs self-attention on its assigned subset of heads completely independently, since heads in standard self-attention do not interact during the attention computation. The attention output from those heads is then passed through the row-parallel output projection. Each GPU accumulates a partial sum of the full-dimension output, and then an all-reduce assembles the complete post-attention representation. This output is added to the residual stream (which is identical on all GPUs because both tensors are synchronized), and then the FFN applies the same column-row pattern.

The total communication per transformer layer under this design is two all-reduces in the forward pass (one after attention, one after FFN) and two dual collectives in the backward pass. For a 48-layer model with hidden dimension 8192 on 8 GPUs, this amounts to roughly 96 all-reduce operations per forward pass, each transmitting an activation tensor of about 64 MB at batch size 4 and sequence length 2048. The aggregate communication per forward pass is around 6 GB, manageable on NVLink but expensive on InfiniBand.

An important extension to the Megatron design is sequence parallelism, described in "Reducing Activation Recomputation in Large Transformer Models" by Korthikanti et al. (2022). The observation is that some operations in a transformer layer, specifically layer normalization and dropout, are not tensor-parallelized in the basic Megatron design; they operate on the full activation tensor replicated on every GPU. This wastes memory proportional to the tensor-parallel degree, since every GPU stores the same copy of the full activation during these replicated operations.

Sequence parallelism addresses this by distributing the sequence dimension across GPUs during the replicated operations. The key insight is that layer normalization and dropout operate independently on each token, so tokens can be distributed across GPUs (each GPU handling $s/t$ tokens) without any cross-GPU dependencies. During the tensor-parallel regions (attention and FFN), the model needs the feature dimension partitioned, not the sequence dimension. The transition between these two parallel modes requires collective operations: an all-gather to expand from sequence-parallel (token-sharded) to tensor-parallel (feature-replicated) mode, and a reduce-scatter to transition back.

These two new collectives replace the single all-reduce at the end of each sub-layer. The net result is that the total communication volume is unchanged (an all-reduce equals a reduce-scatter plus an all-gather), but now the activation memory is sharded throughout the layer rather than being replicated during layer normalization and dropout. Each GPU holds only $1/t$ of the activation memory throughout the full layer, reducing peak activation memory by a factor of $t$. For large models with high tensor-parallel degrees, this memory savings can be the difference between a model that fits in memory and one that does not.

The combination of tensor parallelism and sequence parallelism as described in Korthikanti et al. became the standard for training large models in the Megatron ecosystem, and its ideas influenced subsequent parallelism designs from other research groups.