Distributed Training Basics — DDP and FSDP

Start with a memory failure, not a parameter threshold

A training step must keep more than the inference weights alive. When a model fits for generation but fails at the first optimizer step, the new allocation may be Adam's moments rather than an attention activation. When it fails only at long context, sharding the optimizer may barely help. Distributed training is a choice about which state is replicated, which work is partitioned, and when communication is paid.

This lesson assumes a working single-device training loop, tensor shapes, mixed precision and gradient accumulation. You will derive the DDP/ZeRO memory table, trace collectives, convert a complete small loop to a rank-aware launch, distinguish FSDP1 from FSDP2, and design a measurement and recovery plan. The CPU arithmetic and single-process gradient oracle are exercised in this review. The CUDA/distributed procedures are complete reference examples for PyTorch 2.9.0, not claimed GPU runs. There is no requirement to rent hardware to study them.

Memory accounting: derive the table from bytes

Let P be the number of parameters and R the number of data-parallel ranks. For one illustrative mixed-precision Adam representation, count BF16 parameters and gradients at 2P bytes each, plus FP32 master weights and two FP32 moments at 12P bytes. That is 16P bytes of persistent model state, not 16P bytes of peak GPU memory. Native autocast with FP32 stored parameters, optimizer implementations, master-copy policies and gradient dtypes can produce a different inventory.

All GB here are decimal. At four ranks the last three rows become 49, 38.5 and 28 GB. The ZeRO paper, v3 develops this staged removal of replicated state; PyTorch FSDP and DeepSpeed ZeRO-3 share the fully sharded idea, not identical APIs or all runtime choices.

Peak memory also includes activations, an unsharded unit, prefetched units, reduction buffers, kernel workspace, CUDA context and allocator reserve. A model-wide FSDP unit can materialize nearly all weights at once even though optimizer state remains sharded: this is a poor peak-memory choice, not literally DDP. Block-sized units trade lower gather peaks for more collective launches. Activation checkpointing discards chosen forward intermediates and recomputes them in backward; it does not shrink the persistent 16P term or promise a universal fourfold reduction.

For 50B parameters on eight devices, persistent state alone is 100 GB per device under this model. For 70B it is 140 GB. Neither fits an 80 GB device merely by turning on activation checkpointing. The old exercise that claimed 50–70B fit on eight 80 GB devices under these assumptions was arithmetically impossible. More devices, a different optimizer/state representation, offload or a genuinely different parallelization must change the budget.

# CPU arithmetic; no process group or device allocation.
def state_bytes(params, ranks, mode):
    if params <= 0 or ranks < 1:
        raise ValueError("positive parameter count and rank count required")
    factors = {"single": 16., "ddp": 16., "zero1": 4 + 12/ranks,
               "zero2": 2 + 14/ranks, "fsdp": 16/ranks}
    return params * factors[mode]
 
def peak_estimate(params, ranks, mode, activation_bytes, gather_bytes,
                  other_bytes):
    if min(activation_bytes, gather_bytes, other_bytes) < 0:
        raise ValueError("memory terms cannot be negative")
    return (state_bytes(params, ranks, mode) + activation_bytes
            + gather_bytes + other_bytes)
 
assert state_bytes(7e9, 8, "fsdp") == 14e9
assert state_bytes(50e9, 8, "fsdp") > 80e9

For an eight-device 40 GB exercise, the zero-activation, zero-overhead upper bounds are 2.5B parameters for DDP, 7.27B for ZeRO-1, 10.67B for ZeRO-2 and 20B for fully sharded state. They are not recommended capacities. If a measured activation term is 8 GB and checkpointing reduces it to 2 GB, subtract the observed 2 GB plus other peaks before solving for P; do not invent that ratio in the estimator.

DDP: replicated state, partitioned examples

DDP in PyTorch 2.9 synchronizes gradients, but does not partition your inputs. Each rank has a complete model and optimizer, performs forward/backward on its own examples, and updates using reduced gradients. The same initial state, consistent optimizer steps and compatible collective order keep replicas aligned. Buffers such as BatchNorm statistics have separate behavior; BatchNorm is local unless using an appropriate synchronized implementation. Most decoder LLMs instead use LayerNorm/RMSNorm, but that does not make DDP's buffer contract irrelevant to other models.

If every rank's loss is a mean with equal denominators, averaging rank gradients gives the global mean. With two ranks holding 2 and 6 valid tokens and local mean losses 3 and 1, the correct global loss is (2*3+6*1)/8=1.5, not 2. For a rank-local sum S and a globally reduced valid-token count T, backpropagate R*S/T when DDP averages rank gradients. A zero-token rank contributes a differentiable zero and still follows the collective schedule. Never divide by zero or silently skip backward on just one rank.

