Tensor, Pipeline, and Expert Parallelism Compared

Scaling large language models (LLMs) and Mixture-of-Experts (MoE) architectures requires breaking the model or data across multiple accelerators. The choice of parallelism strategy is not merely a configuration option; it dictates the memory footprint, communication overhead, and computational efficiency of the system. Tensor, pipeline, and expert parallelism address three distinct bottlenecks: intra-operator memory limits, inter-operator latency, and sparse computation utilization.

Tensor Parallelism: Sharding Within Operators

Tensor parallelism (TP) splits individual operations across devices. It is primarily used when a single layer’s weights exceed the memory capacity of a single GPU or when parallelizing the matrix multiplications within a layer yields significant speedup.

Consider a standard Transformer attention head with a query matrix $Q$ of shape $(seq\_len, head\_dim)$. In tensor parallelism with degree $N$, this matrix is split column-wise. Each device holds a slice of the weight matrix. During the forward pass, each device computes its partial result. The results must then be concatenated to form the final output. This concatenation requires an All-Gather collective operation.

The mechanism relies on the assumption that the computational graph within a layer is dense. Every device participates in every operation within the layer. This creates a tight coupling between devices. The communication cost is high because an All-Reduce is performed after every major tensor operation (e.g., linear layers, attention outputs).

