Tensor Parallelism: Slicing the Silicon

Training Large Language Models (LLMs) effectively requires moving beyond simple data parallelism. When a model's weight matrices are too large to fit into a single GPU's VRAM—or when intermediate activation memory explodes—we must slice the model itself.

Tensor Parallelism (TP) is an intra-layer model parallelism technique. Unlike Pipeline Parallelism (which slices the model vertically by layers) or Data Parallelism (which replicates the model), TP slices the model horizontally by splitting individual tensors (weight matrices) across multiple devices. This allows a cluster of GPUs to function as a single, massive accelerator, processing a single operation simultaneously.

This article explains the mechanics of TP with a forward-only PyTorch illustration and discusses its hardware trade-offs.

1. The Hardware Mechanics: Splitting the Matrix

At the register level, neural network training is dominated by Matrix Multiplication (MatMul): $Y = X \cdot W$. In Tensor Parallelism, we exploit the properties of linear algebra to distribute this specific operation across $N$ GPUs.

There are two primary ways to shard this computation, each requiring specific communication primitives to synchronize the results:

Strategy A: Column-Linear Parallelism

In this scheme, we split the weight matrix $W$ along its columns.

• Partitioning: If we have 2 GPUs, we split $W$ into $[W_1, W_2]$.
• Execution: We replicate the input $X$ across both GPUs. GPU 1 computes $Y_1 = X \cdot W_1$ and GPU 2 computes $Y_2 = X \cdot W_2$.
• Result: Each GPU holds a partial width of the output result (e.g., the first half of the output features).
• Communication: If a following operation needs the full output $Y$, an All-Gather can concatenate the partial results. A following row-parallel layer can consume the shards directly.

Strategy B: Row-Linear Parallelism

Here, we split the weight matrix $W$ along its rows.

• Partitioning: $W$ is split into two row blocks $W_1$ and $W_2$ (stacked vertically).
• Execution: To make the math work, the input $X$ must also be split along its last dimension (columns) into $[X_1, X_2]$. GPU 1 computes $Y_1 = X_1 \cdot W_1$.
• Result: Each GPU holds a partial sum of the final result. $Y = Y_1 + Y_2$.
• Communication: To get the valid final output, we must sum the results from all GPUs. This requires an All-Reduce operation.

The "Megatron-LM" Optimization

Efficiency in TP comes from combining these two strategies to minimize communication. In a standard Transformer MLP block (Linear $\rightarrow$ GeLU $\rightarrow$ Linear), we can arrange the splits to avoid synchronization in the middle.

1. Layer 1 (Column Parallel): The weight matrix is split by column. The output is split activations on each GPU.
2. Non-Linearity (GeLU): Since GeLU is an element-wise operation ($f(x)$), it can be applied independently to the partial outputs on each GPU. No communication is needed here.
3. Layer 2 (Row Parallel): The second weight matrix is split by row. It accepts the split output of Layer 1 directly as its split input.
4. Final Sync: Only after the second layer do we perform an All-Reduce to sum the partial results.

For this two-linear-layer MLP, the intermediate gather is avoided and only the row-parallel output reduction is needed in the forward path. Attention and backward communication add other collectives.

2. Bare Metal Implementation: Pure PyTorch

The following forward-only example shows a row-parallel linear layer with torch.distributed. All ranks must have an initialized process group and the same replicated input. It omits training: gradient-correct TP requires autograd-aware collective operations.

import torch
import torch.nn as nn
import torch.distributed as dist
 
class RowParallelLinear(nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        # 1. Setup World Info
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        
        # 2. Calculate Shard Size
        # We split the input dimension (rows of the weight matrix) across GPUs
        # Note: We assume input_size is divisible by world_size for simplicity
        assert input_size % self.world_size == 0
        self.input_shard_size = input_size // self.world_size
        
        # 3. Initialize Sharded Weights
        # Each GPU only holds a fraction of the total weights!
        # Shape: [output_size, input_shard_size]
        # Note: PyTorch Linear weights are typically (out_features, in_features)
        self.weight = nn.Parameter(torch.randn(output_size, self.input_shard_size))
        
        # Bias is usually handled by one rank or replicated and reduced.
        # Simplified here: no bias for clarity.
 
    @torch.no_grad()  # Forward-only illustration; training needs autograd-aware collectives.
    def forward(self, x):
        # x input shape: [batch_size, input_size]
        
        # 4. Scatter Input (Simulating the split)
        # In a real Transformer, x might already be sharded from the previous 
        # Column-Parallel layer. Here we manually shard it to simulate the state.
        # Split x along the last dimension (columns)
        input_shards = list(x.chunk(self.world_size, dim=-1))
        local_input = input_shards[self.rank].contiguous()
        
        # 5. Local MatMul
        # Compute the partial result for this GPU's slice of data
        # local_output shape: [batch_size, output_size]
        # We use .t() because PyTorch Linear weights are stored transposed
        local_output = torch.matmul(local_input, self.weight.t())
        
        # 6. All-Reduce (The Communication Bottleneck)
        # Sum the partial results from all GPUs into every GPU
        # The next dependent operation waits for the reduction result.
        dist.all_reduce(local_output, op=dist.ReduceOp.SUM)
        
        return local_output
 
# Usage pseudo-code:
# dist.init_process_group(backend='nccl')
# layer = RowParallelLinear(1024, 512).cuda()
# y = layer(x)

In this code, the dist.all_reduce line is the critical path. The next dependent operation waits for the reduction. The collective algorithm and overlap depend on the runtime, message size and topology.

3. Pros and Cons: A Hardware Perspective

Pros

1. Memory Reduction: For divisible sharded layers, TP can reduce their parameter, gradient and corresponding optimizer-state storage toward $1/N$ per rank. Replicated weights, activations and communication buffers keep total memory above that ideal fraction.
2. Access to Huge Models: It enables the training of models where a single layer's weights are simply too large to fit into one GPU's VRAM.
3. Concurrent layer work: TP ranks compute parts of the same layer concurrently. Whether this improves latency over pipeline parallelism depends on batch, model and communication.

Cons

1. Communication on the Critical Path: Sharded layers often need collectives within forward and backward passes. A dependent operation waits for the result; optimized schedules can overlap some communication with independent work.
2. Bandwidth Wall: Because of the high frequency of synchronization (often multiple times per Transformer block), TP is heavily bound by interconnect bandwidth.
   • Intra-Node (NVLink): TP often benefits from fast intra-node links. H100 NVLink 4 advertises 900 GB/s bidirectional per GPU (about 450 GB/s each direction before protocol and collective overhead).
   • Inter-Node (Ethernet/InfiniBand): Cross-node TP can be slower when collective traffic dominates; it is not categorically unusable. In a reported 16-H100 DeepSeek-V3 experiment, SiPipe's pipeline-based design achieves up to 4.5× the throughput of the paper's pure-TP baseline ¹. This result is workload-specific.
3. Implementation Complexity: Unlike FSDP which wraps standard PyTorch modules, TP requires shard-aware layers and collectives, whether implemented manually or through a framework.

Summary Table: When to use TP?

The 3D Parallelism Role

Large training systems often combine TP within a fast-connected GPU group with PP and DP, but the placement and degrees depend on the model and cluster.

Related

• Pipeline Parallelism: Pumping Data Upstream — the inter-layer half of the stack.
• Expert Parallelism: The Router's Dilemma — how MoE models shard the feed-forward instead.

References

• Megatron-LM tensor-parallelism paper
• PyTorch distributed API
• NVIDIA HGX H100 NVLink bandwidth