Tensor Parallelism in Large Language Models: How Megatron-LM Partitions Multi-Layer Perceptrons and Attention Heads

Training and serving modern large language models requires navigating severe hardware memory and compute constraints. While standard Distributed Data Parallelism (DDP) replicates the entire model across multiple accelerators, modern frontier architectures containing tens or hundreds of billions of parameters exceed the physical memory capacity of any single GPU.

Even with 80 GB or 144 GB of high-bandwidth memory (HBM), mixed-precision training with 16-bit weights and 32-bit Adam optimizer states consumes between 16 and 20 bytes of static memory per parameter before allocating a single byte for activations or KV caches.

To scale beyond single-device memory walls without incurring prohibitive pipeline bubbles or communication bottlenecks, modern deep learning frameworks rely on Tensor Parallelism (TP). Pioneered by NVIDIA Research in the Megatron-LM framework (Shoeybi et al., 2019), Tensor Parallelism splits individual weight matrices and matrix multiplication operations (GEMMs) across multiple GPUs within a Transformer layer.

Understanding how Megatron-LM constructs tensor-parallel Multi-Layer Perceptrons (MLPs) and Multi-Head Attention (MHA) blocks reveals the mathematical elegance and communication efficiency underpinning distributed LLM infrastructure.

The Parallelism Taxonomy and the Intra-Node Domain

Distributed training and inference for foundation models generally combine three orthogonal forms of parallelism:

1. Data Parallelism (DP) and ZeRO/FSDP: The model is replicated or sharded across GPUs, and each GPU processes a distinct batch of input tokens. Gradients or parameters are synchronized across devices via collective communications like All-Reduce or Reduce-Scatter/All-Gather as detailed in DeepSpeed ZeRO (Rajbhandari et al., 2020).
2. Pipeline Parallelism (PP): The sequential layers of a Transformer are partitioned across devices (e.g., layers 1 to 8 on GPU 0, layers 9 to 16 on GPU 1). This introduces pipeline scheduling mechanisms (such as 1F1B) and pipeline bubbles where GPUs sit idle waiting for boundary activations.
3. Tensor Parallelism (TP): Individual matrix multiplications within a single Transformer layer are sharded across a group of GPUs. Every device computes a shard of the layer concurrently for the same batch of tokens.

Because Tensor Parallelism executes multiple communication collectives inside every single Transformer layer, it is extremely sensitive to latency. Consequently, TP is almost exclusively deployed intra-node, where GPUs communicate across ultra-high-bandwidth interconnects like NVIDIA NVLink and NVSwitch (providing 900 GB/s to 1.8 TB/s bidirectional bandwidth per accelerator), rather than across inter-node InfiniBand or Ethernet networks.

The Core Primitives: Column-Parallel and Row-Parallel Linear Layers

The foundation of Megatron-LM is the decomposition of standard linear transformations ($Y = XW$) across a tensor-parallel group of size $N$.

A standard matrix multiplication multiplies an input activation tensor $X \in \mathbb{R}^{B \times H_{\text{in}}}$ by a weight matrix $W \in \mathbb{R}^{H_{\text{in}} \times H_{\text{out}}}$, yielding an output tensor $Y \in \mathbb{R}^{B \times H_{\text{out}}}$, where $B$ is the sequence/batch token dimension, $H_{\text{in}}$ is the input hidden dimension, and $H_{\text{out}}$ is the output hidden dimension.

Megatron-LM implements two primary linear layer primitives: ColumnParallelLinear and RowParallelLinear.

1. Column-Parallel Linear Layer (ColumnParallelLinear)

In a column-parallel layer, the weight matrix $W$ is sliced vertically along its output dimension across $N$ GPUs:

$$W = \begin{bmatrix} W_1 & W_2 & \dots & W_N \end{bmatrix}$$

where each shard $W_i \in \mathbb{R}^{H_{\text{in}} \times (H_{\text{out}} / N)}$.

• Input: The full input activation tensor $X$ is replicated across all $N$ GPUs.
• Computation: Each GPU $i$ independently computes its local matrix multiplication:

$$Y_i = X W_i$$

• Output: Each GPU holds a slice of the output tensor $Y_i \in \mathbb{R}^{B \times (H_{\text{out}} / N)}$.

Crucially, if the subsequent operation in the neural network is an element-wise function (such as a GeLU, SiLU, or SwiGLU activation function), each GPU can apply that activation directly to its local slice $Y_i$ without any network communication.

2. Row-Parallel Linear Layer (RowParallelLinear)

In a row-parallel layer, the weight matrix $W$ is sliced horizontally along its input dimension across $N$ GPUs:

$$W = \begin{bmatrix} W_1 \\ W_2 \\ \vdots \\ W_N \end{bmatrix}$$

where each shard $W_i \in \mathbb{R}^{(H_{\text{in}} / N) \times H_{\text{out}}}$.