A DistributedSampler partitions indices, not all imaginable token-level objectives. Call set_epoch(epoch) before each shuffled epoch. Its padding or dropping policy changes which examples appear; measure evaluation sums and counts separately rather than averaging unequal-rank accuracies. Uneven iteration counts can hang collectives unless a supported uneven-input protocol is deliberately used.

A complete rank-aware DDP reference

Save as train_ddp.py in an independently approved CUDA environment with PyTorch 2.9.0. This uses synthetic tensors, a small ordinary network and equal-sized microbatches, so it downloads no model or data. Every symbol is defined; it is a replacement for the misleading “four-line conversion.” The launch and program are not executed as distributed code here.

# Reference only: PyTorch 2.9.0, NCCL, one process per visible CUDA GPU.
import os
from contextlib import nullcontext
import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset, DistributedSampler
 
def main():
    distributed = "RANK" in os.environ
    local = int(os.environ.get("LOCAL_RANK", "0"))
    torch.cuda.set_device(local)
    device = torch.device("cuda", local)
    if distributed:
        dist.init_process_group("nccl")
    try:
        rank = dist.get_rank() if distributed else 0
        torch.manual_seed(7)  # identical model/data construction on every rank
        data = TensorDataset(torch.randn(128, 8), torch.randn(128, 2))
        sampler = DistributedSampler(data, shuffle=True) if distributed else None
        loader = DataLoader(data, batch_size=4, sampler=sampler,
                            shuffle=sampler is None, drop_last=True)
        model = nn.Sequential(nn.Linear(8, 16), nn.Tanh(), nn.Linear(16, 2)).to(device)
        if distributed:
            model = DDP(model, device_ids=[local])
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
        accumulation = 3
        for epoch in range(2):
            if sampler is not None:
                sampler.set_epoch(epoch)
            optimizer.zero_grad(set_to_none=True)
            for i, (x, y) in enumerate(loader):
                # Correctly normalize the final, possibly shorter accumulation window.
                begin = (i // accumulation) * accumulation
                window = min(accumulation, len(loader) - begin)
                last = i + 1 == begin + window
                context = model.no_sync() if distributed and not last else nullcontext()
                with context:  # forward MUST be inside no_sync too
                    loss = nn.functional.mse_loss(model(x.to(device)), y.to(device)) / window
                    loss.backward()
                if last:
                    nn.utils.clip_grad_norm_(model.parameters(), 1.)
                    optimizer.step()
                    optimizer.zero_grad(set_to_none=True)
            if rank == 0:
                bare = model.module if distributed else model
                tmp, final = "ddp-checkpoint.tmp", "ddp-checkpoint.pt"
                torch.save({"model": bare.state_dict(), "optimizer": optimizer.state_dict(),
                            "next_epoch": epoch + 1}, tmp)
                os.replace(tmp, final)
            if distributed:
                dist.barrier()
    finally:
        if distributed:
            dist.destroy_process_group()
 
if __name__ == "__main__":
    main()

# Reference launch only, not run in this review.
torchrun --standalone --nproc_per_node=2 train_ddp.py
# The same file also has a one-GPU path:
python train_ddp.py

This epoch-boundary example saves model/optimizer state, not a production exact-resume bundle. For arbitrary stochastic models and mid-epoch continuation also save per-rank CPU/CUDA RNG, sampler/data position, scheduler/scaler state, effective-update count and configuration hashes. Only rank zero writes a replicated checkpoint; sharded checkpointing is a different protocol. A barrier does not rescue a crashed rank or make a remote object-store write atomic.

Accumulation, clipping and scaling measurements

Effective example batch is microbatch × accumulation × data_parallel_ranks. Accumulation alone adds no network on one GPU but takes more sequential compute. In DDP, naive accumulation still reduces on every backward; no_sync defers those reductions until the final backward. Keep clipping, optimizer, scheduler and AMP scaler updates at the effective-update boundary. FSDP has different synchronization/memory behavior; do not transplant DDP's API by name.

To compare 1, 2 and 4 devices, specify strong scaling (fixed global batch) or weak scaling (fixed local batch). Warm up, synchronize device timing, measure a window of full updates including optimizer steps, aggregate the total processed tokens, and report both throughput and peak allocated/reserved memory per rank. Exclude compilation/setup only if separately reported. Compare loss/gradients at a matched objective before claiming speedup. Expected speedup is measured T1/TR; efficiency is T1/(R*TR). No 90–95% scaling, 15% overhead or two-device 1.85× result follows from parameter count alone.

FSDP: materialize only what is needed

At a fully sharded unit, ranks all-gather parameter shards before using the layer. With reshard-after-forward enabled they free the full copies, gather again before backward, then reduce-scatter gradients. Optimizer steps operate on the owned shards. Each rank still processes a data slice; FSDP is sharded data parallelism, not tensor parallelism that partitions the layer's matmul itself.

FSDP1 and FSDP2: explicitly different versioned paths

For the FSDP1 wrapper in PyTorch 2.9, a transformer auto-wrap policy can name your actual block class; construct the optimizer after wrapping. Inspect that each desired block became a unit, select MixedPrecision deliberately and use the documented state-dict/checkpoint path. FSDP1 flattening and FSDP2 per-parameter DTensors have different state representations; do not infer a release's stability from an old blog date.

The following FSDP2 reference replaces the wrapper with the tagged 2.9 fully_shard API. It defines a tiny block rather than relying on an undeclared LlamaLikeBlock. It performs a full update and uses Distributed Checkpoint (DCP) for model and optimizer state. All ranks participate in save/load. This is source-reviewed but unexecuted CUDA/distributed reference code.

# Reference only: PyTorch 2.9.0, launched with torchrun as above.
import os
import torch
from torch import nn
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
 
class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(8, 8)
    def forward(self, x):
        return x + torch.tanh(self.proj(x))
 
def main():
    local = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local)
    dist.init_process_group("nccl")
    try:
        torch.manual_seed(9)
        device = torch.device("cuda", local)
        mesh = init_device_mesh("cuda", (dist.get_world_size(),))
        model = nn.Sequential(Block(), Block(), nn.Linear(8, 2)).to(device)
        policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16,
                                      reduce_dtype=torch.float32)
        for block in model:
            fully_shard(block, mesh=mesh, mp_policy=policy, reshard_after_forward=True)
        fully_shard(model, mesh=mesh, mp_policy=policy, reshard_after_forward=True)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
        # Distinct data per rank, identical initial parameters.
        torch.manual_seed(100 + dist.get_rank())
        x, y = torch.randn(4, 8, device=device), torch.randn(4, 2, device=device)
        optimizer.zero_grad(set_to_none=True)
        loss = (model(x).float() - y).square().mean()
        loss.backward()
        optimizer.step()
        ms, os_ = get_state_dict(model, optimizer)
        state = {"model": ms, "optimizer": os_}
        dcp.save(state, checkpoint_id="fsdp2-checkpoint")
        dcp.load(state, checkpoint_id="fsdp2-checkpoint")
        set_state_dict(model, optimizer, model_state_dict=state["model"],
                       optim_state_dict=state["optimizer"])
    finally:
        dist.destroy_process_group()
 
