Tensor Model Parallelism

While Data Parallelism replicates the model across devices and processes different data batches, it doesn't help when the model itself is too large to fit into the memory of a single accelerator (GPU). For such scenarios, we need Model Parallelism, which involves partitioning a single model across multiple devices. Tensor Model Parallelism (TMP) is a specific type of model parallelism where individual layers or even specific operations within layers are split across devices.

This approach is particularly relevant for models with massive parameter counts concentrated in specific layers, such as the large embedding tables or the feed-forward network (FFN) layers found in modern Transformer architectures.

The Core Idea: Splitting Tensors

The fundamental principle of Tensor Model Parallelism is to split the weight tensors (and sometimes activation tensors) of a layer across multiple GPUs in a coordinated manner. The computation is then performed partially on each GPU, followed by communication steps to synchronize or combine the results, ultimately yielding the same output as the original, unsplit layer.

Let's consider a standard linear layer, defined by the operation $Y = XA + b$, where $X$ is the input activation, $A$ is the weight matrix, $b$ is the bias, and $Y$ is the output activation. TMP offers different strategies to parallelize this operation.

Column Parallelism

In column parallelism, the weight matrix $A$ is split column-wise across $N$ GPUs. Let $A = [A_1, A_2, ..., A_N]$, where each $A_i$ resides on a different GPU.

1. Input Broadcast: The input $X$ is broadcasted (or available) to all GPUs in the model parallel group.
2. Parallel Computation: Each GPU $i$ computes a partial result $Y_i = X A_i$.
3. Gather/Concatenate: The partial results $Y_i$ are gathered and concatenated along the column dimension to form the final output $Y = [Y_1, Y_2, ..., Y_N]$. The bias $b$ can also be split column-wise, $b = [b_1, b_2, ..., b_N]$, and added independently on each GPU before the gather step.

Column parallelism applied to a linear layer. Input $X$ is shared, weights $A$ and bias $b$ are split column-wise, and partial outputs are concatenated.

This approach requires communication (a gather or all-gather operation) after the parallel computation to assemble the full output tensor $Y$.

Row Parallelism

Alternatively, we can split the weight matrix $A$ row-wise: $A = \begin{bmatrix} A_1 \\ A_2 \\ \vdots \\ A_N \end{bmatrix}$.

1. Input Split/Scatter: The input $X$ is split along its last dimension (feature dimension) across the GPUs: $X = [X_1, X_2, ..., X_N]$. Note this is different from column parallelism where the full $X$ was needed on each GPU.
2. Parallel Computation: Each GPU $i$ computes a partial result using its portion of the input and weights: $X_i A_i$.
3. All-Reduce: The partial results are summed across all GPUs using an all-reduce operation to produce the final output $Y = \sum_i X_i A_i$. The bias $b$ is typically added only on one GPU (e.g., rank 0) after the all-reduce sum.

Row parallelism applied to a linear layer. Input $X$ is split, weights $A$ are split row-wise, partial results are computed, and then summed via all-reduce. Bias $b$ is added after the reduction.

Row parallelism requires communication (an all-reduce operation) after the matrix multiplication. An important benefit is that the output activation $Y$ is replicated across all participating GPUs, which might be the required input format for a subsequent layer (e.g., a layer norm or another row-parallel linear layer).

Combining Column and Row Parallelism in Transformers

Transformer models often utilize a combination of these techniques. For instance, within a standard Feed-Forward Network (FFN) block consisting of two linear layers:

1. The first linear layer might use column parallelism. This splits the larger intermediate dimension across GPUs. The output activations are distributed (sharded).
2. The second linear layer might use row parallelism. It takes the sharded activations as input (no communication needed if the split dimension matches) and produces an output that is summed via all-reduce, making the final output of the FFN block replicated across GPUs, ready for the next part of the Transformer (like residual connections or layer normalization).

This strategic combination helps minimize communication overhead by keeping activations sharded between the two linear layers within the FFN block.

Parallelizing Embeddings

For models with extremely large vocabulary sizes, the embedding table can become a significant memory bottleneck. Tensor Model Parallelism can be applied here by splitting the embedding table row-wise (along the vocabulary dimension) across the GPUs. When looking up embeddings for input token IDs:

1. Each GPU receives the full sequence of input IDs.
2. Each GPU looks up embeddings only for the portion of the vocabulary it holds. If an ID corresponds to a row held by another GPU, it produces a zero vector.
3. An all-reduce operation sums the partial embedding vectors across GPUs. Since only one GPU produced the actual embedding vector for any given token ID (others produced zeros), the sum effectively gathers the correct embeddings.

Implementation Details and Communication

Implementing Tensor Model Parallelism manually requires careful handling of tensor sharding, computation, and synchronization using primitives from the torch.distributed package:

• torch.distributed.broadcast: Sends a tensor from one rank to all others.
• torch.distributed.scatter: Scatters chunks of a tensor across ranks.
• torch.distributed.gather: Gathers tensors from all ranks onto one rank.
• torch.distributed.all_gather: Gathers tensors from all ranks onto all ranks.
• torch.distributed.reduce_scatter: Performs an operation (like sum) across ranks and scatters the result.
• torch.distributed.all_reduce: Performs an operation (like sum) across ranks and makes the result available on all ranks.

Specialized functions or wrappers are often created to encapsulate these operations for specific layer types (e.g., ColumnParallelLinear, RowParallelLinear). Libraries like NVIDIA's Megatron-LM pioneered many of these techniques, and parts of this functionality are being integrated into PyTorch core via modules like torch.distributed.tensor.parallel which provide higher-level APIs to simplify these implementations.

Consider an example for a column-parallel linear layer using helper functions:

import torch
import torch.nn as nn
import torch.distributed as dist

# Assume these helper functions exist for managing parallel groups and communication
from .parallel_utils import (
    get_tensor_model_parallel_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
    copy_to_tensor_model_parallel_region, # Handles input broadcasting/splitting
    gather_from_tensor_model_parallel_region # Handles output gathering/reduction
)

class ColumnParallelLinear(nn.Module):
    def __init__(self, input_size, output_size, bias=True, **kwargs):
        super().__init__()
        world_size = get_tensor_model_parallel_world_size()
        # Ensure output_size is divisible by world_size
        assert output_size % world_size == 0
        self.output_size_per_partition = output_size // world_size
        self.input_size = input_size

        # Weight matrix is split along the output dimension (columns)
        self.weight = nn.Parameter(torch.empty(
            self.output_size_per_partition, self.input_size, **kwargs
        ))
        # Initialize weights... (e.g., using init.kaiming_uniform_)

        if bias:
            # Bias is also split along the output dimension
            self.bias = nn.Parameter(torch.empty(
                self.output_size_per_partition, **kwargs
            ))
            # Initialize bias... (e.g., using init.zeros_)
        else:
            self.register_parameter('bias', None)

    def forward(self, input_):
        # Input might need to be broadcast or is already available if
        # the previous layer's output was replicated (e.g., LayerNorm).
        # This function handles the necessary communication.
        parallel_input = copy_to_tensor_model_parallel_region(input_)

        # Perform local matrix multiplication
        output_parallel = nn.functional.linear(parallel_input, self.weight, self.bias)

        # Gather results from all GPUs in the tensor parallel group
        # Concatenates along the column dimension.
        output_ = gather_from_tensor_model_parallel_region(output_parallel)

        return output_

Implementation sketch of a Column Parallel Linear layer. Note the explicit sharding of weights/bias and the use of communication wrappers.

Trade-offs and Considerations

• Complexity: Implementing TMP requires significant modifications to the model architecture code, unlike DDP which often acts as a wrapper. Debugging parallel logic can be challenging.
• Communication Overhead: TMP introduces communication within the forward and backward passes of individual layers. Operations like all-gather or all-reduce can become bottlenecks, especially if the interconnect bandwidth between GPUs (e.g., NVLink vs PCIe) is limited.
• Granularity: It's most effective for layers with a large number of parameters or large activations relative to the computation performed.
• Composability: TMP is often combined with Data Parallelism (DDP) and Pipeline Parallelism to train truly massive models, forming hybrid parallelism strategies.

In summary, Tensor Model Parallelism is an indispensable technique when individual model layers exceed single-device memory limits. By splitting weights and computation within layers across multiple devices, it enables the training of models far larger than would otherwise be possible, albeit at the cost of increased implementation complexity and communication overhead. It's a fundamental building block for scaling today's largest neural networks.