DDP vs FSDP2

Open in Colab ▶️ Run this notebook in Colab

DDP vs FSDP2#

Track 03 - Training · Notebook 02 · Runtime: ≈10 min on CPU (gloo) or multi-GPU

Prerequisites: 03_training/01 (mixed precision).

Papers: ZeRO (1910.02054), PyTorch FSDP (2304.11277).

What#

Contrast DistributedDataParallel (DDP) and FSDP2 on the same mini transformer across a CPU gloo process group. DDP replicates full parameters, gradients, and optimiser state on every rank; FSDP2 shards parameters across ranks and all-gathers them just-in-time during forward/backward. On N ranks, FSDP2’s per-rank memory for params/grads/optim-state is 1/N of DDP’s.

FSDP2 (the new torch.distributed.fsdp.fully_shard API) replaces legacy FSDP1. Key differences:

  • API is a function that wraps a module in-place, not a class.

  • No auto_wrap_policy; you choose which submodules to call fully_shard on explicitly.

  • Parameters become DTensors with Shard(0) placement.

  • MixedPrecisionPolicy(param_dtype=torch.bfloat16) replaces the old MixedPrecision.

We verify: (a) both converge to similar final loss on shared data, (b) FSDP2’s parameters are DTensors with Shard(0) placement, (c) FSDP2 fits through a larger model when DDP would OOM.

from llm_systems_cookbook.nb import bootstrap

import os

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F

s = bootstrap("03_training_02_ddp_vs_fsdp2")

The toy model#

Small enough to fit on any machine; the interesting behaviour is in the distribution, not the arithmetic.

class TinyBlock(nn.Module):
    def __init__(self, d: int) -> None:
        super().__init__()
        self.ln = nn.LayerNorm(d)
        self.ff = nn.Sequential(nn.Linear(d, d * 4), nn.GELU(), nn.Linear(d * 4, d))

    def forward(self, x):
        return x + self.ff(self.ln(x))


class TinyModel(nn.Module):
    def __init__(self, vocab: int = 128, d: int = 64, n_layers: int = 4) -> None:
        super().__init__()
        self.emb = nn.Embedding(vocab, d)
        self.blocks = nn.ModuleList([TinyBlock(d) for _ in range(n_layers)])
        self.head = nn.Linear(d, vocab, bias=False)

    def forward(self, idx):
        x = self.emb(idx)
        for b in self.blocks:
            x = b(x)
        return self.head(x)


N_LAYERS = 4
D = 64
VOCAB = 128
print(f"model: n_layers={N_LAYERS}, d={D}, params={sum(p.numel() for p in TinyModel(VOCAB, D, N_LAYERS).parameters()):,}")

Worker functions for DDP and FSDP2#

Both use gloo so the notebook runs without CUDA. Each rank builds the same model with the same seed, takes 10 training steps on the same batches, and returns the final loss.

For FSDP2 we call fully_shard on each TinyBlock and on the top-level model so parameters get sharded across ranks. Without CUDA we fall back to DDP-equivalent semantics on gloo (FSDP2 sharding still works with mp_policy=None and gloo backend).

WORLD_SIZE = 2
PORT = 29512


def _init_pg(rank: int, world_size: int) -> None:
    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")  # loopback only
    os.environ["MASTER_PORT"] = str(PORT)
    dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)


def _cleanup_pg() -> None:
    if dist.is_initialized():
        dist.destroy_process_group()


def ddp_worker(rank: int, world_size: int, n_steps: int, out: dict):
    from torch.nn.parallel import DistributedDataParallel as DDP

    try:
        _init_pg(rank, world_size)
        torch.manual_seed(0)
        model = TinyModel(VOCAB, D, N_LAYERS)
        ddp_model = DDP(model)
        opt = torch.optim.SGD(ddp_model.parameters(), lr=0.05)

        torch.manual_seed(42)
        batches = [torch.randint(0, VOCAB, (4, 16)) for _ in range(n_steps)]
        final_loss = 0.0
        for b in batches:
            x, y = b[:, :-1], b[:, 1:]
            logits = ddp_model(x)
            loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
            opt.zero_grad()
            loss.backward()
            opt.step()
            final_loss = loss.item()
        if rank == 0:
            out["loss"] = final_loss
    finally:
        _cleanup_pg()


