Tensor Parallelism › Code Implementation
Language AI HandbookAll chapters Tensor Parallelism › Code ImplementationOne section · Sections in this chapter9 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
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.
Setting Up Distributed Process Groups
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.
def setup_distributed(rank: int, world_size: int) -> None:
"""Initialize the distributed process group."""
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"
dist.init_process_group(
backend="gloo", # Use "nccl" for actual GPUs
rank=rank,
world_size=world_size,
)
def cleanup_distributed() -> None:
"""Clean up the distributed process group."""
dist.destroy_process_group()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.
Column-Parallel Linear Layer
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.
import torch
import torch.nn as nn
import torch.nn.functional as F
class ColumnParallelLinear(nn.Module):
"""
Linear layer split along the output (column) dimension.
Each rank holds a shard of shape (in_features, out_features // world_size).
The input X is replicated on all ranks.
The output Y is sharded: each rank produces columns [rank * shard_size : (rank+1) * shard_size].
"""
def __init__(
self,
in_features: int,
out_features: int,
world_size: int,
rank: int,
bias: bool = True,
) -> None:
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.world_size = world_size
self.rank = rank
# Each rank holds out_features / world_size output columns
assert out_features % world_size == 0, (
f"out_features ({out_features}) must be divisible by world_size ({world_size})"
)
self.shard_size = out_features // world_size
# Weight shard for this rank: shape (shard_size, in_features)
# F.linear expects weight of shape (out_features, in_features)
self.weight = nn.Parameter(
torch.randn(self.shard_size, in_features) * (1.0 / in_features**0.5)
)
if bias:
# Bias shard: each rank has the bias for its output columns
self.bias = nn.Parameter(torch.zeros(self.shard_size))
else:
self.bias = None
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: Full input tensor of shape (batch, seq_len, in_features),
replicated on all ranks.
Returns:
Sharded output of shape (batch, seq_len, shard_size),
each rank has a different column slice.
"""
# Local matrix multiply: X @ W_shard^T + b_shard
# F.linear computes x @ weight.T + bias, weight shape: (shard_size, in_features)
output = F.linear(x, self.weight, self.bias)
return outputNotice 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).
Row-Parallel Linear Layer
The row-parallel layer receives a sharded input (output from a column-parallel layer) and produces a full output via all-reduce.
class RowParallelLinear(nn.Module):
"""
Linear layer split along the input (row) dimension.
Each rank holds a shard of shape (in_features // world_size, out_features).
The input X is sharded: each rank receives its corresponding columns.
The full output Y is produced by summing partial results via all-reduce.
"""
def __init__(
self,
in_features: int,
out_features: int,
world_size: int,
rank: int,
bias: bool = True,
) -> None:
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.world_size = world_size
self.rank = rank
assert in_features % world_size == 0, (
f"in_features ({in_features}) must be divisible by world_size ({world_size})"
)
self.shard_size = in_features // world_size
# Weight shard for this rank: shape (out_features, shard_size)
# F.linear expects weight of shape (out_features, in_features)
self.weight = nn.Parameter(
torch.randn(out_features, self.shard_size)
* (1.0 / in_features**0.5)
)
# Bias lives on rank 0 only and is added after the all-reduce
if bias and rank == 0:
self.bias = nn.Parameter(torch.zeros(out_features))
else:
self.bias = None
def forward(self, x_shard: torch.Tensor) -> torch.Tensor:
"""
Args:
x_shard: Sharded input of shape (batch, seq_len, shard_size).
Each rank provides its input shard (columns from the previous layer).
Returns:
Full output of shape (batch, seq_len, out_features) on all ranks.
"""
# Local matrix multiply: partial_Y = X_shard @ W_shard^T
# Shape: (batch, seq_len, out_features)
partial_output = F.linear(x_shard, self.weight)
# All-reduce: sum partial outputs across all ranks
# After this, every rank has the full output
dist.all_reduce(partial_output, op=dist.ReduceOp.SUM)
# Add bias (only rank 0 has it; others contribute zero)
if self.bias is not None:
partial_output = partial_output + self.bias
return partial_outputThe 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.
Tensor-Parallel Feed-Forward Network
We can now assemble a complete tensor-parallel FFN by chaining column-parallel and row-parallel layers.
class TensorParallelFFN(nn.Module):
"""
Tensor-parallel feed-forward network using column-row chaining.
W1 is column-parallel: each rank holds (d_model, 4*d_model/world_size)
W2 is row-parallel: each rank holds (4*d_model/world_size, d_model)
Communication: one all-reduce in the forward pass (inside RowParallelLinear)
"""
def __init__(
self,
d_model: int,
d_ffn: int,
world_size: int,
rank: int,
) -> None:
super().__init__()
self.column_parallel = ColumnParallelLinear(
in_features=d_model,
out_features=d_ffn,
world_size=world_size,
rank=rank,
)
self.row_parallel = RowParallelLinear(
in_features=d_ffn,
out_features=d_model,
world_size=world_size,
rank=rank,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: Full input tensor (batch, seq_len, d_model) on all ranks.
Returns:
Full output tensor (batch, seq_len, d_model) on all ranks.
"""
# Column-parallel: full input → sharded activations
# Each rank computes GeLU(X @ W1_shard)
hidden = F.gelu(self.column_parallel(x))
# Row-parallel: sharded activations → full output (with all-reduce)
output = self.row_parallel(hidden)
return outputThe 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.
Demonstrating Correctness
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.
def demonstrate_column_parallel_correctness() -> dict:
"""
Verify that column-parallel + row-parallel chaining gives
the same result as a standard sequential linear layer.
This function simulates 2-GPU behavior without actual distributed setup.
"""
torch.manual_seed(42)
d_model = 8
d_ffn = 16
batch_size = 2
seq_len = 3
world_size = 2
# Create full weight matrices (the reference)
W1_full = torch.randn(d_model, d_ffn) * (1.0 / d_model**0.5)
W2_full = torch.randn(d_ffn, d_model) * (1.0 / d_ffn**0.5)
b1_full = torch.zeros(d_ffn)
b2_full = torch.zeros(d_model)
# Full input, replicated on both "GPUs"
x = torch.randn(batch_size, seq_len, d_model)
# Reference (non-parallel) computation
hidden_full = F.gelu(x @ W1_full + b1_full)
y_reference = hidden_full @ W2_full + b2_full
# Simulate rank 0
shard_size_ffn = d_ffn // world_size # = 8
W1_shard_0 = W1_full[:, :shard_size_ffn] # (8, 8)
W2_shard_0 = W2_full[:shard_size_ffn, :] # (8, 8)
b1_shard_0 = b1_full[:shard_size_ffn]
hidden_0 = F.gelu(x @ W1_shard_0 + b1_shard_0) # (2, 3, 8)
partial_y_0 = hidden_0 @ W2_shard_0 # (2, 3, 8)
# Simulate rank 1
W1_shard_1 = W1_full[:, shard_size_ffn:] # (8, 8)
W2_shard_1 = W2_full[shard_size_ffn:, :] # (8, 8)
b1_shard_1 = b1_full[shard_size_ffn:]
hidden_1 = F.gelu(x @ W1_shard_1 + b1_shard_1) # (2, 3, 8)
partial_y_1 = hidden_1 @ W2_shard_1 # (2, 3, 8)
# All-reduce: sum partial results + bias
y_parallel = partial_y_0 + partial_y_1 + b2_full
max_diff = (y_reference - y_parallel).abs().max().item()
mean_diff = (y_reference - y_parallel).abs().mean().item()
return {
"max_absolute_diff": max_diff,
"mean_absolute_diff": mean_diff,
"outputs_match": max_diff < 1e-5,
"y_reference_shape": list(y_reference.shape),
"y_parallel_shape": list(y_parallel.shape),
"d_model": d_model,
"d_ffn": d_ffn,
"world_size": world_size,
"shard_size": shard_size_ffn,
}
results = demonstrate_column_parallel_correctness()Tensor Parallelism Correctness Verification ================================================== Input shape: (batch=2, seq=3, d_model=8) Output shape: [2, 3, 8] FFN hidden dim: 16 split into 2 shards of 8 each Max absolute difference (reference vs parallel): 2.38e-07 Mean absolute difference: 3.34e-08 Outputs match (diff < 1e-5): True
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.
Memory Savings Analysis
A key benefit of tensor parallelism is reduced memory per GPU. Let's compute how memory requirements scale with the tensor-parallel degree.
def compute_memory_savings(
d_model: int,
num_layers: int,
ffn_ratio: int = 4,
bytes_per_param: int = 2, # fp16
) -> dict:
"""
Compute weight memory per GPU for different tensor-parallel degrees.
"""
# Parameter counts for one transformer layer
# Attention: 4 weight matrices of shape (d_model, d_model) for Q, K, V, O
attn_params_per_layer = 4 * d_model * d_model
# FFN: W1 of (d_model, ffn_ratio*d_model), W2 of (ffn_ratio*d_model, d_model)
ffn_params_per_layer = 2 * d_model * (ffn_ratio * d_model)
total_params_per_layer = attn_params_per_layer + ffn_params_per_layer
total_params = total_params_per_layer * num_layers
total_bytes = total_params * bytes_per_param
total_gb = total_bytes / (1024**3)
memory_by_degree = {}
for degree in [1, 2, 4, 8, 16]:
params_per_gpu = total_params / degree
bytes_per_gpu = params_per_gpu * bytes_per_param
gb_per_gpu = bytes_per_gpu / (1024**3)
memory_by_degree[degree] = {
"params_per_gpu": int(params_per_gpu),
"gb_per_gpu": gb_per_gpu,
"reduction_factor": degree,
}
return {
"total_params": total_params,
"total_gb": total_gb,
"memory_by_degree": memory_by_degree,
"d_model": d_model,
"num_layers": num_layers,
}
# GPT-3 scale: d_model=12288, 96 layers
memory_info = compute_memory_savings(d_model=12288, num_layers=96)Model: d_model=12288, 96 layers
Total weight parameters: 173.9B
Total weight memory (fp16): 324.0 GB
TP Degree | Params/GPU | Weight GB/GPU
---------------------------------------------
1 | 173.9B | 324.0
2 | 87.0B | 162.0
4 | 43.5B | 81.0
8 | 21.7B | 40.5
16 | 10.9B | 20.2These numbers show why tensor parallelism is essential at GPT-3 scale. With a tensor-parallel degree of 8, each GPU holds roughly 40 GB of weight memory for GPT-3's attention and FFN layers combined, well within the 80 GB capacity of an A100. Without tensor parallelism, a single GPU would need to hold over 320 GB for those weights alone, far exceeding any available device memory. This estimate covers only the core weight matrices; embedding tables, optimizer states, and gradient buffers add substantially more to the total memory burden. The optimizer states for mixed-precision training (maintaining fp32 master weights alongside fp16 working weights) roughly triple the weight memory, which is why real training systems need tensor-parallel degrees of 8 or more even for models much smaller than GPT-3.
Communication Overhead Estimation
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.
def estimate_communication_overhead(
d_model: int,
seq_len: int,
batch_size: int,
num_layers: int,
num_gpus: int,
bandwidth_gb_per_s: float,
compute_tflops_per_gpu: float,
) -> dict:
"""
Estimate communication time relative to compute time for tensor parallelism.
Each transformer layer requires 2 all-reduces (one for attention, one for FFN).
Each all-reduce transmits one full activation tensor of shape (batch, seq, d_model).
"""
bytes_per_element = 2 # fp16
# Size of one activation tensor
activation_size_bytes = batch_size * seq_len * d_model * bytes_per_element
activation_size_gb = activation_size_bytes / (1024**3)
# In a ring all-reduce with N GPUs, each GPU sends and receives 2*(N-1)/N * data
# For communication time, use the effective bandwidth of the ring
all_reduce_volume_gb = 2 * (num_gpus - 1) / num_gpus * activation_size_gb
all_reduce_time_s = all_reduce_volume_gb / bandwidth_gb_per_s
# Total communication: 2 all-reduces per layer
total_comm_time_s = 2 * num_layers * all_reduce_time_s
# Compute time estimate for transformer
# Each transformer layer: ~24 * d_model^2 * batch * seq FLOPs
# (4 proj matrices in attention: 2 * 4 * d^2 * batch * seq = 8 * d^2 * batch * seq
# + FFN 2 linear layers: 2 * 2 * 4 * d^2 * batch * seq = 16 * d^2 * batch * seq)
flops_per_layer = 24 * (d_model**2) * batch_size * seq_len
total_flops = flops_per_layer * num_layers
# With TP degree = num_gpus, compute is distributed across GPUs
flops_per_gpu = total_flops / num_gpus
# A100 achieves ~60% of peak on typical transformer workloads
effective_tflops = compute_tflops_per_gpu * 0.6
compute_time_s = flops_per_gpu / (effective_tflops * 1e12)
comm_to_compute_ratio = total_comm_time_s / compute_time_s
return {
"activation_size_mb": activation_size_gb * 1024,
"all_reduce_volume_mb": all_reduce_volume_gb * 1024,
"all_reduce_time_ms": all_reduce_time_s * 1000,
"total_comm_time_ms": total_comm_time_s * 1000,
"compute_time_ms": compute_time_s * 1000,
"comm_to_compute_ratio": comm_to_compute_ratio,
"comm_overhead_pct": comm_to_compute_ratio * 100,
}
# NVLink (within node): ~600 GB/s bidirectional
nvlink_overhead = estimate_communication_overhead(
d_model=4096,
seq_len=2048,
batch_size=4,
num_layers=32,
num_gpus=8,
bandwidth_gb_per_s=300.0, # One-directional NVLink
compute_tflops_per_gpu=312.0, # A100 fp16
)
# InfiniBand (across nodes): ~25 GB/s
infiniband_overhead = estimate_communication_overhead(
d_model=4096,
seq_len=2048,
batch_size=4,
num_layers=32,
num_gpus=8,
bandwidth_gb_per_s=12.5, # One-directional InfiniBand
compute_tflops_per_gpu=312.0,
)Communication Overhead Analysis (8 GPUs, d_model=4096, 32 layers)
=================================================================
NVLink InfiniBand
-----------------------------------------------------------------
Activation size (MB) 64.0 64.0
All-reduce volume/op (MB) 112.0 112.0
All-reduce time (ms) 0.365 8.750
Total comm time (ms) 23.3 560.0
Compute time (ms) 70.5 70.5
Comm overhead (%) 33.1 794.5The contrast between NVLink and InfiniBand overhead illustrates why tensor parallelism is typically restricted to within a single DGX node (8 GPUs). NVLink's high bandwidth keeps communication below compute time through the typical within-node range, while InfiniBand communication already exceeds compute time at degree 2 in this simplified model. The analysis also shows why the minimal-communication design of Megatron's column-row chaining matters: at InfiniBand speeds, even the carefully optimized two-all-reduces-per-layer design is expensive. A naive design requiring four or six all-reduces per layer would be completely impractical.
Key Parameters
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 GPUs forms one complete model replica.
- data_parallel_degree (dp): Number of replicas used for data parallelism. Total GPUs .
- sequence_length: Affects activation memory proportionally. Longer sequences amplify the benefit of sequence parallelism.
- all_reduce_backend: Use
ncclfor NVIDIA GPU communication (exploits NVLink/InfiniBand),gloofor CPU.