PyTorch Distributed Training Engine Dissected: Design Evolution from DDP to FSDP

Experiment Goal

Getting DDP/FSDP running with torchrun is easy, but "why does DDP double the VRAM", "why does FSDP have so much inter-card communication", "when are gradients actually synced" — these questions decide whether you can tune the engine optimally at thousand-GPU scale. This article dissects PyTorch distributed training from the framework engine mechanism view:

• How does DDP sync gradients exactly? How do communication and compute overlap?
• What does FSDP change relative to DDP? How are parameter sharding and communication layers designed?
• Which scenarios pick DDP, which pick FSDP?

This is a mechanism-analysis + config-practice route; commands and configs are generic practice examples without fabricated performance benchmarks; concrete implementation follows the installed PyTorch version.

Section 1: DDP — Engineering Design of Gradient Synchronization

DDP (DistributedDataParallel) logic is simple: each card holds a full model copy + its own data; after backward, gradients are AllReduced and each card updates with the averaged gradient. The engineering difficulty lies in sync timing and overhead.

1. Gradient Buckets: Compress the Number of AllReduces

A naive approach AllReduces each parameter's gradient separately — a trillion-parameter model has tens of thousands of gradient tensors and communications, unacceptable. DDP uses gradient buckets:

# Conceptual snippet: DDP packs consecutive parameter gradients into buckets, AllReducing per bucket
# bucket_cap_mb: default 25MB; bucket size affects communication granularity and overlap
ddp_model = torch.nn.parallel.DistributedDataParallel(
    model,
    bucket_cap_mb=25,        # bucket capacity; small models can lower (e.g. 5MB)
    gradient_as_bucket_view=True,   # reuse bucket memory, fewer copies
)

• Bucket = communication granularity: larger buckets → fewer communications but longer per-bucket waits; smaller buckets → finer overlap but more communications. 25MB is the default balance; small models can lower it;
• gradient_as_bucket_view=True: gradients live as views into bucket memory, avoiding extra copies (VRAM and time win).

2. Communication-Compute Overlap: "Communicate While Computing" in Backward

DDP's core optimization is communicate-as-soon-as-ready (gradient ready → bucket ready → async AllReduce), overlapping with subsequent layers' backward compute:

backward pass
├─ layer N computes gradient → bucket ready → AllReduce starts (async, non-blocking)
├─ layer N-1 computes gradient (overlapped with N's communication) → ...
└─ all done → last AllReduce completes → update

• Overlap effectiveness depends on bucket split and layer order: buckets split by parameter registration order; if a bucket spans many layers, its readiness is dragged by the longest layer — why model registration order affects DDP performance;
• Verify overlap: use PyTorch Profiler to check for obvious communication gaps (AllReduce waiting for gradients) in the backward phase.

3. DDP's VRAM Accounting

• Each card holds the full model → VRAM ≈ single-card model + activations + gradients + optimizer states;
• Gradient AllReduce traffic ∝ model size (every card sends its own gradients); at thousand-card scale, communication is the bottleneck, not VRAM;
• Engineering observation: DDP fits "model fits one card, but data/compute needs many cards" (fine-tuning, small-to-medium models); when the model no longer fits one card, move to FSDP.

Section 2: FSDP — From "Full Model Per Card" to "Shard Per Card"

FSDP (Fully Sharded Data Parallel)'s core idea: parameters, gradients, and optimizer states are all sharded along the data-parallel dimension, All-Gathering the needed parameters at compute time.

1. Three-Level Sharding: Parameters / Gradients / Optimizer States

Dimension	DDP	FSDP
Parameters	full copy per card	1/DP shard per card, All-Gather on use
Gradients	full per card	shard per card (Scatter-Reduce keeps only local shard)
Optimizer states	full per card	shard per card (the big VRAM win)
Communication	AllReduce gradients per step	All-Gather params + Scatter-Reduce grads (~2-3× DDP traffic)

# FSDP basic usage
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

fsdp_model = FSDP(
    model,
    sharding_strategy=torch.distributed.fsdp.ShardingStrategy.FULL_SHARD,
    mixed_precision=torch.distributed.fsdp.MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    ),
)

2. Communication Layers: Flatten & All-Gather/Scatter-Reduce

• Parameter flatten: FSDP flattens (managed) parameters into contiguous storage and All-Gathers whole blocks at once instead of per-parameter communication — fewer communications;
• Forward: All-Gather a block's parameters before its compute, release after compute (VRAM fetched on demand);
• Backward: each card computes gradients for its own shard → Scatter-Reduce → keeps only its shard gradient → updates its shard's optimizer state;
• Traffic accounting: All-Gather (params) + Scatter-Reduce (grads) ≈ 2× parameter size per step — more than DDP's 1×. FSDP trades more communication for VRAM; with a weak network it may not be faster.

3. FSDP Tuning Knobs

Knob	Effect	Engineering advice
sharding_strategy	FULL_SHARD / SHARD_GRAD_OP (gradients only, less comm) / NO_SHARD	if VRAM allows, SHARD_GRAD_OP cuts communication
Sharding granularity (auto_wrap_policy)	per-layer/block vs whole-model	per-layer → finer All-Gather, better overlap
cpu_offload	offload params/grads to CPU	extreme VRAM cases; watch CPU↔GPU transfer cost
forward_prefetch / backward_prefetch	prefetch next block early	overlaps communication with compute; defaults usually better

Section 3: Selection Matrix — Which Scenario Uses What

Scenario	Recommendation	Why
Small/medium model (fits one card) + multi-card data scaling	DDP	least communication, simple, VRAM not the constraint
Large model (doesn't fit one card) training	FSDP (+ CPU offload if needed)	VRAM sharding is the hard constraint; trade comm for VRAM
Large-model fine-tuning (LoRA etc.)	FSDP	sharding frees VRAM for activations/trainable params
Very large + extreme VRAM	FSDP + tensor parallelism	FSDP alone can't split single oversized layers (huge Embedding/FFN); stack TP

Engineering observation: selection isn't "newer is better" — DDP is the most communication-efficient option for small/medium models; FSDP's value is "VRAM sharding unlocks bigger models". Mixed scenarios (EP×TP×DP×FSDP, Chapter 10 MoE) are another layer of complexity; understand the DDP/FSDP communication-VRAM trade-off first, then stack parallelism dimensions.

Distributed Engine Checklist

• [ ] DDP: bucket size matches model scale (lower bucket_cap_mb for small models); Profiler shows no communication gaps in backward
• [ ] DDP: gradient_as_bucket_view=True enabled
• [ ] FSDP: sharding granularity matches model structure (per-layer auto_wrap); mixed precision in the right place
• [ ] FSDP: network bandwidth assessed — traffic ≈ 2× params; confirm Chapter 03 network supports it
• [ ] Large-model scenarios: FSDP + TP stacking decided (who shards, who slices layers)
• [ ] Communication time share included in training logs/Profiler observability

Engineering Takeaways

• Engine mechanisms decide tuning direction: DDP's buckets, FSDP's sharding granularity are all "communication count × wait time" trade-offs — understand the mechanism to know which knob to turn.
• VRAM and communication are two sides of one coin: FSDP trades 2× communication for VRAM; on a weak network it can be slower — do the math before choosing.
• Profiler is the engine's stethoscope: communication gaps and All-Gather waits show up in flame graphs — more effective than guessing.

Next article preview: with the training engine ready, the inference side also needs "measurement first" — the next article, LLM Inference Benchmarking: Throughput, TTFT & Cost in Practice, breaks down a unified measurement method for inference services.