Tensor Parallelism (TP)

While Data Parallelism (DP) effectively utilizes multiple devices by replicating the model and processing different data slices, it requires each device to hold the entire model. For truly massive models, even a single replica might exceed the memory capacity of one accelerator. Furthermore, some individual operations within the model, like the large linear transformations in feed-forward networks or attention mechanisms, can become computational bottlenecks. This is where Tensor Parallelism (TP), sometimes called intra-layer model parallelism, becomes necessary.

Instead of splitting the data or the sequence of layers, Tensor Parallelism splits the actual tensors (weights, activations, gradients) within a specific operation across multiple devices. This allows computations on these large tensors to be performed in parallel, distributing both the memory footprint and the computational load of individual layers.

Splitting Linear Layers

The most straightforward application of TP is in large linear layers, which are fundamental components of both the multi-layer perceptron (MLP) and attention blocks in Transformers. Consider a linear transformation $Y = XA$, where $X$ is the input activation and $A$ is the weight matrix. We can parallelize this matrix multiplication across, for example, two devices in two primary ways: column parallelism and row parallelism.

Column Parallelism

In column parallelism, the weight matrix $A$ is split vertically (along its columns) across devices. If we have two devices, we split $A$ into $A = [A_1, A_2]$. The input $X$ is typically broadcast or made available to both devices. Each device then computes a part of the output:

• Device 1 computes: $Y_1 = X A_1$
• Device 2 computes: $Y_2 = X A_2$

The final output $Y$ is obtained by concatenating the results along the column dimension: $Y = [Y_1, Y_2]$.

Column parallelism for $Y = XA$. The weight matrix $A$ is split into $A_1$ and $A_2$. The input $X$ is multiplied by each part on separate GPUs. The results $Y_1$ and $Y_2$ are concatenated to form the final output $Y$.

Mathematically, the forward pass involves $X A = X [A_1, A_2] = [X A_1, X A_2]$. The backward pass requires gradients with respect to $X$ and $A$. The gradient $\frac{\partial L}{\partial A}$ is naturally partitioned ($\frac{\partial L}{\partial A_1}$, $\frac{\partial L}{\partial A_2}$). However, computing the gradient $\frac{\partial L}{\partial X}$ requires summing contributions from both paths: $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y_1} A_1^T + \frac{\partial L}{\partial Y_2} A_2^T$. This summation is typically achieved using an all-reduce communication collective across the devices holding $A_1$ and $A_2$.

Row Parallelism

In row parallelism, the weight matrix $A$ is split horizontally (along its rows). For two devices, $A = \begin{bmatrix} A_1 \\ A_2 \end{bmatrix}$. The input $X$ is also considered partitioned along its columns (often because it's the output of a preceding column-parallel layer). Let's consider the input $X$ available on both devices for simplicity in this isolated view. Each device computes a partial result based on its slice of the weights:

• Device 1 computes: $Y_1 = X A_1$ (Note: This formula is slightly simplified. In practice, $X$ would be partitioned, e.g., $X=[X_1, X_2]$, and Device 1 computes $X_1 A_1$, Device 2 computes $X_2 A_2$. We show the common pattern used in combination with column parallelism later).
• Device 2 computes: $Y_2 = X A_2$

Unlike column parallelism, the final output $Y$ is the sum of the partial results: $Y = Y_1 + Y_2$.

Row parallelism for $Y = XA$. The weight matrix $A$ is split into $A_1$ and $A_2$ (row-wise). Partial results $Y_1$ and $Y_2$ are computed on separate GPUs. An all-reduce operation sums these results to produce the final output $Y$.

Mathematically, $Y = X \begin{bmatrix} A_1 \\ A_2 \end{bmatrix}$ isn't standard matrix multiplication notation if $X$ is treated as a single block. The actual operation in context (e.g., following a column-split layer) is $Y = X_1 A_1 + X_2 A_2$, where $X=[X_1, X_2]$. The forward pass requires an all-reduce operation to perform this summation across devices. The backward pass for $\frac{\partial L}{\partial X}$ involves partitioned gradients based on the partitioned $A$: $\frac{\partial L}{\partial X_1} = \frac{\partial L}{\partial Y} A_1^T$ and $\frac{\partial L}{\partial X_2} = \frac{\partial L}{\partial Y} A_2^T$. The gradient $\frac{\partial L}{\partial A}$ is directly computed on each device for its corresponding weight slice.

Tensor Parallelism in Transformers

Tensor Parallelism is typically applied strategically within Transformer blocks to balance computation and minimize communication overhead. A common pattern, popularized by frameworks like Megatron-LM, combines column and row parallelism within the MLP block and applies similar partitioning to the attention mechanism.

MLP Block: A standard Transformer MLP block computes $Y = \text{Activation}(XA)B + X$ (including residual connection). TP is applied as follows:

1. First Linear Layer ($XA$): Use column parallelism for matrix $A$. Split $A = [A_1, A_2]$. Compute $Z_1 = XA_1$ and $Z_2 = XA_2$ on respective devices. The output is $[Z_1, Z_2]$. This step involves an all-reduce for the gradient calculation in the backward pass, but the forward pass requires no communication if $X$ is already available on both devices.
2. Activation: Apply the activation function (e.g., GeLU, SwiGLU) element-wise to the partitioned outputs: $G_1 = \text{Activation}(Z_1)$, $G_2 = \text{Activation}(Z_2)$. No communication needed.
3. Second Linear Layer ($GB$): Use row parallelism for matrix $B$. Split $B = \begin{bmatrix} B_1 \\ B_2 \end{bmatrix}$. Device 1 computes $Y_1 = G_1 B_1$, and Device 2 computes $Y_2 = G_2 B_2$.
4. Summation: The final output of the linear transformation is $Y = Y_1 + Y_2$. This summation is performed using an all-reduce collective in the forward pass. The gradient calculation in the backward pass requires no all-reduce for the inputs $G_1, G_2$.

This combination cleverly arranges the parallelism such that the communication required for the forward pass (all-reduce after the second linear layer) and the backward pass (all-reduce for $\nabla X$ after the first linear layer) do not overlap unnecessarily, optimizing the flow.

Attention Block: TP can also be applied to the self-attention mechanism.

1. Q, K, V Projections: The weight matrices $W_Q$, $W_K$, and $W_V$ used to project the input into queries, keys, and values are often split using column parallelism, similar to the first linear layer of the MLP. $Q = XW_Q$, $K = XW_K$, $V = XW_V$ are computed in parallel across devices, resulting in partitioned $Q = [Q_1, Q_2]$, $K = [K_1, K_2]$, $V = [V_1, V_2]$.
2. Attention Score Calculation: The attention scores $S = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})$ involve matrix multiplications between the partitioned Q and K tensors. This requires careful implementation and communication (e.g., all-gather operations might be needed depending on the specific splitting strategy) to compute the full attention matrix or its necessary parts on each device.
3. Value Aggregation: The output $O = SV$ involves multiplying the attention scores $S$ with the partitioned value tensors $V$.
4. Output Projection: The final linear layer in the attention block ($O W_O$) is typically parallelized using row parallelism, similar to the second linear layer of the MLP. This requires an all-reduce in the forward pass to combine the results.

