Distributed Training: Data, Tensor & Pipeline Parallelism

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Data Parallelism (DDP)
5. ZeRO: Memory-Efficient Data Parallelism
6. Tensor Parallelism
7. Pipeline Parallelism
8. 3D Parallelism Configuration
9. Communication Optimization
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Compare data, tensor, and pipeline parallelism tradeoffs.
2. Explain ZeRO stages 1, 2, 3 and their memory savings.
3. Derive tensor parallelism for MLP and attention layers.
4. Analyze pipeline parallelism bubble overhead.
5. Configure 3D parallelism for a given model size and cluster.

Notation

• $N$ — number of GPUs
• $B$ — global batch size, $b$ — per-GPU batch size
• TP — tensor parallelism degree
• PP — pipeline parallelism degree
• DP — data parallelism degree

Core Intuition

A single GPU can't hold or efficiently train a 70B model (needs 280GB in FP32 for weights + optimizer). Distributed training splits the work across GPUs in three orthogonal dimensions: DATA (same model, different data), TENSOR (split single layers across GPUs), PIPELINE (split layers across GPUs). Each has different tradeoffs in memory, communication, and efficiency.

Data Parallelism (DDP)

Each GPU: Full model copy + different data batch.

AllReduce: After backward pass, average gradients across all GPUs:

Properties:

• Linear scaling: $N$ GPUs → $N\times$ throughput (ideally).
• No model partitioning needed.
• Requires: model fits on ONE GPU (with activations + optimizer states).

Memory per GPU: Full model + full optimizer + activations. For 7B in BF16 with AdamW: $14 + 56 + \sim30 = 100$ GB. Fits on A100 80GB only with activation checkpointing.

Communication: AllReduce of all gradients once per step. Volume = model size. With Ring-AllReduce: bandwidth-optimal.

ZeRO: Memory-Efficient Data Parallelism

ZeRO (Rajbhandari et al., 2020): Shard optimizer states, gradients, and/or parameters across GPUs:

Stage 1 (Optimizer state sharding):

• Each GPU stores 1/N of the optimizer states.
• Memory savings: 4x (from 16 bytes/param to 4 bytes/param per GPU).
• Communication: same as DDP.

Stage 2 (+ Gradient sharding):

• Gradients also sharded; reduce-scatter instead of AllReduce.
• Memory savings: 8x.
• Communication: same volume, different pattern.

Stage 3 (+ Parameter sharding):

• Parameters also sharded; AllGather before each layer's forward.
• Memory savings: proportional to $N$.
• Communication: 1.5x DDP (AllGather + ReduceScatter).

FSDP (PyTorch): Implementation of ZeRO Stage 3 in PyTorch native.

Tensor Parallelism

Split individual layers across GPUs:

MLP (column-then-row split):

Split $\mathbf{W}_1$ column-wise across $T$ GPUs: each GPU computes partial activation. GeLU applied locally. $\mathbf{W}_2$ split row-wise. Final AllReduce to sum partial results.

Attention (head-parallel):

• Split heads across GPUs ($H/T$ heads per GPU).
• Each GPU computes its heads independently.
• AllReduce after output projection.

Communication: 2 AllReduce ops per transformer layer (one after MLP, one after attention). Low volume but HIGH frequency.

When to use: Layers too large for one GPU, or when low latency per step matters (TP keeps all GPUs busy simultaneously).

Pipeline Parallelism

Split layers sequentially across GPUs:

Naive pipeline: Sequential execution → only 1/P GPUs active at a time (pipeline bubble).

Micro-batching (GPipe): Split batch into $M$ micro-batches; pipeline them:

• Bubble fraction: $\frac{P-1}{M+P-1}$.
• With $M=4P$: bubble = $\frac{P-1}{5P-1} \approx 20\%$.

Interleaved pipeline (1F1B): Each GPU holds non-contiguous layers. Reduces bubble to near-zero with enough micro-batches.

Communication: Point-to-point send/recv of activations between adjacent GPUs. Low bandwidth requirement but adds latency.

3D Parallelism Configuration

Total GPUs = DP × TP × PP.

Example: 70B model on 64 A100 80GB:

• TP = 8 (one node, NVLink connected).
• PP = 2 (across 2 nodes).
• DP = 4 (4 pipeline replicas, different data).
• Total: 8 × 2 × 4 = 64 GPUs.

Rules of thumb:

• TP within a node (requires high bandwidth: NVLink 600 GB/s).
• PP across nodes (low communication: only send activations).
• DP across remaining GPUs (moderate communication: gradient sync).

Memory budget per GPU:

Communication Optimization

Overlap communication with computation:

• Start AllReduce for layer $l$'s gradients while computing layer $l-1$'s gradients.
• Overlap AllGather of next layer's params with current layer's compute (ZeRO-3/FSDP).

Gradient compression: Reduce communication volume:

• FP16 gradients (2x savings over FP32).
• PowerSGD: Low-rank approximation of gradients.
• TopK sparsification: Only communicate largest gradients.

Topology-aware placement: Place TP groups on NVLink-connected GPUs; PP groups on same rack; DP across racks.

Common Pitfalls

Pitfall 1. Using tensor parallelism across nodes (over InfiniBand). TP requires 2 AllReduce per layer — too much latency over network. Keep TP within NVLink nodes.

Pitfall 2. Pipeline parallelism with too few micro-batches. If $M < 2P$: bubble exceeds 33%, wasting GPUs. Use $M \geq 4P$.

Pitfall 3. FSDP/ZeRO-3 without communication overlap. Without overlapping AllGather with compute, every layer stalls waiting for parameters. Always enable prefetch_factor.

Summary

• Data parallelism: Simple; linear scaling; requires model fits on 1 GPU.
• ZeRO (1/2/3): Shard optimizer/gradients/params; enables any model size.
• Tensor parallelism: Split layers across GPUs; lowest latency; highest bandwidth needs.
• Pipeline parallelism: Split model sequentially; tolerates low bandwidth; has bubble overhead.
• 3D parallelism: Combine all three for maximum scale.
• Communication: Overlap with compute; compress gradients; topology-aware placement.

Exercises

Exercise 1. For a 13B model with AdamW in BF16: compute memory per GPU for (a) DDP, (b) ZeRO-2, (c) ZeRO-3 on 8 GPUs.

Exercise 2. For pipeline parallelism with PP=4 and M=16 micro-batches: compute the bubble fraction and GPU utilization.

Exercise 3. Design a 3D parallelism config for a 175B model on 512 A100 80GB GPUs. Specify TP, PP, DP and justify each choice.

Exercise 4. Compute the AllReduce volume per step for DDP vs the AllGather+ReduceScatter volume for FSDP on a 7B model.

Exercise 5. Compare the time-to-train for a 70B model: (a) 64 A100 with 3D parallelism, vs (b) 256 H100 with FSDP only. Account for communication overhead.