FSDP and DDP: Distributed Training Patterns That Actually Scale

Training a model on one GPU is a notebook exercise. Training the same model on eight GPUs — let alone sixty-four across eight nodes — requires picking a parallelism strategy, configuring a collective communication library, pinning threads, choosing a precision scheme, and paying attention to roughly forty knobs nobody warned you about. The first time a distributed training job hangs silently while every GPU sits at 100% utilization, you start to appreciate that “just add more GPUs” is not, in fact, a strategy.

This post is a tour of the two PyTorch parallelism approaches you will actually use — DDP (DistributedDataParallel) and FSDP (Fully Sharded Data Parallel) — plus the techniques that make them practical: activation checkpointing, mixed precision, gradient accumulation, and NCCL tuning. The goal is a working mental model, not a config dump.

The parallelism landscape in thirty seconds

Before anything else, know which axis you’re scaling on:

• Data parallel — every GPU has a full copy of the model; each sees a different slice of the batch; gradients are averaged across GPUs each step. Works as long as the model fits on one GPU.
• Tensor parallel — a single tensor is split across GPUs. An MLP’s weight matrix lives half on GPU 0 and half on GPU 1. Reduces memory per GPU but requires high-bandwidth interconnect (NVLink) to communicate activations every layer.
• Pipeline parallel — different layers live on different GPUs; batches flow through as micro-batches. Reduces memory, adds pipeline bubbles.
• Sequence parallel — a specialization relevant to transformers where activations are sharded along the sequence dimension.
• Fully sharded data parallel (FSDP/ZeRO) — a data-parallel variant that also shards parameters, gradients, and optimizer state across GPUs, gathering them on demand.

DDP and FSDP are the two you reach for first. Tensor and pipeline parallelism show up for truly enormous models (hundreds of billions of parameters) where the memory math doesn’t work with FSDP alone. Libraries like DeepSpeed, Megatron-LM, and PyTorch’s torch.distributed.pipelining handle those; many teams ship fine without ever touching them.

DDP: the simple, fast default

DistributedDataParallel is the baseline. Every rank has an identical model. Each step:

1. Forward pass on local batch.
2. Backward pass computing local gradients.
3. An all-reduce across ranks averages gradients.
4. Optimizer updates local parameters identically on every rank (because gradients are identical after all-reduce).

The critical trick is that the all-reduce happens overlapped with backward. PyTorch’s autograd engine fires hooks as gradients become available layer-by-layer. DDP bucket-packs those gradients and kicks off all-reduce asynchronously while the rest of backward continues. On a well-configured system, the communication is mostly hidden behind the computation.

A minimal DDP launch:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

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

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

model = build_model().to(local_rank)
model = DDP(model, device_ids=[local_rank])

for batch in loader:
    optimizer.zero_grad(set_to_none=True)
    loss = model(batch).loss
    loss.backward()
    optimizer.step()

Launched with torchrun --nproc_per_node=8 train.py (single node) or torchrun --nnodes=4 --nproc_per_node=8 --rdzv_backend=c10d --rdzv_endpoint=head:29500 train.py (multi-node). torchrun handles the rendezvous and sets the environment variables DDP reads.

The fine print:

• device_ids=[local_rank] — essential. Without it, DDP guesses wrong on multi-GPU nodes.
• set_to_none=True in zero_grad — sets gradients to None rather than zeroing tensors. Faster and better at not holding stale memory.
• find_unused_parameters=False (the default in recent PyTorch) — if your forward doesn’t use every parameter every step, flipping this to True lets DDP still sync gradients but costs overhead. Prefer fixing your model to always use all params.
• Distributed sampler — use torch.utils.data.DistributedSampler so each rank sees a disjoint shard of the dataset. Call sampler.set_epoch(epoch) every epoch or every rank will iterate the same order.

DDP works wonderfully when the model fits in GPU memory. The moment it doesn’t, you have choices.

FSDP: when the model doesn’t fit