if __name__ == "__main__":
    main()

The fully_shard declaration specifies dim-0 DTensor sharding and hook-driven collectives; bottom-up calls form separate communication groups. Mixed precision is configured for each application/group, not a promise that every tensor can independently use any format in one call. A sharded save/load smoke test is not exact continuation: the stronger assessment compares the next batch, loss and update after restart, including RNG and data state.

DTensor, topology and framework alternatives

A DeviceMesh assigns ranks to one or more axes. A DTensor placement describes storage along a mesh axis: Shard(dim) owns a piece, Replicate() owns a full copy, and Partial() owns a contribution awaiting reduction. Redistributing Partial to Replicate can invoke an all-reduce; to Shard can invoke reduce-scatter. These names describe semantics, not free metadata operations.

For a 128-device design with 16 nodes of 8 devices, an initial candidate is TP=8 within a node and data-parallel/sharding groups across 16 nodes. TP exchanges fine-grained activation information on the fast local fabric; FSDP groups operate on already tensor-partitioned parameters. Global token batch is 16 × local_sequences × sequence_length × accumulation, not 128 independent data replicas. Pipeline stages may help capacity but add schedule bubbles; expert parallelism adds routed-token all-to-all. Profile and verify supported composition; not every PyTorch distributed feature is built on DTensor, and no parameter threshold requires this exact layout.

Composer v0.29.0 provides training/checkpoint machinery; autoresume requires a save folder and run name and searches particular saved checkpoint locations. LLM Foundry v0.18.0 supplies data-preparation, training, evaluation and benchmarking scripts around Composer. Inspect the selected model/wrapping, state restoration, data-order guarantees and supported precision in the pinned configuration. Neither name guarantees automatic loss-spike repair, deterministic mid-epoch recovery, a particular DBRX/MPT training cost, or “free” operational safety. Hand-written small loops teach mechanisms; frameworks reduce integration work only after their actual contracts are checked.

Communication and a roofline sanity check

For a ring with R ranks and a tensor payload B bytes, per-rank sent volume is approximately 2(R-1)B/R for all-reduce and (R-1)B/R for all-gather or reduce-scatter. Count directions consistently; do not compare sent bytes with an aggregate bidirectional bandwidth. A BF16 7B payload is 14 GB. At R=8, DDP's same-dtype gradient all-reduce is 24.5 GB; two parameter gathers plus one same-dtype reduce-scatter are 36.75 GB, a 1.5× volume ratio in this simplified schedule. FP32 gradients, non-resharding and prefetch alter it.