• Input: The input tensor $X$ must be sharded across GPUs along its channel dimension: $X = \begin{bmatrix} X_1 & X_2 & \dots & X_N \end{bmatrix}$, where $X_i \in \mathbb{R}^{B \times (H_{\text{in}} / N)}$.
• Computation: Each GPU $i$ computes a local partial matrix multiplication:

$$Y_i = X_i W_i \in \mathbb{R}^{B \times H_{\text{out}}}$$

• Output: To obtain the true mathematical output $Y = XW$, the partial matrix products from all $N$ GPUs must be summed together:

$$Y = \sum_{i=1}^{N} Y_i = \sum_{i=1}^{N} X_i W_i = XW$$

• Communication: This summation is executed using an All-Reduce (sum) collective communication primitive across the $N$ GPUs in the tensor-parallel group. After the All-Reduce, all GPUs hold the identical, complete output tensor $Y$.

Partitioning the Multi-Layer Perceptron (MLP) Block

In a standard Transformer architecture, the feed-forward network (MLP) consists of an up-projection matrix $W_{\text{in}} \in \mathbb{R}^{h \times 4h}$, a non-linear activation function $\sigma(\cdot)$, and a down-projection matrix $W_{\text{out}} \in \mathbb{R}^{4h \times h}$, where $h$ is the model hidden dimension:

$$\text{MLP}(X) = \sigma(X W_{\text{in}}) W_{\text{out}}$$

A naive implementation of tensor parallelism might attempt to synchronize activations between every single linear layer, adding catastrophic communication latency.

Megatron-LM solves this by pairing a ColumnParallelLinear layer directly with a RowParallelLinear layer:

1. Up-Projection / Gate (ColumnParallelLinear):

• The weight matrix $W_{\text{in}}$ is split column-wise into $[W_{\text{in}, 1}, W_{\text{in}, 2}, \dots, W_{\text{in}, N}]$, where each $W_{\text{in}, i} \in \mathbb{R}^{h \times (4h / N)}$.
• Each GPU $i$ takes identical replicated input $X$ and computes $Z_i = X W_{\text{in}, i}$.

1. Element-Wise Non-Linearity:

• Each GPU applies the activation function locally: $A_i = \sigma(Z_i)$.
• Because point-wise non-linearities (like GeLU or Swish) satisfy $\sigma([Z_1, Z_2]) = [\sigma(Z_1), \sigma(Z_2)]$, no cross-GPU communication is required.
• For gated architectures like SwiGLU used in LLaMA and modern open models, both the gate and up projections are sharded column-wise, and the element-wise multiplication $\text{SiLU}(X W_{\text{gate}, i}) \odot (X W_{\text{up}, i})$ occurs completely locally on each GPU.

1. Down-Projection (RowParallelLinear):

• The down-projection matrix $W_{\text{out}}$ is split row-wise into $[W_{\text{out}, 1}; W_{\text{out}, 2}; \dots; W_{\text{out}, N}]$, matching the sharded output dimensions of $A_i$.
• Each GPU $i$ computes the local product $O_i = A_i W_{\text{out}, i}$.

1. All-Reduce Collective:

• An All-Reduce sum is executed across the TP group to compute $O = \sum_{i=1}^N O_i$.

By chaining column-parallel into row-parallel, the entire MLP block requires exactly one All-Reduce communication in the forward pass, and exactly one All-Reduce in the backward pass (to synchronize input gradients across the column-parallel layer).

Partitioning Multi-Head and Grouped-Query Attention

The self-attention mechanism presents a similar structural opportunity. In Multi-Head Attention (MHA), the Query ($Q$), Key ($K$), and Value ($V$) projections map the hidden state $X \in \mathbb{R}^{B \times h}$ into $h_{\text{heads}}$ attention heads, each of dimension $d_k = h / h_{\text{heads}}$.

Megatron-LM shards the attention block by partitioning the attention heads across the $N$ GPUs:

$$\frac{h_{\text{heads}}}{N} \text{ heads per GPU}$$

Step-by-Step Attention Parallelism

1. Q, K, V Projections (ColumnParallelLinear):

• The projection matrices $W_Q, W_K, W_V \in \mathbb{R}^{h \times h}$ are sliced column-wise across the head dimension.
• Each GPU $i$ computes the local projections $Q_i, K_i, V_i \in \mathbb{R}^{B \times (h / N)}$ for its local subset of attention heads.
• No communication is required.

1. Local Self-Attention Computation:

• Each GPU independently evaluates the scaled dot-product attention for its assigned heads:

$$\text{Head}_{i} = \text{Softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i$$

• Because attention heads operate independently without cross-head interactions during the softmax and weighted value accumulation, this step requires zero communication.

1. Output Projection ($W_O$) (RowParallelLinear):

• The outputs of the local heads $\text{Head}_i$ are concatenated locally on each GPU.
• The output projection matrix $W_O \in \mathbb{R}^{h \times h}$ is sliced row-wise: $W_O = [W_{O, 1}; W_{O, 2}; \dots; W_{O, N}]$.
• Each GPU multiplies its local attention output by its row-sliced $W_{O, i}$.

1. All-Reduce Collective:

• A single All-Reduce sum aggregates the partial products across all $N$ GPUs to generate the full attention block output.

Grouped-Query Attention (GQA) Constraints

In modern architectures utilizing Grouped-Query Attention (GQA) or Multi-Query Attention (MQA), such as LLaMA 3, Mistral, and Qwen, the number of Key-Value heads ($N_{KV}$) is substantially smaller than the number of Query heads ($N_Q$).

When applying Tensor Parallelism to GQA architectures:

• The number of KV heads $N_{KV}$ must be divisible by the tensor parallel degree $TP$ ($N_{KV} \pmod{TP} = 0$).
• If $TP > N_{KV}$ (for example, attempting $TP=8$ on a model with only 4 KV heads), KV heads must either be duplicated across ranks or the model must employ sequence/context parallelism to distribute compute.

Sequence Parallelism: Eliminating Redundant Activation Memory

In the standard Megatron-LM formulation, the operations outside the MLP and Attention blocks—specifically Layer Normalization (or RMSNorm), Dropout, and residual additions—are duplicated across all GPUs in the TP group. Each GPU holds identical copies of the full activation tensor $X \in \mathbb{R}^{s \times b \times h}$ (where $s$ is sequence length and $b$ is batch size).

As sequence lengths grew to 32K, 128K, and beyond, this duplicated activation footprint became a dominant memory bottleneck.

In 2022, NVIDIA researchers introduced Sequence Parallelism (SP) in Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022).

Standard Megatron-LM:
[LayerNorm (Replicated)] -> [ColumnParallel Linear] -> [RowParallel Linear] -> [All-Reduce]

Megatron-LM with Sequence Parallelism:
[LayerNorm (Sharded s/TP)] -> [All-Gather] -> [ColumnParallel Linear] -> [RowParallel Linear] -> [Reduce-Scatter] -> [LayerNorm (Sharded s/TP)]

Transforming the Collectives

Sequence Parallelism observes that LayerNorm and Dropout operate element-wise along the hidden dimension $h$, independent across sequence tokens. Therefore, the activation tensor can be partitioned along the sequence dimension $s$ into slices of size $s / N$:

1. Before Column-Parallel Linear: An All-Gather collective gathers the sequence slices across the TP group, restoring the full sequence tensor $s \times b \times h$ right before the column-parallel projection.
2. After Row-Parallel Linear: Instead of performing an All-Reduce (which sums and replicates the output), the framework executes a Reduce-Scatter collective. This sums the partial results while scattering the output along the sequence dimension, leaving each GPU with only an $s / N$ slice of activations.

Mathematically, an All-Reduce operation is composed of a Reduce-Scatter followed by an All-Gather, transferring $2 \times \frac{N-1}{N} \times \text{data size}$ in ring topologies.

By replacing each All-Reduce with a Reduce-Scatter at the end of a block and an All-Gather at the beginning of the next block, Sequence Parallelism incurs zero additional communication volume while reducing the activation memory of LayerNorm and Dropout by a factor of $N$ (up to a $5\times$ reduction in overall layer activation memory).

Communication Budget and Scaling Limits

A complete Transformer layer parallelized with Megatron-LM tensor parallelism exhibits a very specific communication profile:

| Layer Component | Forward Communication | Backward Communication | | :--- | :--- | :--- | | Multi-Head Attention | 1 All-Reduce (or 1 Reduce-Scatter + 1 All-Gather) | 1 All-Reduce (or 1 Reduce-Scatter + 1 All-Gather) | | Feed-Forward Network (MLP) | 1 All-Reduce (or 1 Reduce-Scatter + 1 All-Gather) | 1 All-Reduce (or 1 Reduce-Scatter + 1 All-Gather) | | Total per Transformer Layer | 2 Collectives | 2 Collectives |

For a model with $L$ layers, a single training step requires $4L$ high-volume collective operations across the TP group.

Because these collectives occur synchronously inside the critical execution path of every layer, Tensor Parallelism is governed by strict communication limits:

• Intra-Node Scaling ($TP \le 8$): Within a single 8-GPU server chassis (e.g., HGX H100/H200 or B200), NVLink mesh bandwidth provides sub-microsecond latency and hundreds of gigabytes per second of transfer speed. Here, $TP=2, 4, 8$ achieves near-linear compute scaling and efficient memory distribution.
• Inter-Node Scaling ($TP > 8$): Extending Tensor Parallelism across network switches (via InfiniBand or RoCE) introduces network latency that quickly dominates GPU compute time, causing severe GPU underutilization.

Consequently, modern large-scale training systems cap Tensor Parallelism at the physical node boundary ($TP=8$), composing it with Pipeline Parallelism (PP) and Data Parallelism with ZeRO/FSDP across servers to train models across thousands of nodes.

Sources

• Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053
• Narayanan, D., et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. arXiv:2104.04473
• Korthikanti, V. A., et al. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198
• Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054
• Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135