Model parallelism (tensor, pipeline)

When models grow too large to fit on a single GPU's memory, model parallelism splits the model itself across multiple devices, unlike data parallelism which replicates the entire model.

Connections

• Data Parallelism - complementary approach (replicate model, split data)
• Gradient Accumulation - reduces memory for batch size
• Mixed Precision Training - reduces memory per parameter
• Transformer Architecture - common target for model parallelism
• GPU Memory Hierarchy - hardware constraint driving the need
• Activation Checkpointing - trades compute for memory

Model parallelism says: "If the model doesn't fit on one device, cut the model into pieces and put each piece on a different device."

Two main strategies exist:

1. Tensor Parallelism: Split individual layers (cutting matrices horizontally/vertically)
2. Pipeline Parallelism: Split model depth-wise (early layers on GPU-0, later layers on GPU-1)

Tensor Parallelism

Derivation: How to Split a Matrix Multiply

Consider a simple linear layer: $Y = XW$ where:

• $X \in \mathbb{R}^{b \times d_{in}}$ (batch × input dimension)
• $W \in \mathbb{R}^{d_{in} \times d_{out}}$ (weight matrix)
• $Y \in \mathbb{R}^{b \times d_{out}}$ (output)

Why can we split this? Matrix multiplication is associative with concatenation.

Column-wise Split (most common)

Split $W$ into $N$ column slices: $W = [W_1 | W_2 | \ldots | W_N]$ where each $W_i \in \mathbb{R}^{d_{in} \times (d_{out}/N)}$.

Derivation: $Y = XW = X[W_1 | W_2 | \ldots | W_N] = [XW_1 | XW_2 | \ldots | XW_N]$

So device $i$ computes $Y_i = XW_i$, then we concatenate: $Y = [Y_1 | Y_2 | \ldots | Y_N]$.

Key insight: Each device needs the full input $X$ but only a fraction of $W$. Memory saved on weights, but need to broadcast $X$ to all devices.

Why this step? The column split works because matrix-vector multiplication distributes: each column block of $W$ produces an independent column block of $Y$.

What about the backward pass?

Gradients flow: $\frac{\partial L}{\partial W} = X^T \frac{\partial L}{\partial Y}$

If $Y$ is split column-wise, then $\frac{\partial L}{\partial Y}$ is also split. Device $i$ computes: $\frac{\partial L}{\partial W_i} = X^T \frac{\partial L}{\partial Y_i}$

Critical: To compute $\frac{\partial L}{\partial X}$ (needed for earlier layers): $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} W^T = \sum_{i=1}^{N} \frac{\partial L}{\partial Y_i} W_i^T$

This requires an all-reduce to sum contributions from all devices.

Why this step? The chain rule demands we sum over all output dimensions to get gradients w.r.t. inputs Since outputs are split, we must communicate to aggregate.

A Transformer MLP has two linear layers: $\text{MLP}(x) = \text{GELU}(xW_1)W_2$

where $W_1: d_{model} \to 4d_{model}$ (expand) and $W_2: 4d_{model} \to d_{model}$ (project).

Column-parallel on $W_1$, Row-parallel on $W_2$:

1. Device $i$ of $N$: Store $W_1^{(i)} \in \mathbb{R}^{d_{model} \times (4d_{model}/N)}$
2. Compute locally: $h_i = \text{GELU}(x W_1^{(i)})$
3. No communication yet! Each device has a slice of the hidden state
4. Store $W_2^{(i)} \in \mathbb{R}^{(4d_{model}/N) \times d_{model}}$ (row slice)
5. Compute: $y_i = h_i W_2^{(i)}$
6. All-reduce: $y = \sum_{i=1}^{N} y_i$

Why this pattern? The column-parallel followed by row-parallel design minimizes communication: only one all-reduce per MLP block, not one per layer.

Memory saved: $W_1$ and $W_2$ are each split by $N$, so roughly $1/N$ of the MLP parameters per device.

Why this step? The activation after $W_1$ is an intermediate that doesn't need to be gathered. We only gather the final output after $W_2$, reducing communication overhead.

Attention has projection matrices $W_Q, W_K, W_V \in \mathbb{R}^{d_{model} \times d_{model}}$ and output projection $W_O$.

Split across attention heads: If8 heads, put 2 heads on each of 4 devices.

Device $i$ computes: $\text{head}_{2i}, \text{head}_{2i+1} = \text{Attention}(Q_i, K_i, V_i)$

