Multi-GPU Parallelism (Tensor, Pipeline, Data)

Computation Graphs & ML Infrastructure advanced How AI Works 13 min read

Overview

Modern LLMs are too large for a single GPU. A 70B parameter model at FP16 requires 140 GB of VRAM — nearly double what even the highest-end single GPU (H200 at 141 GB) can hold when you account for KV cache and overhead. Training these models is even more demanding, requiring memory for optimizer states, gradients, and activations that can multiply the base memory requirement by 4-16x. Multi-GPU parallelism is the set of strategies for distributing this work across multiple GPUs.

There are three fundamental approaches to parallelism, each splitting the work along a different axis: data parallelism replicates the model and splits the dataset, tensor parallelism splits individual layers across GPUs, and pipeline parallelism assigns different groups of layers to different GPUs. In practice, production systems combine all three in what’s known as 3D parallelism to maximize throughput and minimize idle time.

Understanding when and how to apply each strategy is critical for AI engineers working on large-scale training and high-throughput inference. The wrong parallelism configuration can waste GPU resources, create communication bottlenecks, or fail to meet latency requirements.

How It Works

Data Parallelism (DP)

Data parallelism is the simplest and most widely used form of parallelism. The core idea: replicate the entire model on every GPU, split the training batch into mini-batches (one per GPU), and synchronize gradients after each backward pass.

Data Parallelism

Each GPU holds a complete copy of the model and processes a different slice of the data batch. After computing gradients locally, an all-reduce operation averages gradients across all GPUs, ensuring model replicas stay synchronized. Effective batch size = per-GPU batch size × number of GPUs.

The training loop with data parallelism works as follows:

  1. Scatter different data mini-batches to each GPU
  2. Each GPU performs forward and backward pass independently
  3. All-reduce operation synchronizes gradients across GPUs
  4. Each GPU updates its model copy with the averaged gradients
# PyTorch DistributedDataParallel (DDP) example
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# Initialize process group
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# Wrap model in DDP
model = MyModel().to(local_rank)
model = DDP(model, device_ids=[local_rank])

# Training loop — each GPU gets different data
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler)

for batch in dataloader:
    loss = model(batch)
    loss.backward()       # Gradients computed locally
    optimizer.step()      # DDP handles all-reduce automatically
    optimizer.zero_grad()

Strengths: Simple to implement, near-linear throughput scaling, minimal code changes required. Limitation: Every GPU must hold a full model copy, so it doesn’t help when a single model exceeds one GPU’s VRAM. Communication cost is proportional to model size (all-reduce sends gradients equal to the parameter count).

For inference, data parallelism means running independent model replicas that each handle separate requests. This scales throughput by approximately 50% per additional GPU (accounting for overhead) with no inter-GPU communication during the forward pass — making it the simplest optimization for scaling inference volume.

Tensor Parallelism (TP)

Tensor parallelism splits individual layer computations across GPUs. Rather than each GPU running a complete layer, each GPU computes a portion of every layer and exchanges intermediate results. This is the finest-grained form of parallelism and the most communication-intensive.

Tensor Parallelism

Individual weight matrices within a layer are partitioned across GPUs. For a linear layer Y = XW, the weight matrix W is split column-wise or row-wise across GPUs. Each GPU computes its partition, then results are combined via all-reduce or all-gather. This reduces per-GPU memory proportionally to the TP degree.

In a transformer, tensor parallelism typically splits the attention heads and the feed-forward network (FFN) across GPUs:

# Conceptual tensor parallelism for a linear layer
# With TP=2, weight matrix W is split into [W1, W2]

# GPU 0 computes:                 GPU 1 computes:
# Y0 = X @ W1                    Y1 = X @ W2

# Column-parallel: concatenate   [Y0, Y1] → Y
# Row-parallel: all-reduce        Y0 + Y1 → Y

# In vLLM, tensor parallelism is configured at startup:
# python -m vllm.entrypoints.openai.api_server \
#   --model meta-llama/Llama-3.1-70B \
#   --tensor-parallel-size 4

Performance impact: Tensor parallelism delivers the best latency improvement — a 3x improvement in time-to-first-token (TTFT) and consistent gains in per-token latency. However, it requires high-bandwidth interconnects (NVLink at 900 GB/s, not PCIe at 64 GB/s) because GPUs exchange data after every layer. Tensor parallelism should be confined to GPUs within a single node connected by NVLink.

Pipeline Parallelism (PP)

Pipeline parallelism assigns contiguous groups of layers to different GPUs, creating a sequential pipeline. Data flows from the first GPU (holding the early layers) through to the last GPU (holding the final layers and the output head).

Pipeline Parallelism