FSDP — inspired by DeepSpeed’s ZeRO — is the answer to “my model is too big for one GPU.” It shards three things across ranks:

• Parameters: each GPU holds 1/N of the weights.
• Gradients: each GPU holds 1/N of the gradients.
• Optimizer state: each GPU holds 1/N of the optimizer state (Adam’s first and second moments, typically 2× the parameter memory in fp32).

For an Adam-optimized fp16 training run, the memory breakdown is roughly:

• Parameters: 2 bytes × P
• Gradients: 2 bytes × P
• Optimizer state (fp32 Adam m, v, and fp32 master copy): ~12 bytes × P
• Activations: depend on batch × sequence × hidden

Without any sharding, per-GPU cost is ~16 bytes × P just for model state, before activations. With FSDP on N GPUs, it drops to ~16P/N. A 70B parameter model with N=8 goes from ~1.1 TB per GPU (unfittable) to ~140 GB (fits on an H100 with some activation memory to spare).

The execution model:

1. Before a forward pass of a layer, all-gather the parameters from all ranks so every GPU has the full weight matrix.
2. Run the forward pass.
3. Free the gathered parameters — each GPU keeps only its shard.
4. In backward, all-gather again for the gradient computation.
5. After computing gradients for this layer, do a reduce-scatter — averaging the gradients and distributing the shards back to their owning ranks.
6. Free the gathered gradients.
7. Optimizer updates only the local shard.

Notice the extra communication compared to DDP: FSDP does all-gather + reduce-scatter per layer, while DDP does only all-reduce per step. That’s why FSDP is slower per step than DDP — but it unlocks models that wouldn’t otherwise fit. The bandwidth hit is real; NVLink or InfiniBand matters more with FSDP than with DDP.

FSDP2: the reshaped API

The original torch.distributed.fsdp.FullyShardedDataParallel has been joined (and is gradually being superseded) by FSDP2, whose per-parameter sharding API (fully_shard) is simpler and composes better with other parallelism strategies. A minimal FSDP2 wrap:

1
2
3
4
5
6
7
8
9

from torch.distributed._composable.fsdp import fully_shard
from torch.distributed.device_mesh import init_device_mesh

mesh = init_device_mesh("cuda", (world_size,))

model = build_model().to("cuda")
for block in model.transformer_blocks:
    fully_shard(block, mesh=mesh)
fully_shard(model, mesh=mesh)  # root must be wrapped too

The key concept is wrapping policy: which modules get their own FSDP shard boundary. Wrap too coarsely (the whole model in one shard) and you lose memory savings because the full thing gets gathered each forward. Wrap too finely (every Linear) and communication overhead explodes. For transformers, the established pattern is to wrap each transformer block — a Pareto point that most codebases have converged on.

For older FSDP1 code, transformer_auto_wrap_policy handles this automatically:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

model = FSDP(
    model,
    auto_wrap_policy=partial(transformer_auto_wrap_policy, transformer_layer_cls={TransformerBlock}),
    mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=local_rank,
)

Sharding strategy: the memory vs speed dial

FSDP exposes several strategies:

• FULL_SHARD (ZeRO-3 equivalent) — shard params, gradients, and optimizer state. Maximum memory savings.
• SHARD_GRAD_OP (ZeRO-2) — shard gradients and optimizer state only; replicate params. Less memory savings, fewer collectives.
• HYBRID_SHARD — full shard within a node, replicate across nodes. Minimizes inter-node traffic while keeping intra-node memory savings. Excellent for multi-node where NVLink is fast but IB is slower.
• NO_SHARD — equivalent to DDP. Useful for A/B comparisons.

For multi-node training on models that fit in a node’s combined GPU memory, HYBRID_SHARD is frequently the right pick. You pay intra-node communication (NVLink, fast) and avoid cross-node parameter gathering. The tradeoff vanishes when your model doesn’t fit in a single node — then you need FULL_SHARD.

Mixed precision: not optional