# Conceptual TP Linear Layer
# W is split along dim=1 (output features)
class TensorParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, rank, world_size, weight):
        super().__init__()
        # Split weights across ranks
        self.weight = nn.Parameter(torch.split(weight, out_features // world_size, dim=1)[rank])
        
    def forward(self, x):
        # Local matmul
        out = torch.matmul(x, self.weight)
        # Concatenate output slices across all GPUs
        out = all_gather(out) 
        return out

TP is most effective for models where the layer size is large but the number of layers is moderate. It minimizes memory usage per device but maximizes network traffic. If the network bandwidth is insufficient, TP becomes the bottleneck.

Pipeline Parallelism: Sharding Across Layers

Pipeline parallelism (PP) divides the model sequentially by layer. Device 0 processes layers 1–K, Device 1 processes layers K+1–2K, and so on. This approach addresses the limitation that TP cannot scale beyond the number of layers in the model.

The primary challenge in PP is pipeline bubbles—periods where a device is idle because it is waiting for data from the previous stage or sending data to the next. In a naive implementation (GPipe), the first stage computes layer 1, passes it to stage 2, and waits. Stage 2 computes layer 2, passes it to stage 3, etc. While stage 1 is idle, stage 2 is computing.

To mitigate this, PP uses micro-batching. The global batch size is divided into smaller micro-batches that flow through the pipeline like a production line. Advanced scheduling techniques, such as 1F1B (One Forward, One Backward), ensure that once the pipeline is full, every stage is computing during every time step, minimizing idle time.

Time Step | Stage 1 (Layer 1) | Stage 2 (Layer 2) | Stage 3 (Layer 3)
----------|-------------------|-------------------|------------------
T1        | Micro-batch 1 (F) | Idle              | Idle
T2        | Micro-batch 2 (F) | Micro-batch 1 (F) | Idle
T3        | Micro-batch 3 (F) | Micro-batch 2 (F) | Micro-batch 1 (F)
T4        | Micro-batch 1 (B) | Micro-batch 2 (B) | Micro-batch 1 (B) ...

PP increases memory efficiency by distributing activations across devices, but it introduces significant latency due to the sequential nature of layer execution. It also requires caching activations for backward passes, which can be memory-intensive if the pipeline depth is shallow.

Expert Parallelism: Sharding Sparse Experts

Expert parallelism (EP) is specific to Mixture-of-Experts (MoE) models. In an MoE model, each token is routed to a small subset of "experts" (feed-forward networks) rather than passing through all parameters. The total parameter count is large, but the active parameters per token are small.

EP shards the experts themselves across devices. If there are 64 experts and 8 GPUs, each GPU holds 8 experts. When a token is processed, the router determines which experts should handle it. The token’s activation must then be sent to the device holding the relevant expert. This requires an All-to-All communication pattern, where each device sends different parts of its data to different other devices.

# Conceptual MoE Router with EP
def moe_forward(tokens, experts, router_weights):
    # 1. Determine expert assignment for each token
    # Returns: token_expert_map (token_id -> expert_id)
    
    # 2. Gather tokens to their respective expert devices
    # This is an All-to-All collectives operation
    gathered_tokens = all_to_all(tokens, token_expert_map)
    
    # 3. Each device processes only its local experts
    local_outputs = [expert(tok) for tok, expert in zip(gathered_tokens, local_experts)]
    
    # 4. Scatter outputs back to original token order
    final_output = all_to_all(local_outputs, reverse_map)
    return final_output

EP is highly efficient for inference and training because it allows the model to scale parameter count without increasing active computation. However, it introduces two critical challenges: load imbalance and communication overhead. If the router assigns too many tokens to a few experts, some devices become bottlenecks. Additionally, the All-to-All communication can be expensive if the token count per expert is small, leading to high overhead relative to computation.

Synthesis and Trade-offs

Choosing between these strategies depends on the model architecture and hardware constraints.

1. Memory vs. Communication: TP reduces memory usage but increases communication frequency (All-Reduce). PP reduces memory usage but increases latency (pipeline bubbles). EP reduces active compute but increases communication complexity (All-to-All).
2. Load Balance: TP and PP generally have balanced workloads. EP can suffer from load imbalance if the routing algorithm is not optimized.
3. Implementation Complexity: TP is straightforward to implement using standard distributed primitives. PP requires complex scheduling and activation management. EP requires custom routing logic and efficient All-to-All implementations.

In practice, modern large-scale systems often combine these strategies. For example, Megatron-LM uses TP within layers and PP across layers. Recent MoE implementations like Megatron-DeepSpeed integrate EP with TP and PP to maximize throughput.

The choice is not mutually exclusive but compositional. Tensor parallelism handles the width of the model, pipeline parallelism handles the depth, and expert parallelism handles the sparsity. Engineers must balance the communication bandwidth, memory capacity, and computational density of their cluster to select the optimal configuration.

Common Pitfalls

When implementing these parallelism strategies, engineers frequently encounter specific pitfalls that degrade performance or stability:

1. Network Bandwidth Saturation in TP: Tensor parallelism requires frequent All-Reduce operations. If the interconnect bandwidth (e.g., NVLink or InfiniBand) is insufficient for the tensor sizes involved, the GPU compute units will stall waiting for data transfer, leading to near-zero utilization. Always profile communication overhead against compute time.
2. Pipeline Bubble Inefficiency in PP: Even with 1F1B scheduling, deep pipelines can suffer from residual bubbles, especially during the warm-up and cool-down phases. Furthermore, if micro-batch sizes are not tuned correctly, the overhead of managing many small batches can outweigh the benefits of overlapping compute and communication.
3. Load Imbalance in EP Routing: Expert parallelism relies heavily on the router's ability to distribute tokens evenly. If the routing policy is naive or if certain experts become "hotspots" due to data skew, some GPUs will be overloaded while others remain idle. This leads to significant synchronization delays as the system waits for the slowest device to finish.

Practical Takeaways

Use these mental models to guide your architecture decisions:

1. Start with Tensor Parallelism for Width: If your model fits in memory per layer but you need faster iteration, use TP. It is the easiest to implement and provides immediate memory relief for large layers.
2. Use Pipeline Parallelism for Depth: When the model is too large for a single node's memory, even with TP, introduce PP. Accept the latency penalty in exchange for the ability to scale beyond single-node limits.
3. Reserve Expert Parallelism for Sparse Models: Only use EP if you are explicitly training or serving a MoE architecture. The complexity of All-to-All communication and routing is not justified for dense models.

FAQ

Can I use Tensor Parallelism and Pipeline Parallelism together? Yes, this is the most common configuration for large models. TP handles sharding within each layer, while PP handles sharding across layers. Frameworks like Megatron-LM support this hybrid approach seamlessly.

When should I use Expert Parallelism? Use Expert Parallelism only when working with Mixture-of-Experts (MoE) models. It is not applicable to standard dense transformer architectures. It is particularly useful when you want to scale the total parameter count (model size) without proportionally increasing the active computation cost per token.

Does Expert Parallelism always improve performance? No. EP introduces significant communication overhead via All-to-All collectives. If the number of tokens per expert is low, the communication cost may exceed the computation savings, leading to slower performance compared to dense models or simpler parallelism strategies.

Conclusion

Tensor, pipeline, and expert parallelism are not competing alternatives but complementary tools for scaling distributed systems. Tensor parallelism optimizes intra-layer memory, pipeline parallelism optimizes inter-layer latency, and expert parallelism optimizes sparse compute utilization. Effective system design requires selecting the right combination based on network topology, memory constraints, and model architecture to achieve maximum throughput.