The model’s layers are divided into P stages across P GPUs. GPU 0 processes layers 1-N/P, GPU 1 processes layers N/P+1 to 2N/P, and so on. Data flows through the pipeline, with each GPU’s output becoming the next GPU’s input. Micro-batching reduces idle time (the “bubble”) by overlapping computation across stages.

The naive approach has a significant problem: while GPU 0 processes the first micro-batch’s early layers, GPUs 1-3 sit idle. This idle time is called the pipeline bubble. The solution is micro-batching — splitting each batch into smaller micro-batches that flow through the pipeline in an overlapping fashion:

# Pipeline parallelism execution timeline (4 GPUs, 4 micro-batches)
# F = Forward pass, B = Backward pass, number = micro-batch

Time →  ──────────────────────────────────────────────
GPU 0:  |F1|F2|F3|F4|  |B4|B3|B2|B1|
GPU 1:     |F1|F2|F3|F4|  |B4|B3|B2|B1|
GPU 2:        |F1|F2|F3|F4|  |B4|B3|B2|B1|
GPU 3:           |F1|F2|F3|F4|  |B4|B3|B2|B1|

# Bubble ratio ≈ (P - 1) / (M + P - 1)
# where P = pipeline stages, M = micro-batches
# More micro-batches → smaller bubble → better utilization

Strengths: Only requires point-to-point communication between adjacent GPUs (low bandwidth requirement), works across nodes connected by standard networking. Enables serving models that don’t fit on a single GPU. Limitation: Pipeline bubble reduces utilization; more effective with many micro-batches. Adds latency for individual requests since data must traverse all stages.

For inference, pipeline parallelism cuts TTFT P99 latency by 2.5-3x at high concurrency by distributing the KV cache across GPUs, allowing a larger aggregate cache.

3D Parallelism: Combining All Three

Production systems at scale combine all three strategies. A typical configuration for training a 175B model on 64 GPUs might use:

# 3D Parallelism configuration example
# 64 GPUs organized as: TP=4 × PP=4 × DP=4

# Within each node (8 GPUs connected by NVLink):
#   - Tensor Parallelism groups of 4 GPUs
# Across nodes (connected by InfiniBand):
#   - Pipeline Parallelism across 4 stages
#   - Data Parallelism with 4 replicas

# NVIDIA NeMo configuration:
trainer:
  devices: 8
  num_nodes: 8
model:
  tensor_model_parallel_size: 4    # TP within node
  pipeline_model_parallel_size: 4  # PP across nodes
  # DP = total_gpus / (TP × PP) = 64 / 16 = 4

The general principle: use tensor parallelism within a node (high bandwidth NVLink), pipeline parallelism across nodes (lower bandwidth network), and data parallelism for everything else. This hierarchy matches the communication requirements of each strategy to the available interconnect bandwidth.

Expert Parallelism (Bonus)

Mixture-of-Experts (MoE) models like Mixtral introduce a fourth dimension: expert parallelism. Each expert is placed on a different GPU, and the router dispatches tokens to the appropriate expert’s GPU. This is orthogonal to the other three strategies and enables models with massive total parameter counts while keeping per-token computation manageable.

Practical Applications

Inference Serving with vLLM

vLLM is the most popular LLM serving framework and supports all three parallelism strategies. For serving a 70B model, you might use tensor parallelism across 4 H100s within a node:

# Serve Llama 70B with tensor parallelism
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 4 \
  --max-model-len 8192

# Scale throughput with data parallelism (run multiple replicas)
# Deploy 2 replicas behind a load balancer, each with TP=4
# Total: 8 GPUs, TP=4 × DP=2

Training with DeepSpeed ZeRO

DeepSpeed’s ZeRO (Zero Redundancy Optimizer) is an advanced form of data parallelism that eliminates memory redundancy by partitioning optimizer states (ZeRO-1), gradients (ZeRO-2), and parameters (ZeRO-3) across GPUs instead of replicating them. ZeRO-3 achieves the memory efficiency of model parallelism with the simplicity of data parallelism.

When to Use Each Strategy

# Decision framework for parallelism strategy
Constraint                → Strategy
──────────────────────────────────────────────────
Model fits on one GPU     → Data Parallelism (scale throughput)
Model doesn't fit         → Tensor or Pipeline Parallelism
Need lowest latency       → Tensor Parallelism (requires NVLink)
Cross-node distribution   → Pipeline Parallelism
High request volume       → Data Parallelism (independent replicas)
Training at scale         → 3D Parallelism (TP + PP + DP)
MoE architecture          → Expert Parallelism + TP/PP/DP

Communication Bandwidth Requirements

The interconnect between GPUs is often the limiting factor:

Key Takeaways

Further Reading