where $Q_i = XW_Q^{(i)}$, etc., with $W_Q^{(i)} \in \mathbb{R}^{d_{model} \times (d_{model}/4)}$.

Then concatenate and apply $W_O$ (row-parallel), followed by all-reduce.

Why this step? Attention heads are independent, so they naturally parallelize. Each device computes a subset of heads, reducing memory and compute per device.

Pipeline Parallelism

Naive Pipeline: The Bubble Problem

Setup: Model with 8 layers, 4 GPUs. GPU-0 has layers 0-1, GPU-1 has layers 2-3, etc.

Forward pass:

1. GPU-0 processes batch → sends activations to GPU-1
2. GPU-1 processes → sends to GPU-2
3. etc.

Problem: GPU-1 sits idle while GPU-0 computes. GPU-2 idles while GPU-0 and GPU-1 work. This creates a "bubble" of unused compute time.

Utilization: With $D$ devices and $B$ micro-batches, bubble time is approximately $\frac{D-1}{B+D-1}$ of total time. For $D=4, B=1$, that's 75% idle!

Why this happens: Sequential dependency. Later stages wait for earlier stages. This is the fundamental trade-off of pipeline parallelism.

GPipe: Micro-Batching Solution

Timeline:

GPU-0: [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]
GPU-1:      [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]
GPU-2:           [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]
GPU-3:                [F₀] [F₁] [F₂] [F₃] [B₀] [B₁] [B₂] [B₃]

$F_i$ = forward micro-batch $i$, $B_i$ = backward micro-batch $i$

Bubble fraction: $\frac{D-1}{M+D-1}$ where $M$ = number of micro-batches.

With $M=16, D=4$: bubble = $3/19 \approx 16\%$. Much better!

Cost: Need to store activations for all $M$ micro-batches for backward pass. Memory increases linearly with $M$.

Why this step? More micro-batches keep later stages busy while early stages are still processing. The pipeline fills up, reducing idle time at the cost of activation memory.

Model: 48 layers, 4 GPUs (12 layers each), micro-batch size 1, total batch 32.

Scenario 1: $M=8$ micro-batches

• Bubble time: $(4-1)/(8+4-1) = 27\%$
• Activation memory: 8× single micro-batch
• Throughput: ~5.5 samples/sec

Scenario 2: $M=32$ micro-batches

• Bubble time: $(4-1)/(32+4-1) = 8.6\%$
• Activation memory: 32× single micro-batch
• Throughput: ~6.8 samples/sec

Trade-off: Higher $M$ improves efficiency but increases memory. GPipe uses activation checkpointing (recompute activations during backward) to reduce memory at cost of 33% more compute.

Why this example? Shows the quantitative impact of pipeline depth. Real systems must balance $M$ based on available memory and desired throughput.

PipeDream: 1F1B Schedule

Problem with GPipe: All forward passes, then all backward passes. Memory spike for activations.

PipeDream solution: Interleave forward and backward: "1 Forward, 1 Backward" schedule.

Timeline:

GPU-0: [F₀] [F₁] [F₂] [F₃] [B₀] [F₄] [B₁] [F₅] [B₂] [F₆] [B₃] ...
GPU-1:      [F₀] [F₁] [F₂] [B₀] [F₃] [B₁] [F₄] [B₂] [F₅] [B₃] ...

Why this helps: Backward passes for early micro-batches happen before later forward passes. Activations are freed soner. Steady-state memory: only $D$ micro-batches' activations (# of pipeline stages), not $M$ (total micro-batches).

Complication: Earlier layers see gradients from older micro-batches (version inconsistency). PipeDream uses weight stashing: keep multiple weight versions during training to ensure each micro-batch's forward and backward use the same weights.

Why this step? By overlapping forward and backward passes, we reduce the peak memory from $O(M)$ to $O(D)$, making deeper pipelines feasible without running out of memory.

Why this feels right: Pipeline seems simpler—just split layers, no complex matrix slicing.

Why it's wrong: Pipeline has unavoidable bubble time and inter-device communication between every layer partition. Tensor parallelism has communication within each layer but can process batches in parallel across all devices simultaneously (no bubble in steady state).

When each wins:

• Tensor parallelism: Better for wide models (large hidden dimensions), high inter-device bandwidth (NVLink). Communication cost is proportional to hidden size, not depth.
• Pipeline parallelism: Better for deep models with limited bandwidth. Communication happens only at partition boundaries (fewer times per forward pass).