Training in full FP32 on modern GPUs is leaving 2–4× throughput on the table. Mixed precision runs most operations in BF16 or FP16 and keeps a FP32 master copy of weights for the optimizer update.

BF16 vs FP16:

• BF16 has FP32’s exponent range with reduced mantissa. Doesn’t need loss scaling. Works out of the box on Ampere+ (A100, H100). Default choice.
• FP16 has more precision but narrow range. Requires dynamic loss scaling to avoid underflow in gradients. Legacy; only pick it on hardware without BF16 support.

With DDP + torch.amp:

1
2
3
4
5
6
7
8

scaler = torch.cuda.amp.GradScaler()  # FP16 only; not needed for BF16

for batch in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(batch).loss
    loss.backward()
    optimizer.step()

With FSDP, precision is configured via MixedPrecision:

1
2
3
4
5

MixedPrecision(
    param_dtype=torch.bfloat16,     # params cast to bf16 during forward/backward
    reduce_dtype=torch.float32,     # gradients reduced in fp32 for numerical stability
    buffer_dtype=torch.bfloat16,
)

The reduce_dtype=float32 is important. Reducing gradients in fp16/bf16 across many ranks can accumulate enough numerical error to hurt convergence; fp32 reduce costs more bandwidth but is reliably accurate. Most teams keep this at fp32.

FP8 is the new frontier, available on H100 via Transformer Engine. Roughly 2× throughput vs BF16 for attention + MLP operations, at the cost of careful per-tensor scaling. Production ready for inference, maturing for training — if you have H100s and a team willing to debug numerics, worth evaluating.

Activation checkpointing: trading compute for memory

Every layer’s activations have to be kept around for the backward pass. For transformer models, activations often dwarf parameters in memory — a 70B parameter model with long sequences can need hundreds of GB of activation memory.

Activation checkpointing (also called gradient checkpointing) discards activations during forward and recomputes them during backward. You pay ~30% more compute for roughly sqrt(L)× activation memory reduction on L layers. For large models, this is the difference between “fits” and “doesn’t.”

1
2
3
4
5
6
7
8

from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper, CheckpointImpl
)

for i, block in enumerate(model.transformer_blocks):
    model.transformer_blocks[i] = checkpoint_wrapper(
        block, checkpoint_impl=CheckpointImpl.NO_REENTRANT,
    )

NO_REENTRANT is the modern API; it composes with FSDP, AMP, and distributed training cleanly. The older reentrant checkpoint has subtle gotchas — don’t use it unless you have to.

Selective activation checkpointing is the current best practice: checkpoint only expensive-to-store activations (attention matrices) rather than every layer. Saves most of the memory at a fraction of the recompute cost. FlashAttention’s fused kernels also eliminate the need to materialize the attention matrix at all, reducing the pressure that makes checkpointing necessary in the first place.

Gradient accumulation: simulating bigger batches

You want a global batch size of 2048 but can only fit 64 per GPU. Solution: run 32 micro-batches, accumulate gradients, then step. The math is identical to a real batch of 2048.

1
2
3
4
5
6
7

for i, batch in enumerate(loader):
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(batch).loss / GRAD_ACCUM_STEPS
    loss.backward()
    if (i + 1) % GRAD_ACCUM_STEPS == 0:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

With DDP there’s a subtlety: you don’t want to all-reduce gradients on every micro-batch — only the last one. Use model.no_sync() as a context manager to skip gradient sync on intermediate steps:

1
2
3
4
5
6
7
8
9

for i, batch in enumerate(loader):
    is_last = (i + 1) % GRAD_ACCUM_STEPS == 0
    context = nullcontext() if is_last else model.no_sync()
    with context:
        loss = model(batch).loss / GRAD_ACCUM_STEPS
        loss.backward()
    if is_last:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

FSDP has equivalent semantics via set_reshard_after_backward(False) on intermediate micro-batches, though the exact API depends on the FSDP version.