The implementation details within the attention block can be intricate, involving optimized kernels and communication patterns to handle the sequence length and head dimensions efficiently.

Communication Costs

A significant drawback of Tensor Parallelism is the increased communication overhead compared to Data Parallelism. While DP typically involves one all-reduce per training step (for gradients), TP introduces communication within the forward and backward passes of each Transformer block.

• Column Parallelism: Requires communication (all-reduce) for gradient computation w.r.t input $X$ during the backward pass.
• Row Parallelism: Requires communication (all-reduce) for summing outputs during the forward pass.
• Attention: Can involve additional communication like all-gather depending on implementation.

These communication operations (like all-reduce) involve synchronizing and exchanging data across all devices participating in the TP group. The volume of data exchanged depends on the size of the activations or gradients being communicated. This communication cost scales with the number of devices in the TP group and can become a performance bottleneck if the interconnect bandwidth between devices (e.g., NVLink, InfiniBand) is insufficient relative to the computation speed.

Advantages and Considerations

• Enables Larger Models: TP is essential when model layers themselves are too large for a single device's memory.
• Complements Other Strategies: It can be combined with DP and PP (Pipeline Parallelism) to form hybrid parallelism approaches (e.g., using TP within a pipeline stage, which is then replicated using DP).
• Potential for Overlap: Sophisticated implementations might overlap communication with computation to hide latency.

However, consider these points:

• Implementation Complexity: Requires careful code modifications or reliance on specialized libraries like NVIDIA's Megatron-LM or Microsoft's DeepSpeed, which provide optimized implementations of TP for Transformers.
• Communication Bottleneck: Performance is highly sensitive to the interconnect bandwidth and topology of the compute cluster. Scaling TP to a large number of devices can lead to diminishing returns if communication dominates.

Here's a PyTorch-style snippet illustrating the idea of splitting a weight matrix for column parallelism (this is highly simplified and omits communication and gradient handling):

import torch
import torch.nn as nn

# Assume world_size is the number of GPUs for TP
# Assume rank is the current GPU's rank in the TP group

class ColumnParallelLinear(nn.Module):
    def __init__(self, input_size, output_size, world_size, rank):
        super().__init__()
        self.input_size = input_size
        # Each GPU handles a slice of the output features
        self.output_size_per_partition = output_size // world_size
        self.world_size = world_size
        self.rank = rank

        # Initialize only the part of the weight matrix for this rank
        self.weight = nn.Parameter(
            torch.randn(self.output_size_per_partition, self.input_size)
        )
        # Bias is also partitioned
        self.bias = nn.Parameter(
            torch.randn(self.output_size_per_partition)
        )

    def forward(self, x):
        # Assume input x is available on all GPUs (broadcast or result of
        # all-reduce)
        # Matrix multiplication happens only on the local partition of the weight
        # Output_partition = X * A_partition^T + b_partition
        # Note: PyTorch linear layers expect weight as
        # (out_features, in_features)
        output_partition = nn.functional.linear(x, self.weight, self.bias)

        # In a real implementation, this output_partition would need to be
        # gathered across GPUs if the next layer isn't row-parallel.
        # For the Megatron-LM MLP pattern (Col->Row), no forward comms needed
        # here.

        # Backward pass communication (all-reduce for grad_X) is handled
        # by custom autograd functions in libraries like Megatron-LM.

        return output_partition # Returns only the slice computed by this GPU

In summary, Tensor Parallelism is a powerful technique for distributing the computation and memory of individual large layers across multiple devices. While it introduces significant communication overhead, it is often a necessary component, alongside Data and Pipeline Parallelism, for training state-of-the-art large language models that push the boundaries of single-accelerator capabilities. Frameworks like Megatron-LM and DeepSpeed abstract away much of the complexity, providing optimized building blocks for implementing TP effectively.