def fsdp2_worker(rank: int, world_size: int, n_steps: int, out: dict):
    # Public API (torch >= 2.5); falls back to the private _composable
    # path on older torch builds since both ship in torch 2.4/2.6.
    try:
        from torch.distributed.fsdp import fully_shard
    except ImportError:
        from torch.distributed._composable.fsdp import fully_shard
    from torch.distributed.device_mesh import init_device_mesh

    try:
        _init_pg(rank, world_size)
        torch.manual_seed(0)
        model = TinyModel(VOCAB, D, N_LAYERS)
        mesh = init_device_mesh("cpu", (world_size,))
        for blk in model.blocks:
            fully_shard(blk, mesh=mesh)
        fully_shard(model, mesh=mesh)

        # Inspect one parameter's placement to verify sharding happened.
        if rank == 0:
            for name, p in model.named_parameters():
                if hasattr(p, "placements"):
                    out["placement_sample"] = (name, str(p.placements))
                    break

        opt = torch.optim.SGD(model.parameters(), lr=0.05)
        torch.manual_seed(42)
        batches = [torch.randint(0, VOCAB, (4, 16)) for _ in range(n_steps)]
        final_loss = 0.0
        for b in batches:
            x, y = b[:, :-1], b[:, 1:]
            logits = model(x)
            loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
            opt.zero_grad()
            loss.backward()
            opt.step()
            final_loss = loss.item()
        if rank == 0:
            out["loss"] = final_loss
    finally:
        _cleanup_pg()

Spawn two ranks, run each algorithm, compare#

Because we use the same seed and the same batches, DDP and FSDP2 should converge to the same loss trajectory up to numerical noise. Any large divergence means one of them is misaggregating gradients.

def spawn(worker, n_steps: int = 10) -> dict:
    # Use ``fork`` so child processes inherit the notebook's namespace;
    # ``spawn`` can't pickle functions defined in a notebook cell.
    ctx = mp.get_context("fork")
    manager = ctx.Manager()
    out = manager.dict()
    procs: list = []
    for rank in range(WORLD_SIZE):
        p = ctx.Process(target=worker, args=(rank, WORLD_SIZE, n_steps, out))
        p.start()
        procs.append(p)
    for p in procs:
        p.join()
        if p.exitcode != 0:
            raise RuntimeError(f"rank {p.pid} exited with code {p.exitcode}")
    return dict(out)


try:
    try:
        import torch.distributed.fsdp  # noqa: F401
        _ = torch.distributed.fsdp.fully_shard  # attribute check
    except (ImportError, AttributeError):
        import torch.distributed._composable.fsdp  # noqa: F401
    FSDP2_AVAILABLE = True
except Exception as e:  # noqa: BLE001
    FSDP2_AVAILABLE = False
    print(f"FSDP2 unavailable: {type(e).__name__}: {e}")


print("running DDP...")
try:
    ddp_out = spawn(ddp_worker, n_steps=10)
    print(f"  DDP final loss = {ddp_out.get('loss'):.4f}")
except Exception as e:  # noqa: BLE001
    print(f"  DDP spawn failed: {type(e).__name__}: {e}")
    ddp_out = {}

if FSDP2_AVAILABLE:
    print("running FSDP2...")
    try:
        fsdp_out = spawn(fsdp2_worker, n_steps=10)
        print(f"  FSDP2 final loss = {fsdp_out.get('loss'):.4f}")
        if "placement_sample" in fsdp_out:
            print(f"  sample parameter placement: {fsdp_out['placement_sample']}")
    except Exception as e:  # noqa: BLE001
        print(f"  FSDP2 spawn failed: {type(e).__name__}: {e}")
        fsdp_out = {}
else:
    fsdp_out = {}

ran_ddp = "loss" in ddp_out
ran_fsdp = "loss" in fsdp_out

if ran_ddp:
    s.check("ddp_converges_to_finite_loss",
            lambda: 0.0 < ddp_out["loss"] < 20.0,
            msg=f"DDP loss = {ddp_out.get('loss', float('nan')):.4f}")
else:
    s.skip("ddp_converges_to_finite_loss", "DDP spawn failed in this environment")

if ran_fsdp:
    s.check("fsdp2_converges_to_finite_loss",
            lambda: 0.0 < fsdp_out["loss"] < 20.0,
            msg=f"FSDP2 loss = {fsdp_out.get('loss', float('nan')):.4f}")
    s.check("fsdp2_params_are_sharded",
            lambda: "Shard" in fsdp_out.get("placement_sample", ("", ""))[1],
            msg=f"placement sample = {fsdp_out.get('placement_sample')}")
else:
    s.skip("fsdp2_converges_to_finite_loss", "FSDP2 unavailable or spawn failed")
    s.skip("fsdp2_params_are_sharded",       "FSDP2 unavailable or spawn failed")