NCCL: the communication layer you must tune

NCCL (NVIDIA Collective Communications Library) is the collective communication backend that actually moves bits between GPUs. It implements all-reduce, all-gather, reduce-scatter, broadcast, and friends. Get it wrong and your training is 2–10× slower than it should be.

The environment variables that matter most:

NCCL_DEBUG=INFO                    # verbose logs during startup — essential for debugging
NCCL_IB_DISABLE=0                  # use InfiniBand if present
NCCL_SOCKET_IFNAME=eth0            # TCP interface for out-of-band bootstrap
NCCL_IB_HCA=mlx5                   # InfiniBand device pattern
NCCL_P2P_DISABLE=0                 # use peer-to-peer (NVLink/PCIe) if present
NCCL_NVLS_ENABLE=1                 # NVLink SHARP on H100+ — significant speedup
NCCL_ALGO=Ring,Tree                # let NCCL pick; don't hardcode unless benchmarking
NCCL_MIN_NCHANNELS=4               # minimum communication channels

At the start of your job, check the NCCL banner for topology. You want to see NVLink being used within a node and IB being used across nodes. If NCCL falls back to sockets for inter-node, your throughput is going to be painful and you probably have a networking misconfiguration (wrong HCA, wrong subnet, missing ibstat confirmation).

Common pitfalls:

• Mismatched NCCL versions across containers. NCCL negotiates at startup; mismatched versions sometimes work, sometimes hang with no error.
• MTU wrong on TCP fallback. If your IB isn’t detected, NCCL uses TCP. Default MTU of 1500 is disastrous for multi-GB gradients. Enable jumbo frames (MTU 9000) on your Ethernet.
• PXE boot weirdness. Some clusters have CPU affinity misconfigured so ranks end up on the wrong NUMA node relative to their assigned GPU. Check NCCL_DEBUG=INFO output for the GPU ↔ NIC ↔ NUMA assignment.
• GDRCopy disabled. On systems with GPU Direct RDMA, confirm NCCL_NET_GDR_LEVEL=5 and that the driver is loaded. Otherwise gradients round-trip through CPU memory.

Use nccl-tests to sanity check your hardware before running a full job. all_reduce_perf -b 8 -e 1G -f 2 -g 8 tells you the all-reduce bandwidth your setup actually achieves. Compare to the hardware peak (NVLink 3 = 600 GB/s, InfiniBand NDR = 400 Gbps per link). If you’re getting a fraction of peak, find out why before scaling to 100 nodes.

Debugging: the tooling that works

Distributed training failure modes:

Hangs with no progress. Almost always a collective mismatch — one rank entered a different collective than the others. Common causes: divergent control flow in the model, variable-length batches creating different all-gather shapes, a rank dying silently. Fix: TORCH_DISTRIBUTED_DEBUG=DETAIL, which logs every collective and which ranks entered it.

Slow first iteration, fast after. NCCL buffer warmup. Sometimes also CUDA graph capture. Ignore unless first-iteration cost hurts you (e.g., in benchmarking).

Loss is fine on 1 GPU, diverges on 8. Usually a batch norm issue (per-GPU statistics) or a data loader seeding problem (same data on every rank). Double-check that the sampler seeds per rank and per epoch.

Memory growth each step. CUDA caching allocator fragmentation, or retained references (often to the loss tensor). torch.cuda.memory_summary() per rank at the end of each epoch helps. PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True fixes many fragmentation issues in newer PyTorch.

Works on 8 GPUs, hangs on 16. A rank is timing out waiting for collective completion. Increase the NCCL timeout: torch.distributed.init_process_group(timeout=timedelta(minutes=30)). The default 30 minutes is often too short for first-iteration warmup on large models.

The essential tools:

