Distributed Training: DP / DDP / FSDP / Tensor Parallel

L7-01 explained why AI uses GPUs. But:

A single 80GB H100 fits at most a 7B model with full training state. 70B → multiple cards. GPT-4-scale (~1.7T MoE) → tens of thousands.

This piece: how do those cards work together?

The Two Bottlenecks

Distributed training is constrained by:

1. Memory — can the model + gradients + optimizer state fit?
2. Communication — keeping cards in sync across the network

The architectural choices below all trade between these two.

Memory Math: Why Models Get So Big

Training a 7B model in FP16 needs about:

Item	Size	Multiplier
Weights	14 GB	7B × 2 bytes
Gradients	14 GB	same shape as weights
Adam state (m, v)	28 GB	2× weights
Activations	~30 GB	varies with batch/seq
Total	~80 GB	Just fits H100

Adam state is bigger than the model itself. For 70B → 700GB+ total → multi-GPU mandatory.

Strategy 1: Data Parallel (DP / DDP)

Simplest. Replicate model on each GPU, split data:

GPU 0: full model | mini-batch 1
GPU 1: full model | mini-batch 2
GPU 2: full model | mini-batch 3
GPU 3: full model | mini-batch 4

After each step:

1. Each GPU computes its gradient
2. All-reduce gradients (average across all cards)
3. Each card updates its model

DP (legacy)

PyTorch’s nn.DataParallel — runs on one machine, single process splits work to multiple GPUs. Slow (GIL, single-process bottleneck). Don’t use it.

DDP (modern)

nn.parallel.DistributedDataParallel — one process per GPU, communicates via NCCL.

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model)
# ... train as normal

Limit: model must fit in one GPU’s memory. 7B yes. 70B no.

Strategy 2: Model / Tensor Parallel

Model too big to fit one GPU → split the model.

Pipeline Parallelism (Pipeline Parallel)

Split by layer:

GPU 0: layers 1-3
GPU 1: layers 4-6
GPU 2: layers 7-9
GPU 3: layers 10-12

Forward: GPU 0 → GPU 1 → GPU 2 → GPU 3 → output Backward: same in reverse

Problem: while GPU 0 works on micro-batch 1, GPU 1, 2, 3 are idle. Bubbles.

Solution: micro-batches + 1F1B scheduling — see the pipeline visualization (/visualize/pipeline).

Tensor Parallel

Split inside a single layer.

A standard linear: $y = Wx$ where $W \in \mathbb{R}^{N \times M}$. Split W column-wise across 4 GPUs:

GPU 0: W_0 (columns 0 to M/4)
GPU 1: W_1 (columns M/4 to M/2)
GPU 2: W_2 (...)
GPU 3: W_3 (...)

Each GPU computes a partial output, then all-gather to merge.

Used by: Megatron-LM, Megatron-Core. Heavy intra-node communication — needs NVLink within node.

Strategy 3: ZeRO / FSDP — Sharded Optimizer

Microsoft’s DeepSpeed ZeRO (2019) had an insight:

Adam state and gradients don’t need to be replicated — shard them across GPUs.

ZeRO Stage 1

Shard only optimizer state:

GPU 0: optimizer state for layers 1-3
GPU 1: optimizer state for layers 4-6
GPU 2: optimizer state for layers 7-9
GPU 3: optimizer state for layers 10-12

(all GPUs hold full weights + full gradients)

Adam state was 50% of memory → save ~40%.

ZeRO Stage 2

Also shard gradients. Save more memory.

ZeRO Stage 3 (= FSDP)

Shard weights too — each card only stores 1/N of the weight, fetches others when needed.

GPU 0: weight shards for layers 1-3
GPU 1: weight shards for layers 4-6
...

Forward: when needed, fetch the full layer from peer cards (all-gather)
         use, then discard

PyTorch built this in as FSDP (Fully Sharded Data Parallel):

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model)
# train as normal

Effect:

• 70B model trainable on 8× A100 80GB
• Almost like having a single GPU with 8× the memory
• Communication cost slightly higher (need extra all-gather)

Strategy 4: 3D Parallelism

The most complex setups combine all of the above:

Data Parallel (across nodes) × Tensor Parallel (within node) × Pipeline Parallel (across stages)

This is what GPT-4 / Llama 3 use. Configuration examples:

512 GPU cluster:
- 8-way Tensor Parallel (within node, 8 cards)
- 8-way Pipeline Parallel (across 8 nodes)
- 8-way Data Parallel (8 replicas)
8 × 8 × 8 = 512 ✓

Tuning the ratios = optimization art.

Communication Primitives

NCCL provides these:

Primitive	What it does
All-Reduce	Sum data across all GPUs, give result to all (DDP gradient sync)
All-Gather	Collect data from all GPUs, all have full copy (FSDP weight fetch)
Reduce-Scatter	All-Reduce + Scatter (used in ZeRO)
Broadcast	One GPU → all GPUs (parameter init)

Each has cost scaling with N (number of GPUs) × data size. Reducing the number of comm ops is the main optimization.

Practical Recipe

Scenario	Strategy
Small model (< 7B), single node	DDP
Mid model (7B-70B), single node multi-GPU	FSDP (Stage 3)
Big model (70B-200B), multi-node	FSDP + Pipeline
Frontier (1T+)	3D Parallelism
Inference only	KV cache + tensor parallel

Engineering Pain Points

1. Hyperparameters: micro-batch size, gradient accumulation, etc.

Each strategy has knobs. Wrong settings → 50% performance drop. DeepSpeed’s auto-tuning helps.

2. Faults

10,000 GPU cluster — one card dies a day on average. Need:

• Checkpoint frequently (e.g., every 100 steps)
• Auto-recovery (skip failed steps, restart from checkpoint)
• Health monitoring (GPU temp, network, NCCL warnings)

Meta’s Llama 3 training had a hardware failure every few hours over months.

3. Communication bottleneck

H100: 989 TFlops compute
NVLink:  900 GB/s (within node)
InfiniBand: 200 Gb/s (across node, = 25 GB/s)

Cross-node bandwidth is 36× slower than within node. So you want as much intra-node communication as possible (Tensor Parallel inside, Data Parallel outside).

4. Logging and debugging

8000-GPU training run — how do you know if it’s actually progressing?

• Loss curve (smoothed over many ranks)
• Throughput (tokens/sec)
• Gradient norms (look for explosions)
• Memory pressure (any rank running low?)

Tools: Weights & Biases, TensorBoard, custom dashboards.

Real-World Examples

Llama 3.1 (Meta)

• 405B params
• 16k H100 trained over 3 months
• 15T tokens training data
• Strategy: 3D parallelism (TP + PP + DP)
• Estimated cost: $60-100M

DeepSeek V3 (2024)

• 671B MoE (37B active per token)
• Only ~2k H800 (export-restricted variant)
• Used FP8 mixed precision + various clever tricks
• Reportedly trained for ~$6M
• → much better cost efficiency than Llama 3.1 — engineering matters

GPT-4 (rumored)

• ~1.7T MoE
• Massive scale, undisclosed
• Estimated cost: $50-100M

Open Source Tools

Library	What it does
PyTorch DDP / FSDP	Built-in, sufficient for most cases
DeepSpeed	Microsoft, ZeRO + optimization tricks
Megatron-LM	NVIDIA, Tensor Parallel + Pipeline
ColossalAI	All strategies in one library
Accelerate	HuggingFace, simple wrapper

Recommendation: small projects → Accelerate / FSDP; large-scale production → DeepSpeed or Megatron.

Next recommended: L7-03 Inference Optimization or L7-04 Quantization.