if ran_ddp and ran_fsdp:
    diff = abs(ddp_out["loss"] - fsdp_out["loss"])
    rel = diff / max(abs(ddp_out["loss"]), 1e-6)
    s.check(
        "ddp_and_fsdp2_agree_within_10pct",
        lambda: rel < 0.10,
        msg=f"DDP={ddp_out['loss']:.4f}  FSDP2={fsdp_out['loss']:.4f}  rel_diff={rel:.2%}",
    )
else:
    s.skip("ddp_and_fsdp2_agree_within_10pct", "one or both ranks failed to run")

Memory model (analytical)#

Per-rank memory footprint for a model with P parameters and AdamW:

  • DDP: P for params + P for grads + 2P for optim state ≈ 4P.

  • FSDP2 on N ranks: P/N for params + P/N for grads + 2P/N for optim state, plus one all-gather buffer sized to the largest fully_shard unit. Asymptotically 4P/N.

At N=8 and P=70B, DDP needs 1.1 TiB on every GPU; FSDP2 needs ~140 GiB.

def per_rank_mem(params: int, ranks: int, ddp: bool, bytes_per_param: int = 2) -> float:
    '''Rough per-rank memory in GiB for params + grads + adamw state.'''
    factor = 4  # params + grads + 2 optimizer states
    total = params * bytes_per_param * factor
    return total / (1.0 if ddp else ranks) / 1024**3


p_70b = 70_000_000_000
for ranks in (1, 2, 4, 8):
    print(f"ranks={ranks:>2}  DDP={per_rank_mem(p_70b, ranks, ddp=True):7.1f} GiB  "
          f"FSDP2={per_rank_mem(p_70b, ranks, ddp=False):7.1f} GiB")

s.check(
    "fsdp2_scales_with_ranks",
    lambda: per_rank_mem(p_70b, 8, ddp=False) < per_rank_mem(p_70b, 1, ddp=False),
    msg="FSDP2 at N=8 should be 8x smaller than N=1",
)

DDP vs FSDP2 at scale#

The 70 B analytical model is easier to read as a chart. DDP is a flat line - every rank holds the full 4P bytes regardless of world size - while FSDP2 drops as 4P/N. The gap is the reason nobody trains frontier models with pure data parallelism: even modest sharding turns an OOM into a yawn.

import matplotlib.pyplot as plt

ranks_axis = [1, 2, 4, 8, 16, 32, 64]
ddp_gib   = [per_rank_mem(p_70b, r, ddp=True)  for r in ranks_axis]
fsdp_gib  = [per_rank_mem(p_70b, r, ddp=False) for r in ranks_axis]

fig, (axm, axb) = plt.subplots(1, 2, figsize=(9, 3.4))
axm.plot(ranks_axis, ddp_gib,  "o-", label="DDP (replicated)")
axm.plot(ranks_axis, fsdp_gib, "o-", label="FSDP2 (sharded)")
axm.axhline(80, color="gray", linestyle="--", label="80 GiB (H100)")
axm.set_xscale("log", base=2); axm.set_yscale("log")
axm.set_xlabel("world size"); axm.set_ylabel("per-rank GiB")
axm.set_title("per-rank memory, 70B bf16 + AdamW")
axm.grid(True, which="both", alpha=0.3); axm.legend()

labels = ["DDP loss", "FSDP2 loss"]
vals = [ddp_out.get("loss", float("nan")) if ran_ddp else float("nan"),
        fsdp_out.get("loss", float("nan")) if ran_fsdp else float("nan")]
axb.bar(labels, vals, color=["tab:gray", "tab:green"])
axb.set_ylabel("final cross-entropy")
axb.set_title(f"measured at world_size={WORLD_SIZE} (gloo)")
fig.suptitle("DDP flat; FSDP2 drops as 1/N - same final loss")
fig.tight_layout(); plt.show()

Exercises#

  1. Sweep N_LAYERS and plot per-rank peak memory (when running on CUDA) for DDP vs FSDP2. FSDP2 should stay flat; DDP grows linearly.

  2. Enable mixed precision in FSDP2 by passing mp_policy=MixedPrecisionPolicy(param_dtype=torch.bfloat16) to fully_shard. Confirm params are stored in bf16 at rest and cast to fp32 for reductions.

  3. Try reshard_after_forward=True (default) vs False. The former releases sharded params immediately after forward to minimise peak; the latter keeps them for backward to avoid a second all-gather. Measure the memory vs step-time tradeoff.

References#

  • torch.distributed.fsdp.fully_shard docs.

  • PyTorch FSDP paper §3 for the sharding algorithm.

s.summary()
s.save()