Reality: Modern systems (Megatron-LM, DepSpeed) combine both. Use tensor parallelism within a node (fast NVLink) and pipeline parallelism across nodes (slower InfiniBand).

The fix: Profile your model and hardware. If communication time< bubble time, prefer tensor. If memory per layer is the bottleneck, prefer pipeline.

Why this feels right: Smaller batches → less activation memory → model fits on one GPU.

Why it's wrong: Model parameters and optimizer states are independent of batch size. A 175B parameter model needs 700GB for weights alone (FP32). Reducing batch size to 1 doesn't change this.

The confusion: Activation memory (forward pass intermediate results) does scale with batch size. But parameter memory dominates for large models.

The numbers:

• Parameters: 175B × 4 bytes = 700GB (fixed)
• Optimizer states (Adam): 175B × 12 bytes = 2.1TB (fixed)
• Activations: ~batch_size × seq_len × hidden_dim × num_layers (variable)

For GPT-3, even with batch size 1, parameter + optimizer memory exceds any single GPU.

The fix: Model parallelism is necessary when parameter count exceeds device memory. Activation memory reduction (gradient accumulation, checkpointing) is complementary.

Communication Patterns

Tensor Parallelism (per layer): $\text{Comm}_{\text{tensor}} = 2bsh$

where $b$ = batch size, $s$ = sequence length, $h$ = hidden dimension.

Factor of 2: all-gather after forward (to broadcast activations) + reduce-scatter after backward (to aggregate gradients).

Pipeline Parallelism (per stage boundary): $\text{Comm}_{\text{pipeline}} = bsh$

Only forward activation sent to next stage and backward gradient sent to previous stage. No all-reduce within the pipeline.

Why this matters: For tensor parallelism, communication grows with hidden size. For pipeline, it's constant per partition boundary. This guides the hybrid strategy.

Goal: Train 530B parameter model (Megatron-Turing NLG).

Hardware:

• 4480 GPUs (A100 80GB)
• Nodes with 8 GPUs connected by NVLink (900 GB/s)
• Nodes connected by InfiniBand (200Gb/s)

Configuration:

• Tensor parallelism: $N=8$ within each node (fast NVLink makes all-reduce cheap)
• Pipeline parallelism: $P=35$ stages across nodes (reduces cross-node communication)
• Data paralelism: 16-way across pipeline replicas

Why this works:

• Tensor paralelism exploits fast intra-node bandwidth
• Pipeline paralelism limits slow inter-node bandwidth use
• Data parallelism scales training across the full cluster

Memory per GPU:

• Model: 530B / (8 × 35) = 1.9B parameters per GPU ≈ 7.6GB (FP32)
• With optimizer states: ~23GB per GPU
• Leaves room for activations and workspace

Why this step? This is a real-world example of how paralelism strategies compose. Each dimension of paralelism addresses a different bottleneck.

Imagine you're trying to build the world's biggest LEGO castle, but it's so huge that all the pieces won't fit on your table. What do you do?

Option 1 (Tensor Parallelism): You and your friend each take half the pieces for same wall. You both build your half, then snap them together. For the next wall, you split pieces again. You work together on each part, but share the load.

Option 2 (Pipeline Parallelism): You build the first floor, then pass it to your friend who builds the second floor, while you start the first floor of a new castle. Like an assembly line—everyone has their own stage, and castles flow through.

The problem? In Option 1, you both need to see the instruction book (communication). In Option 2, your friend waits while you finish (bubble time).

Real AI systems do both! Fast splits (tensor) with friends at your table, assembly line (pipeline) between tables across the room.

TENSOR = Together Each Node Slices One Row → All devices work on the same layer, split horizontally

PIPELINE = Pass Inputs Progressively Each Layer In Next Element → Data flows through layers like a pipe

When to use: "Tight bandwidth → Tensor; Poor links → Pipeline"

Flashcards

What is the fundamental difference between data parallelism and model parallelism? :: Data parallelism replicates the entire model on each device and splits the data batch. Model parallelism splits the model itself across devices because it's too large to fit on one device.

What is pipeline parallelism? :: Splitting the model depth-wise across devices, where each device holds different layers and data flows sequentially like a pipeline (inter-layer parallelism).

Concept Map

Hinglish (regional understanding)