At an explicitly hypothetical effective unidirectional bandwidth of 100 GB/s, those serialized transfer times are .245 and .3675 seconds, before latency and contention. They are not observed step overhead: overlap can hide part of communication, while dependencies and many small messages can expose it. NVLink is an intra-node link family; InfiniBand is one inter-node fabric. A model name alone does not say whether the rented topology includes NVSwitch or has sufficient NIC bandwidth per rank. PCIe generation, lanes, NUMA placement and shared NICs matter. GPU marketing aggregate bandwidth is not a measured all-gather rate. Activation checkpointing adds recomputation and can change the gather schedule; it does not inherently reduce collective count.

Failure investigation and recovery, not invented war stories

An oversized wrapping unit: inspect wrappers, compare persistent state against gather-time peaks, then change unit boundaries. An unused-parameter path: reproduce the actual graph, configure DDP's unused/static-graph options for that graph, and benchmark; do not blindly disable a needed option. A collective hang: compare rank progress, data counts, collective order and failures before blaming the fabric. A NaN: record pre-reduction per-rank finite checks and gradient norms so reduction does not conceal origin. Do not automatically skip a suspicious batch or lower LR without recording the change and investigating data, precision and hardware.

Public BLOOM/OPT training chronicles are useful case-study resources, but the original lesson's exact failure counts, temperatures, restart totals and first-person incidents were not established by the cited passages. The restored learning task is a concrete incident log: timestamp/rank/step, last complete checkpoint, first anomalous operation, reversible diagnosis, and post-restart parity. It does not impersonate an observed multi-month training run.

Practice and worked answers

1. Explain DDP versus FSDP. DDP replicates state, partitions examples through your input pipeline, and reduces gradients. Fully sharded training distributes state, reconstructs units as needed and reduces gradients to shards. Both can process different examples; their memory/communication schedules differ.
2. Name the ZeRO stages and budget a 7B model. Optimizer; then gradients; then parameters. The eight-rank state values are 38.5, 26.25 and 14 GB versus DDP's 112 GB. Add activations and all temporary peaks before claiming fit.
3. What four changes start a DDP port, and what did the slogan omit? Initialize group, assign device, wrap model and partition input. It omitted epoch seeding, objective normalization, accumulation boundaries, rank-aware I/O, process cleanup and recovery.
4. What if rank one has no valid tokens? Use a differentiable zero local loss numerator, globally sum counts, and preserve required collectives. If the entire global count is zero, all ranks agree to reject the update; no rank independently exits.
5. Why more GPUs sometimes fail to fix OOM? Fixed local batch/context keeps local activations approximately unchanged, while state shrinks. Adding ranks can still help total memory; it is false that it can never fix an activation-heavy OOM. Measure the dominant term.
6. Solve the 128-device stretch. Use the 16×8 candidate above, compute shard and activation peaks, benchmark collectives at target payloads, then compare TP/PP alternatives at fixed global tokens. No 30% communication threshold is a universal law.
7. Design the scaling experiment. Record 1/2/4-rank full-update times after warmup, world-size-aware tokens/sec and per-rank peaks. Compare gradients with a one-process concatenated-batch oracle. Publish no speedup until these measurements exist.
8. FSDP2 in one sentence? In the tagged 2.9 implementation, fully_shard manages DTensor parameter shards and grouped hook-driven collectives rather than the FSDP1 FlatParameter wrapper representation.

Spaced review: tomorrow reconstruct the four-row byte table and trace a gather/reduce-scatter schedule without notes; on day seven explain unequal-token normalization and resume requirements. Bring back the 16-byte assumption, the distinction between state and activation memory, and the actual torchrun launch contract—not an unmeasured price or scaling factor. Next: choosing and operating a rented GPU turns these budgets into a reference operating procedure; it does not imply a resource has been provisioned.

Previous: ← DL S050 · Next: DL S052 →

Original optional resource links

The original optional video/reading destinations are retained here for continuity, not as current API, price, benchmark or publication-date evidence. The versioned primary citations in the teaching above govern the corrected claims. These supplementary destinations were not all fetched or watched; a missing or changed resource does not replace the local derivations, procedures and answers.

• https://www.youtube.com/watch?v=toUSzwR0EV8
• https://arxiv.org/abs/1910.02054
• https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
• https://engineering.fb.com/2021/07/15/open-source/fsdp/
• https://pytorch.org/blog/pytorch2-4/
• https://github.com/pytorch/torchtitan
• https://github.com/pytorch/pytorch/issues/88838
• https://github.com/mosaicml/llm-foundry
• https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm
• https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md
• https://github.com/facebookresearch/metaseq/tree/main/projects/OPT/chronicles.md