Tensor Parallelism: Column, Row, and Megatron Patterns

Training a large language model with hundreds of billions of parameters is not something a single GPU can do on its own. A single NVIDIA A100 has 80 GB of memory, and a model like GPT-3 requires roughly 350 GB just to store its weights in 16-bit precision, before accounting for optimizer states, gradients, and activations. The simplest response to this constraint is data parallelism, where each GPU holds a complete copy of the model but processes a different slice of the minibatch. But data parallelism fails precisely when the model itself does not fit on one device. You cannot copy what you cannot hold.

Tensor parallelism solves this differently. Instead of replicating the entire model across GPUs, it splits individual tensors across multiple devices. Each GPU holds only a shard of each weight matrix, performs computation on its shard, and then communicates with peer GPUs to assemble the full result. The model is not replicated; it is partitioned. This lets you train models whose individual layers would overflow any single device's memory.

The idea sounds straightforward, but realizing it efficiently requires care. A naive partitioning strategy might simply distribute layers to different GPUs (inter-layer parallelism), but this creates severe load-balancing problems and requires activations to travel from one device to another between every layer. Tensor parallelism takes a different approach: it operates within a single layer, distributing the matrix multiplications that dominate transformer computation. By splitting weight matrices in mathematically compatible ways, tensor parallelism allows each GPU to do useful work simultaneously on every layer, with communication happening at precisely defined synchronization points rather than at every layer boundary.

This chapter covers the mechanics of tensor parallelism in depth: how weight matrices are split along columns or rows, what communication is required to combine shards into correct outputs, how sequences of column-parallel and row-parallel layers are chained to minimize synchronization, and how the Megatron-LM system from NVIDIA Research formalized these ideas into a practical training framework used for some of the largest models in existence. We will build understanding from first principles, working through the math of why the splits are correct before looking at implementation.

As we discussed in earlier chapters on the transformer architecture, the dominant compute in a transformer layer comes from two sources: the multi-head attention mechanism and the two-layer feed-forward network (FFN). Tensor parallelism in Megatron style targets both of these components, splitting the weight matrices in carefully chosen directions to enable maximum parallelism with minimum communication.

Before diving into how tensor parallelism works, it helps to be precise about the problem it solves and why the solutions take the forms they do.

A standard linear layer computes $Y = XA$, where $X \in \mathbb{R}^{b \times d}$ is a batch of $b$ tokens each with hidden dimension $d$, and $A \in \mathbb{R}^{d \times h}$ is a weight matrix mapping to some output dimension $h$. The output $Y \in \mathbb{R}^{b \times h}$ is dense: every output element depends on every input element. At first glance, this density suggests that any partitioning of the computation must involve data dependencies, and therefore communication, across devices.

But the density of the output does not mean the computation is inseparable. Matrix multiplication has a fundamental block structure that allows it to be decomposed in multiple ways. Understanding this structure is the conceptual key to tensor parallelism.

