Data Parallelism

Data parallelism is the simplest and most widely used distributed training strategy. Every GPU has a full copy of the model, processes a different slice of the batch, and synchronizes gradients via AllReduce. This chapter traces the evolution from naive DP to DDP to FSDP.

---

1. The Core Idea

GPU 0: model copy + batch shard 0 → gradients 0 ──┐
GPU 1: model copy + batch shard 1 → gradients 1 ──┤── AllReduce (average) ──→ all GPUs get averaged gradients
GPU 2: model copy + batch shard 2 → gradients 2 ──┤                          → optimizer step (identical on all)
GPU 3: model copy + batch shard 3 → gradients 3 ──┘

Mathematical equivalence: Training with batch size $B$ on $P$ GPUs, each processing $B/P$ samples, with averaged gradients is mathematically identical to training with batch $B$ on one GPU. Same convergence, $P\times$ throughput.

---

2. Evolution: DP → DDP → FSDP

DataParallel (DP) — Deprecated

PyTorch’s original nn.DataParallel: one process, multiple GPUs. GPU 0 is the “master” — it scatters input, gathers output, and does all gradient reduction.

model = nn.DataParallel(model)  # don't use this

Three fatal flaws:

1. GIL bottleneck: Single Python process → single GIL → GPU 0 becomes the bottleneck
2. Imbalanced memory: GPU 0 holds all gradients for reduction → OOM while others have spare memory
3. Redundant copies: Scatter/gather through GPU 0 wastes NVLink bandwidth

Never use DP. Use DDP.

DistributedDataParallel (DDP) — Standard

One process per GPU. No master — all GPUs are equal. Gradients synchronized via AllReduce (ring algorithm through NCCL).

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

dist.init_process_group('nccl')
model = DDP(model.cuda(), device_ids=[local_rank])

# Training loop is identical to single-GPU
for batch in dataloader:
    loss = model(batch)
    loss.backward()       # DDP hooks AllReduce into backward automatically
    optimizer.step()
    optimizer.zero_grad()

How DDP works:

1. Forward pass: each GPU processes its batch shard independently (no communication)
2. Backward pass: as gradients are computed, DDP buckets them and starts AllReduce
3. Gradient bucketing: gradients are grouped into ~25 MB buckets. AllReduce starts per-bucket as soon as that bucket’s gradients are ready — overlapping communication with backward computation
4. After AllReduce: all GPUs have identical averaged gradients → optimizer step produces identical parameters

FSDP / ZeRO — Memory-Efficient

DDP replicates the full model on every GPU. For a 70B model: 140 GB per GPU just for BF16 parameters — doesn’t fit on any single GPU.

FSDP (Fully Sharded Data Parallel) / ZeRO shard model state across GPUs:

Stage	What’s sharded	Memory per GPU	Communication
DDP (no sharding)	Nothing	Full model + optimizer	AllReduce (gradients)
ZeRO Stage 1	Optimizer states	~60% of DDP	Same as DDP
ZeRO Stage 2	+ Gradients	~40% of DDP	ReduceScatter + AllGather
ZeRO Stage 3 / FSDP	+ Parameters	~1/P of DDP	AllGather (fwd) + ReduceScatter (bwd)

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model, auto_wrap_policy=transformer_auto_wrap_policy)

FSDP tradeoff: Less memory per GPU (can train larger models), but more communication (AllGather before each layer’s forward pass).

---

3. Communication Analysis

DDP: AllReduce Gradients

Communication per step: one AllReduce of all gradients.

$$\text{Volume} = \frac{2N(P-1)}{P} \times \text{bytes/param} \approx 2N \times 2 \text{ bytes (BF16)}$$

For 7B model, BF16: $2 \times 7 \times 10^9 \times 2 = 28$ GB per step.

On 8 GPUs with NVLink (450 GB/s): $28 / 450 \approx 62$ ms — but mostly hidden behind backward (gradient bucketing).

FSDP: AllGather + ReduceScatter Per Layer

For each of $L$ layers:

• Forward: AllGather parameters ($\sim N/L$ params) — needs full params for compute
• Backward: ReduceScatter gradients ($\sim N/L$ params) — each GPU keeps its gradient shard

Total per step: $L \times \text{AllGather} + L \times \text{ReduceScatter}$

Volume is the same as DDP ($\approx 2N$), but spread across $2L$ smaller operations instead of one large AllReduce. The advantage: each AllGather can overlap with the previous layer’s computation.

---

4. Scaling Behavior

Strong Scaling (Fixed Total Batch)

$P$ GPUs, each processes $B/P$ samples. As $P$ increases:

• Compute per GPU: $\propto 1/P$ (linear speedup ✅)
• Communication: $\approx$ constant (AllReduce volume ≈ $2N$, independent of $P$ ✅)
• But: Per-GPU batch $B/P$ shrinks → GPU utilization drops when batch is too small

Practical limit: When per-GPU batch < ~4 samples, GPU utilization drops below 50%. Solution: gradient accumulation (process multiple micro-batches before AllReduce).

Weak Scaling (Fixed Per-GPU Batch)

$P$ GPUs, each processes $B$ samples. Total batch = $PB$. As $P$ increases:

• Compute per GPU: constant ✅
• Communication: constant ✅
• Total batch grows: may need learning rate scaling ($\text{lr} \propto \sqrt{PB}$)

This is how large-scale training works: 1000 GPUs, each with batch 4, total batch 4000. Communication overhead is <5% on good networks.

---

For details on FSDP/ZeRO memory sharding, see the Training Systems > Memory Optimization module. For how AllReduce works, see Communication Primitives.

5. When Data Parallelism Is Not Enough

Problem	Why DP can’t solve it	Solution
Model doesn’t fit on one GPU	DP replicates the full model	Tensor/Pipeline parallelism or FSDP
Sequence too long for one GPU	Attention memory $O(S^2)$ per GPU	Context/Sequence parallelism
MoE experts too many for one GPU	Each expert is a separate FFN	Expert parallelism
Communication becomes bottleneck	Only happens at >1000 GPUs or slow networks	Gradient compression, hierarchical AllReduce

In practice: Data parallelism is the outer dimension — you always use it. Other parallelism strategies are the inner dimensions, composed with DP.

---

6. Practical DDP Template

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup():
    dist.init_process_group('nccl', timeout=timedelta(seconds=300))
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    return local_rank

def train():
    local_rank = setup()

    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    sampler = DistributedSampler(dataset, shuffle=True)
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=32,
                           num_workers=4, pin_memory=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # ensure different shuffle per epoch
        for batch in dataloader:
            batch = batch.cuda(local_rank, non_blocking=True)
            loss = model(batch)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()

# Launch: torchrun --nproc_per_node=8 train.py

---

Self-Check

• Explain why DDP is mathematically equivalent to single-GPU training with larger batch
• Explain the three flaws of nn.DataParallel and why DDP fixes them
• Explain gradient bucketing and how it overlaps communication with backward
• Explain ZeRO Stage 1/2/3 and the memory-communication tradeoff
• Calculate AllReduce volume for a 7B model in BF16
• Explain strong vs weak scaling and when each hits limits
• Write a complete DDP training script with DistributedSampler

---

References

• PyTorch DDP Tutorial
• PyTorch FSDP Tutorial
• ZeRO: Memory Optimizations (Rajbhandari et al., 2019)