• py-spy dump --pid <pid> — stack trace a hung Python process without disturbing it.
• nvidia-smi watches — GPU utilization, memory, temperature. If all GPUs are at 100% util but nothing is progressing, you’re deadlocked.
• nsys profile — NVIDIA Nsight Systems traces kernel-level activity. Heavy, but tells you precisely where compute-communication overlap is failing.
• torch.profiler — lighter, in-process. Export Chrome traces and look at the timeline. Gaps between kernels are where you’re losing time.

Scaling practices that matter

Some patterns that separate solid multi-node training from the kind that wakes you up:

Use torchrun, not custom launch scripts. The elastic launcher handles rendezvous, per-rank environment setup, and restart-on-failure. Its --rdzv_backend=c10d uses a TCP-based rendezvous that doesn’t need an external coordinator.

Checkpoint asynchronously. Synchronous checkpointing on a 70B model takes minutes — precious training time. Use torch.distributed.checkpoint with async save. It copies tensors to pinned CPU memory and writes in a background thread.

Save sharded checkpoints. Legacy full-state-dict checkpointing gathers the whole model to rank 0 before saving. OOMs for big models. Sharded checkpointing has each rank save its shard; resharding happens at load time if the world size changes.

Pin your seeds per rank. Set torch.manual_seed(base_seed + rank) and torch.cuda.manual_seed(base_seed + rank). Every rank must have a different seed or data augmentation is identical everywhere.

Limit the num_workers on the dataloader. One worker per rank is often fine; more can oversubscribe CPU cores since every rank on a node competes for the same pool. Pin memory (pin_memory=True) unless profiling shows it’s not helping.

Prefetch on the GPU. torch.cuda.Stream for overlapped H2D copies. Helps when CPU-side data loading is a bottleneck (large image inputs, heavy augmentation).

Profile at low scale first. If DDP is slow on 8 GPUs, it will be slow on 128. Find the bottleneck at small scale where you can iterate in minutes, not hours.

FSDP vs DeepSpeed: which to pick

DeepSpeed is the alternative that popularized ZeRO sharding. The comparison today:

• Ecosystem: FSDP is in PyTorch core. No extra library to maintain, no version matching, clean interop with other PyTorch tools.
• Features at the bleeding edge: DeepSpeed has had features like ZeRO-Infinity (CPU/NVMe offloading) for years. FSDP now supports CPU offload; NVMe offload is less mature.
• Extensions: DeepSpeed ships 3D parallelism (data + tensor + pipeline) in one package. For models requiring all three, it’s still simpler. Megatron-LM + FSDP composition is catching up.
• Developer experience: FSDP feels more “PyTorch-native.” DeepSpeed feels like a framework.

For most teams in 2026 starting fresh, FSDP (specifically FSDP2 where supported) is the default. Reach for DeepSpeed or Megatron-LM when you’re training genuinely enormous models or need features FSDP doesn’t have yet.

A pragmatic rollout

For a team starting distributed training:

1. Get single-GPU training working. Don’t distribute what you haven’t verified.
2. Add DDP on a single multi-GPU node. Benchmark. If you’re getting less than ~90% of single-GPU throughput per GPU, something is wrong — fix it before scaling further.
3. Add mixed precision (BF16). Free 2× speedup on Ampere+.
4. If the model fits, stay on DDP. Simpler is better.
5. If the model doesn’t fit, add FSDP with transformer auto-wrap and activation checkpointing. Measure memory and throughput.
6. If you’re multi-node and the model fits per-node, switch to HYBRID_SHARD.
7. Tune NCCL with nccl-tests before claiming scaling works.
8. Add sharded checkpointing and async save before you lose a 12-hour run.
9. Profile with Nsight on a representative step. Look for gaps, not just utilization.
10. Only reach for tensor parallel or pipeline parallel after FSDP + activation checkpointing + HYBRID_SHARD stops working. Most teams never need to.

Distributed training is unforgiving but not mysterious. The mental model stays stable across years — collectives, overlap, sharding, precision — even as the specific APIs churn. A few weeks of discomfort learning it pays for itself every time a model is too big for one GPU, which is increasingly the default rather than the exception.