When $d$ and $h$ are large (e.g., $d = 12288$ for GPT-3's hidden size), the weight matrix $A$ can be enormous. A single linear layer in GPT-3's attention has four weight matrices for Q, K, V, and output projections, each of size $12288 \times 12288$, consuming roughly 576 MB each in fp16. Across 96 layers, the attention weights alone exceed 200 GB. The FFN layers add another 150 GB or more. No single device, regardless of generation, can hold this.

The key insight enabling tensor parallelism is that matrix multiplication is separable. Different columns of $A$ contribute to different columns of $Y$ without any cross-dependency, and different rows of $A$ can each be multiplied with the corresponding input slice independently, with results accumulated at the end. This separability is the mathematical foundation for both column-parallel and row-parallel splitting strategies. The question is not whether the computation can be split, but how to split it such that the communication required to assemble final results is minimized and falls at convenient points in the computation graph.

The choice of splitting direction is not arbitrary. Different layers and positions in the transformer have different mathematical structures, and the clever design of tensor parallelism is that it exploits these structures to pair column splits with row splits in a way that eliminates intermediate communication. The column-parallel output of one layer is exactly the right input format for a row-parallel next layer, creating a natural pipeline that requires only a single all-reduce per sub-layer rather than a communication step between every layer.

Column parallelism splits a weight matrix along its output (column) dimension. Given $N$ GPUs with indices $0, 1, \ldots, N-1$, we partition the weight matrix $A \in \mathbb{R}^{d \times h}$ into $N$ column slices:

where each $A_i \in \mathbb{R}^{d \times (h/N)}$.

Each GPU $i$ stores its slice $A_i$ and computes:

where:

• $X \in \mathbb{R}^{b \times d}$ is the full input, replicated on all GPUs
• $A_i \in \mathbb{R}^{d \times (h/N)}$ is the column shard on GPU $i$
• $Y_i \in \mathbb{R}^{b \times (h/N)}$ is the partial output on GPU $i$

The partial outputs $Y_i$ together form the full output: concatenating $[Y_0, Y_1, \ldots, Y_{N-1}]$ along the column dimension recovers the full $Y = XA \in \mathbb{R}^{b \times h}$. This is correct because matrix multiplication with a partitioned weight matrix produces a partitioned result:

The critical observation is that each GPU's computation $Y_i = XA_i$ is completely independent. No communication is needed during the forward pass to produce $Y_i$; each device can compute its shard without knowledge of what the other devices are doing. The price is that the input $X$ must be fully available on every GPU before computation begins. In practice, this is handled by an all-gather operation at the start of the layer or by keeping the input in the replicated state (which happens automatically when tensor parallelism is paired with data parallelism).

Why is replicating the input acceptable? Because the input to each transformer layer is the residual stream, which after an all-reduce from the previous layer is identical on all GPUs. There is no communication cost to have a replicated input in steady-state operation, since the previous layer's row-parallel all-reduce already distributed it.

Column parallelism is natural for the query, key, and value projections in multi-head attention. Recall that multi-head attention splits the model's hidden dimension into $N_h$ heads, where each head operates on a $d_k = d / N_h$ dimensional subspace. When you have $N$ GPUs and split the weight matrix along the head dimension, each GPU naturally handles $N_h / N$ attention heads. The column-parallel split corresponds exactly to distributing heads across devices.

This alignment between tensor parallelism and multi-head attention is not a coincidence. The original Megatron-LM paper was designed with this correspondence in mind. Because attention heads are completely independent of one another during the self-attention computation (they attend to the full sequence in parallel subspaces without interacting), you can assign entire heads to individual GPUs and have each GPU compute attention for its assigned heads without ever communicating with other GPUs during the attention computation itself. The interaction between heads happens only when their outputs are projected back to the model's hidden dimension, and that projection is handled by the row-parallel output projection.

For an attention layer with hidden size $d = 4096$ and 32 heads on 4 GPUs, each GPU holds the Q, K, V projection weights for 8 heads: matrices of size $4096 \times 1024$ rather than $4096 \times 4096$. Each GPU attends over its 8 heads independently and produces attention outputs of size $b \times 1024$. The output projection then acts on these 1024-dimensional features row-parallel, accumulating the full 4096-dimensional output across GPUs via all-reduce.

The feed-forward network (FFN) also benefits from column parallelism. A standard transformer FFN applies two linear transformations with a nonlinearity between them:

where:

• $X \in \mathbb{R}^{b \times d}$ is the input batch with $b$ tokens and hidden dimension $d$
• $W_1 \in \mathbb{R}^{d \times 4d}$ is the expansion weight matrix, increasing the dimension by a factor of 4
• $b_1 \in \mathbb{R}^{4d}$ is the expansion bias
• $\sigma(\cdot)$ is the activation function (GeLU in most modern transformers, ReLU in older ones)
• $W_2 \in \mathbb{R}^{4d \times d}$ is the contraction weight matrix, reducing back to $d$
• $b_2 \in \mathbb{R}^{d}$ is the output bias

Column parallelism applied to $W_1$ splits the expansion across GPUs. GPU $i$ receives the full input $X$ and computes $XW_{1,i}$ using its column shard $W_{1,i} \in \mathbb{R}^{d \times d}$ (for a 4-way split of a $4d$-dimensional FFN). The activation function is then applied locally because GeLU and ReLU operate element-wise on each dimension independently. This is a subtle but important point: element-wise nonlinearities are trivially parallelizable because each output element depends on exactly one input element, which lies entirely within the local shard.

After the nonlinearity, each GPU holds $\hat{Z}_i = \sigma(XW_{1,i}) \in \mathbb{R}^{b \times d}$. These shards together form the full intermediate activation $\hat{Z} = [\hat{Z}_0 \mid \hat{Z}_1 \mid \cdots \mid \hat{Z}_{N-1}] \in \mathbb{R}^{b \times 4d}$, ready to be consumed by the row-parallel second layer.

Row parallelism splits a weight matrix along its input (row) dimension. Given the weight matrix $A \in \mathbb{R}^{d \times h}$, we partition along the row dimension into $N$ slices:

where each $A_i \in \mathbb{R}^{(d/N) \times h}$.

Each GPU $i$ stores its row shard $A_i$ and receives a corresponding slice $X_i \in \mathbb{R}^{b \times (d/N)}$ of the input, then computes:

where:

• $X_i \in \mathbb{R}^{b \times (d/N)}$ is the input shard on GPU $i$
• $A_i \in \mathbb{R}^{(d/N) \times h}$ is the row shard on GPU $i$
• $Z_i \in \mathbb{R}^{b \times h}$ is the partial result on GPU $i$

The full output requires summing the partial results:

This is an all-reduce operation: each device produces a partial sum over the full output dimension $h$, and summing these partial results gives the correct final output. Unlike column parallelism, where the output shards can be consumed independently (each GPU's shard feeds into a corresponding downstream shard), row parallelism requires synchronization to assemble the full output.

The correctness of this decomposition follows directly from the block form of matrix multiplication. If we write the full input as $X = [X_0 \mid X_1 \mid \cdots \mid X_{N-1}]$ (conceptually partitioned along the column dimension) and the weight matrix as stacked row blocks, then:

Each term $X_i A_i$ is a partial contribution to the full output, and summing them gives the correct result. Note that each partial result $Z_i = X_i A_i$ has the full output shape $\mathbb{R}^{b \times h}$ rather than a shard of it. This is the defining feature of row parallelism: the partial results are over the full output space, and correctness requires their sum, not their concatenation.

The design of row parallelism makes most sense when you see it as the second half of a pair. The output of a column-parallel layer is a feature-sharded tensor, where GPU $i$ holds the columns $[i \cdot h/N, (i+1) \cdot h/N)$ of the full hidden state. This is precisely the input format that row parallelism expects: GPU $i$ should receive input shard $X_i$ corresponding to its row shard $A_i$. When you place a row-parallel layer directly after a column-parallel layer, the output of the first layer is already distributed in exactly the right way for the second layer to consume without any additional communication.

This chaining property is what makes tensor parallelism efficient. Without it, you would need an all-gather after the column-parallel layer to reassemble the full activation before the row-parallel layer, adding an extra communication step. With careful layer pairing, the communication is deferred to the very end, where a single all-reduce assembles the final output.

The real power of tensor parallelism emerges from chaining column-parallel and row-parallel layers such that the output of one feeds naturally into the other, eliminating unnecessary communication.

Consider the two-layer FFN with $W_1$ (expansion) followed by $W_2$ (contraction):

1. Column-parallel $W_1$: GPU $i$ receives the full input $X$ and computes $Z_i = XW_{1,i}$ where $W_{1,i}$ is the $i$-th column shard of $W_1$. After applying the nonlinearity, each GPU holds $\hat{Z}_i = \text{GeLU}(Z_i)$, a shard of the intermediate activations. No communication has occurred at this point.

2. Row-parallel $W_2$: GPU $i$ holds the $i$-th row shard $W_{2,i}$ of $W_2$, where this partitioning is along the input dimension of $W_2$ (which is the same dimension along which $\hat{Z}_i$ is partitioned). Each GPU computes $Y_i = \hat{Z}_i W_{2,i}$, a partial sum over the full output space.

3. All-reduce: The partial sums $Y_i$ are summed across GPUs to produce the full output $Y = \sum_i Y_i$.

The key observation is that the output of column parallelism (sharded along the feature dimension) is exactly the right input format for row parallelism (which expects sharded inputs along the same dimension). This means no communication is needed between the two layers. The only communication happens at the end of the row-parallel layer, in the all-reduce that assembles the final output.

For a two-layer block (FFN or attention), this yields a total of one all-reduce in the forward pass. Without this careful pairing, a naive implementation would require an all-gather after the column-parallel layer (to reassemble the full activations before the second layer), costing an extra communication step. This design choice, seeming minor in the abstract, has a substantial impact on training throughput: every saved all-gather at model scale represents gigabytes of data that do not need to traverse the interconnect on every forward pass.

A subtle but important aspect of the FFN chaining is where the nonlinearity fits. The GeLU (or ReLU) function is applied between the two linear layers. In column-parallel mode, each GPU applies GeLU locally to its activation shard $Z_i$. This works correctly because GeLU is an element-wise function: $\text{GeLU}(z)$ depends only on the scalar value $z$, not on any other element in the tensor. There are no cross-feature dependencies that would require communication.

This would not be the case for operations like softmax or layer normalization, which are reduction operations over features. Softmax requires the full feature vector to compute the normalizing denominator, and layer normalization requires statistics (mean and variance) over the feature dimension. These operations cannot be applied to feature shards independently. This is why tensor parallelism requires careful attention to which operations are tensor-parallelizable and which require full-feature access, a consideration that motivated the sequence parallelism extension discussed later.

Multi-head attention can be decomposed similarly to the FFN. The QKV projections are column-parallel: each GPU handles a subset of attention heads independently. The output projection maps from the concatenated head outputs back to the model dimension, and it is row-parallel.

Why does this work? Because attention heads are completely independent of one another in standard self-attention. Each head attends to the full sequence but produces its output in a separate subspace of the representation. When you distribute heads across GPUs, each GPU attends over its subset of heads without needing to know what other GPUs' heads are doing. The interaction between heads happens only at the output projection layer, which aggregates all head outputs and is naturally handled by the row-parallel reduce.

Concretely, the QKV weight matrices are split column-wise across GPUs, with each GPU holding the Q, K, V weights for its assigned heads. Each GPU computes attention scores and applies them to values independently, since heads do not communicate during the attention computation. The attention output for each GPU's heads is then passed through the row-parallel output projection, which accumulates partial sums across GPUs in an all-reduce.

This design has a pleasing symmetry with the FFN decomposition: both sub-layers require exactly one all-reduce in the forward pass, giving two all-reduces total per transformer layer. The all-reduces happen at the output projection of attention and at $W_2$ of the FFN. Between the two sub-layers, residual connections are added to the accumulated results. Because every GPU holds the full output after each all-reduce, the residual stream is synchronized and identical on all devices before the next sub-layer begins, making sure that the replicated-input assumption for the next column-parallel layer holds automatically.

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.

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.

Let's trace through a specific numerical example to make the computation concrete. Suppose we have a feed-forward network with input dimension $d = 4$, hidden dimension $h = 8$, and we are using $N = 2$ GPUs.

The full weight matrix $W_1 \in \mathbb{R}^{4 \times 8}$ is split column-wise:

where $W_{1,0} \in \mathbb{R}^{4 \times 4}$ lives on GPU 0 and $W_{1,1} \in \mathbb{R}^{4 \times 4}$ lives on GPU 1.

Given an input $X \in \mathbb{R}^{2 \times 4}$ (batch of 2 tokens, each with 4 features), both GPUs have the full input. GPU 0 computes:

and GPU 1 computes:

After the GeLU nonlinearity applied locally on each GPU, the activations $\hat{Z}_0$ and $\hat{Z}_1$ together form the full intermediate representation $[\hat{Z}_0 \mid \hat{Z}_1] \in \mathbb{R}^{2 \times 8}$. Neither GPU needed to communicate with the other to reach this point.

The second weight matrix $W_2 \in \mathbb{R}^{8 \times 4}$ is split row-wise:

where $W_{2,0} \in \mathbb{R}^{4 \times 4}$ on GPU 0 and $W_{2,1} \in \mathbb{R}^{4 \times 4}$ on GPU 1.

GPU 0 multiplies its activation shard by its weight shard: $Y_0 = \hat{Z}_0 W_{2,0} \in \mathbb{R}^{2 \times 4}$. GPU 1 multiplies its activation shard by its weight shard: $Y_1 = \hat{Z}_1 W_{2,1} \in \mathbb{R}^{2 \times 4}$.

The final output is obtained by all-reducing (summing) across GPUs:

This is identical to the full computation $Y = \text{GeLU}(XW_1) W_2$, but computed with weight matrices that fit on individual GPUs. The only communication is the final all-reduce that sums $Y_0$ and $Y_1$.

What makes this example particularly illuminating is the absence of intermediate communication. Between the column-parallel step and the row-parallel step, GPUs 0 and 1 proceed completely independently. GPU 0 never needs to know what $Z_1$ looks like. GPU 1 never needs to know what $Z_0$ looks like. They only need to coordinate once, at the very end, to produce the final output. This one communication per two linear layers is the core efficiency of Megatron-style tensor parallelism.

Let's implement tensor parallelism from scratch using PyTorch's distributed computing primitives. We'll simulate multiple GPUs using CPU processes, then show how the same logic applies on actual multi-GPU hardware.

First, we set up the distributed environment. In real training, this is handled by the launcher (e.g., torchrun), but for this tutorial, we'll configure it explicitly.

The process group is the fundamental abstraction for collective communication in PyTorch. Each GPU (or CPU process in simulation) has a rank within the group, and collective operations like all-reduce act across all ranks in the group. Real Megatron training uses NCCL backend for GPU-to-GPU communication via NVLink or InfiniBand; we use Gloo here for CPU-based simulation.

Now let's implement a column-parallel linear layer. Each process holds a column shard of the weight matrix and computes on the full input.

Notice that the forward method contains no communication. This is the defining property of column parallelism: each rank can compute its output shard entirely from local data. The weight shard has shard_size rows (corresponding to the output columns assigned to this rank) and in_features columns (the full input dimension). F.linear computes x @ weight.T + bias, which for our weight shape gives x @ (shard_size, in_features).T = x @ (in_features, shard_size), producing the correct output shard of shape (batch, seq_len, shard_size).

The row-parallel layer receives a sharded input (output from a column-parallel layer) and produces a full output via all-reduce.

The all-reduce call is the only communication in the row-parallel forward pass. Notice the bias handling: only rank 0 holds the bias term, which is added after the all-reduce. This is correct because the bias is a full-output-dimensional vector (not split), and adding it once (on rank 0) is equivalent to adding it to the summed output. Adding it on every rank would multiply the bias by world_size, which would be incorrect.

We can now assemble a complete tensor-parallel FFN by chaining column-parallel and row-parallel layers.

The elegance of this implementation is in what is absent: there is no explicit communication between the column-parallel and row-parallel layers. The sharded output of column_parallel flows directly into row_parallel as a sharded input, with no reshaping, no synchronization, and no data movement. The single dist.all_reduce inside row_parallel handles all necessary communication for the entire two-layer block.

To verify our tensor-parallel implementation produces the same result as a standard sequential FFN, we simulate two-GPU behavior by running the full computation twice with rank-0 and rank-1 shards.

The near-zero numerical difference confirms that the column-parallel and row-parallel chaining produces exactly the same result as the sequential computation. Any remaining difference is due to floating-point arithmetic order: when you sum the same values in a different sequence, rounding errors at the bit level can introduce tiny discrepancies that are entirely correct behavior. The correctness proof guarantees mathematical equivalence; the implementation delivers numerical equivalence up to floating-point precision.

A key benefit of tensor parallelism is reduced memory per GPU. Let's compute how memory requirements scale with the tensor-parallel degree.

Understanding when tensor parallelism becomes the bottleneck requires estimating how much time the all-reduce operations consume relative to compute. Let's build this analysis.

The key parameters for tensor parallelism configuration are:

• tensor_parallel_degree (tp): Number of GPUs across which each layer is split. Must divide all weight dimensions evenly. Typically 1, 2, 4, or 8 within a single node.
• pipeline_parallel_degree (pp): Number of stages in pipeline parallelism, used alongside tensor parallelism. A cluster of $tp \times pp$ GPUs forms one complete model replica.
• data_parallel_degree (dp): Number of replicas used for data parallelism. Total GPUs $= tp \times pp \times dp$.
• sequence_length: Affects activation memory proportionally. Longer sequences amplify the benefit of sequence parallelism.
• all_reduce_backend: Use nccl for NVIDIA GPU communication (exploits NVLink/InfiniBand), gloo for CPU.

Tensor parallelism is powerful, but it comes with real constraints that practitioners must account for when designing training infrastructure.

The most immediate limitation is that every dimension you want to shard must be divisible by the tensor-parallel degree. This sounds like a minor implementation detail, but it shapes model architecture choices from the ground up. A model with a hidden dimension of 4096 and 32 attention heads can be tensor-parallelized with degrees 1, 2, 4, 8, 16, and 32 (all of which divide both 4096 and 32 evenly), but a model with a hidden dimension of 3840 and 30 heads is more constrained. If you want degree 8, you need all key dimensions to be divisible by 8.

Large research teams deploying Megatron-style training consistently choose hidden dimensions and head counts that are highly divisible: powers of 2 or multiples of 128 are standard. GPT-3's hidden dimension of 12288 is $96 \times 128$, divisible by 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, and more. This kind of divisibility is not accidental; it is designed to accommodate tensor parallelism at multiple degrees without architectural changes.

FFN dimensions are similarly constrained. The standard FFN expansion factor of 4 gives a hidden dimension of $4d$, and the product $4d$ must be divisible by the tensor-parallel degree. In practice, since $4$ itself is divisible by 2 and 4, any hidden dimension divisible by 8 will work with degrees up to 8. At degree 16, you need $4d$ to be divisible by 16, so $d$ must be divisible by 4, which is true for all powers of 2.

The second major limitation is communication bandwidth dependency. The all-reduce operations that make tensor parallelism correct are not free; they consume network bandwidth proportional to the activation tensor size at every layer. On modern DGX nodes, NVLink provides 600 GB/s bidirectional bandwidth, which is typically sufficient for tensor-parallel degrees up to 8 within a node. Crossing node boundaries over InfiniBand introduces a 10-20x bandwidth reduction, making tensor parallelism inefficient at those scales.

This is why Megatron-style systems typically combine tensor parallelism (within a node) with pipeline parallelism (across nodes). Pipeline parallelism communicates only at pipeline stage boundaries, transmitting only the activations for one microbatch at the boundary layer rather than all-reducing full activations at every layer. This reduces cross-node communication volume by orders of magnitude compared to applying tensor parallelism across nodes. The combination of intra-node tensor parallelism and inter-node pipeline parallelism became the standard architecture for large-scale training, codified in the Megatron v2 and v3 papers.

Tensor parallelism increases the degree of synchronization required within a group. With $N$ GPUs in a tensor-parallel group, all $N$ devices must synchronize at every all-reduce. If one GPU stalls due to a hardware issue, memory pressure, or kernel scheduling delay, all other GPUs in the group wait. This tight coupling means that hardware failures or stragglers have immediate impact on training throughput.

In contrast, data parallelism is much more loosely coupled: each GPU processes its own microbatch independently, and gradient averaging happens once per step rather than twice per transformer layer. Pipeline parallelism falls in between: stages must synchronize at pipeline boundaries, but each stage can compute independently otherwise. The tight synchronization of tensor parallelism is one reason it is restricted to within a single node, where communication is fast enough that synchronization delays are measured in microseconds rather than milliseconds.

Implementing tensor parallelism correctly requires careful attention to weight initialization. When a weight matrix is split across $N$ GPUs and each shard is initialized independently from a normal distribution, the effective initialization variance seen by the full layer computation differs from what a single-GPU initialization would produce.

Consider column parallelism: each GPU initializes its shard $W_i \in \mathbb{R}^{d \times (h/N)}$ from $\mathcal{N}(0, \sigma^2)$. The full output $Y = XA$ has variance proportional to $d \sigma^2$ (by standard variance analysis of matrix-vector products). This is the same as if the full matrix had been initialized with variance $\sigma^2$ on a single GPU. So far, so good.

For row parallelism, the situation is different. Each GPU computes a partial sum $Z_i = X_i A_i$, and the all-reduce computes $Y = \sum_i Z_i$. The variance of $Y$ is $N$ times the variance of a single $Z_i$ (since the $Z_i$ are independent). To maintain the correct variance, Megatron-LM scales the initialization of row-parallel weight shards by $1/\sqrt{N}$. This ensures that the aggregated output has the same distribution as the single-GPU case. Without this adjustment, deeper layers would see activations with inflated variance, potentially destabilizing training.

Choosing the right tensor-parallel degree in practice involves balancing several competing factors. The key considerations are:

• Memory capacity: Use a degree large enough that the model fits on each GPU with room for optimizer states and activations.
• Divisibility: All key dimensions (hidden size, number of heads, FFN intermediate size) must be divisible by the degree.
• Communication budget: On NVLink hardware, degrees 2-8 are typically efficient. Higher degrees diminish returns because the per-layer communication overhead grows while the per-GPU compute shrinks.
• Node boundary alignment: Keep tensor parallelism within a single node to exploit NVLink. Use pipeline or data parallelism for cross-node scaling.

The interaction between these factors means that choosing a tensor-parallel degree is fundamentally a hardware-aware decision. A model that trains efficiently with degree 8 on an A100 DGX node might train better with degree 4 on older hardware with slower NVLink, because the same communication volume represents a larger fraction of compute time.

Despite these limitations, tensor parallelism has proven essential for training the largest language models. Models like GPT-3, Megatron-Turing NLG (530 billion parameters), and subsequent systems from Google, Meta, and Anthropic all rely on tensor parallelism as a core component of their training infrastructure. The technique scales naturally with the number of high-bandwidth links available and integrates cleanly with data and pipeline parallelism to enable distributed training at scales that were previously considered impractical.

We will explore how tensor parallelism integrates with pipeline parallelism and data parallelism in the next chapter, which covers 3D parallelism strategies for training at extreme scale.

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.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about tensor parallelism.