DL S051 · Distributed Training Basics — DDP and FSDP

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 51 / 80 · Module M09 · ~2 hours

The story we're starting with

Your 100M-param Llama-clone fits fine on one A100 40GB — model, optimizer states, activations, all in ~15 GB with room to spare. But scale up. A 1B-param model needs ~20 GB just for parameters + optimizer state in mixed precision. A 7B model needs ~140 GB — three A100 80GBs. A 70B model needs ~1.5 TB — twenty A100 80GBs. Suddenly "just fit it in one GPU" is not an option.

There are two big ideas for multi-GPU training. Data parallelism — replicate the model on every GPU and split the batch across them. Model parallelism — shard the model itself across GPUs. Modern LLM pretraining is a mix: DDP (data parallel, simple) works up to ~1B params. FSDP / ZeRO (sharded parameters + optimizer state) is what you need past that. Beyond ~70B you also need tensor parallelism and pipeline parallelism.

For your 100M M09 run you'll use one GPU and no parallelism — but you should understand DDP and FSDP because (a) the next model you scale up will need them, and (b) reading loss curves and throughput numbers from other people's runs is impossible without this vocabulary.

Prerequisites

• S018 (GPU + mixed precision) — you understand .to("cuda") and bf16.
• S050 (architecture choices) — for parameter-count intuition.

1 · Memory accounting — where does the VRAM go?

Before any parallelism talk, you need to understand what's living in GPU memory during training. There are four things:

Rule of thumb for a bf16 Adam-trained transformer: ~16 bytes per parameter of steady-state memory (ignoring activations).

Concrete:

Activations on top of the steady state can add another 10–100 GB depending on batch size and seq length. For a 7B model at 4k context, activations easily push you to 200 GB total memory needed. And you haven't fit a single training example yet.

This is why distributed training exists.

2 · DDP (DistributedDataParallel) — the easy one

Mental model. Copy the whole model to every GPU. Split each batch into world_size micro-batches, one per GPU. Each GPU runs forward + backward on its slice. At the end of backward, all_reduce the gradients so every GPU sees the mean gradient across all GPUs. Optimizer step is identical on all GPUs (same params, same gradients, same result).

2.1 Memory: no savings

Every GPU holds the entire 16 B/param footprint. DDP does not reduce per-GPU memory. It only lets you process a bigger effective batch by parallelizing.

2.2 Compute: near-linear scaling

If comms are fast (NVLink for intra-node, InfiniBand for inter-node), throughput scales ~90–95% of ideal up to 8 GPUs. The all_reduce cost is O(N) bytes per step, amortized across the compute of a step. For small models it can become the bottleneck; for real LLMs it's usually <15% of step time.

2.3 Four-line switch

Single-GPU training loop → DDP is genuinely 4 lines:

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
 
dist.init_process_group("nccl")                       # 1
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)                     # 2
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])           # 3
 
# In your DataLoader, use DistributedSampler
sampler = DistributedSampler(dataset)                 # 4
loader = DataLoader(dataset, sampler=sampler, ...)

Launch:

torchrun --nproc_per_node=8 train.py

That's it. DDP is the "just make it multi-GPU" button. Use it for anything ≤1B params on a single node.

3 · ZeRO — the memory-saving family

DDP's limitation: replicated everything. ZeRO (Rajbhandari 2019, DeepSpeed team) asks: what if we shard the redundant stuff?

Three stages, in escalating aggressiveness:

3.1 ZeRO-1: shard optimizer states

Each GPU owns 1/N-th of the optimizer state (Adam's m, v, master weights). Forward/backward still work on the full replicated params. During optimizer step, each GPU updates its slice; then all-gather the updated weights so every GPU has the new full model.

Savings: 12 B/param optimizer state → 12/N B/param. On 8 GPUs, a 7B model drops from 84 GB optimizer state to 10.5 GB per GPU. Total per-GPU memory: 16 → ~5 B/param.

3.2 ZeRO-2: also shard gradients

Same idea, applied to gradients. Each GPU only stores its 1/N slice of gradients after backward.

Savings: another 2/N B/param.

3.3 ZeRO-3: also shard parameters (= FSDP)

Now even the parameters themselves are sharded. During forward/backward, each GPU temporarily all-gathers the params it needs just for the current layer, then releases them.

Savings: full 16/N B/param footprint. On 8 GPUs, a 7B model uses 14 GB per GPU — fits on one A100.

PyTorch's FSDP = ZeRO stage 3. Same idea, different codebase.

4 · FSDP — how it actually works

FSDP wraps each block (typically a transformer layer). During forward:

1. All-gather this layer's params from every rank → each rank has the full layer.
2. Run forward on the layer.
3. Free the gathered params (keep only the local shard).
4. Move to next layer, repeat.

During backward, mirror it: all-gather params, compute gradients, reduce-scatter gradients (each rank ends up with its shard of the gradient), free the gathered params, move on.

The math per step:

• Forward all-gathers: N params × 2 bytes = 2N bytes moved across the network, per step.
• Backward all-gathers + reduce-scatters: ~2 × 2N = 4N bytes.
• Total comms: ~6N bytes per training step.

Compare DDP: only 2N bytes (one all-reduce of gradients). FSDP has 3× the comms.

Why anyone accepts 3× the comms: because otherwise the model doesn't fit at all. FSDP unlocks training models bigger than any single GPU's VRAM.

4.1 PyTorch FSDP boilerplate

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial
 
wrap_policy = partial(transformer_auto_wrap_policy,
                      transformer_layer_cls={LlamaLikeBlock})
 
model = FSDP(
    model,
    auto_wrap_policy=wrap_policy,
    device_id=local_rank,
    mixed_precision=torch.distributed.fsdp.MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    ),
)

Same launch: torchrun --nproc_per_node=8 train.py.

5 · Comms cost — the case for NVLink and InfiniBand

FSDP's ~6N bytes/step comms budget has to complete inside the compute time of a step, or your GPUs stall waiting for data.

Concrete example: a 7B model, 8 GPUs.

• Comms per step ≈ 6 × 7e9 × 2 bytes = 84 GB moved.
• PCIe 4.0: 64 GB/s bidirectional. Comms alone: 1.3 seconds.
• NVLink 3.0 (A100): 600 GB/s. Comms alone: 0.14 seconds.
• InfiniBand HDR: 200 Gbps = 25 GB/s. Comms alone: 3.4 seconds.

Step compute time for a 7B model on 8×A100 at MFU 40%: ~0.5 seconds.

• On NVLink: comms 0.14 s vs compute 0.5 s → mostly overlapped, ~10% overhead.
• On PCIe only: comms 1.3 s vs compute 0.5 s → stall of 0.8 s per step. 60% throughput lost.
• On IB inter-node: 3.4 s stall → 85% throughput lost.

This is why serious training runs use NVLink for intra-node and InfiniBand (or fancier RDMA) for inter-node. Consumer GPUs with only PCIe are a bad fit for FSDP.

6 · When to use what — decision tree

For our M09 100M run: one GPU, no parallelism. Total training time on one A100 = ~3 hours for 2B tokens. Adding a second GPU literally isn't worth it.

If you scale to a 1B model at Chinchilla-optimal (20B tokens): ~150 A100-hours on one GPU. That's when DDP on 2-4 GPUs starts making sense to cut wall clock.

7 · A real training-script pattern

Here's a skeleton that works single-GPU today, DDP tomorrow, FSDP next year:

import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
 
def setup():
    if "RANK" in os.environ:
        dist.init_process_group("nccl")
        rank = dist.get_rank()
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)
    else:
        rank = 0; local_rank = 0
    return rank, local_rank
 
def build_model(cfg, local_rank):
    model = LlamaLike(cfg).to(local_rank)
    if dist.is_initialized():
        model = DDP(model, device_ids=[local_rank])
    return model
 
def train():
    rank, local_rank = setup()
    model = build_model(CONFIG, local_rank)
    ...

Launch single-GPU: python train.py. Launch multi-GPU: torchrun --nproc_per_node=4 train.py. Same code path.

8 · Gradient accumulation — the poor person's DDP

Can't afford more GPUs? Simulate a bigger batch by running k forward/backward passes and accumulating gradients before stepping:

optimizer.zero_grad()
for i, (x, y) in enumerate(loader):
    loss = model(x, y) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Effective batch size = micro_batch × accum_steps. Comes at zero comms cost — you're just running the training loop accum_steps times per optimizer update.

For our 100M run on one A100, we might set micro_batch=32, accum_steps=8 → effective batch 256. Roughly what published Llama configs use per-GPU. Costs no extra memory (activations are freed between micro-batches).

9 · War stories

10 · The 2025 receipts — FSDP2, DTensor, and Composer

PyTorch's distributed story matured hard in 2024–2025. If you last read the FSDP docs in 2023, half of what you know is now outdated. Here's the diff.

10.1 · FSDP2 (per-parameter sharding, PyTorch 2.4+)

Original FSDP shards per FlatParameter — it concatenates all params in a wrapped module into one giant flat tensor and shards that. Simple, but has three problems: (1) you can't state_dict() a single layer without unsharding the whole group; (2) mixed precision is per-FlatParameter, not per-tensor — so a layer that needs fp32 forces the whole group to fp32; (3) LoRA and other partial-tuning setups are painful.

FSDP2, shipped in PyTorch 2.4 (July 2024) and stable in 2.5+, shards per parameter tensor using the new DTensor (distributed tensor) primitive. The API is different — no wrap policy, you call fully_shard(module) and it does per-param sharding under the hood. Benefits:

• Cleaner state_dict: model.state_dict() returns DTensors you can materialize per-param.
• Per-tensor mixed precision (fp32 norms, bf16 linears, fp8 experts — all in one model).
• Composes with tensor parallelism natively via parallelize_module.
• 5–15% throughput improvement on 8+ GPU workloads (per PyTorch 2.5 release notes).

When to use FSDP1 vs FSDP2 in mid-2025: if you're starting fresh, use FSDP2. If you have a working FSDP1 script, stay on it — the migration has sharp edges around checkpoint format (DCP works with both, but the on-disk layouts differ). See PyTorch's torch.distributed.fsdp.fully_shard docs.

10.2 · DTensor — the new abstraction under everything

torch.distributed.tensor.DTensor is a tensor that knows how it's sharded (Shard(dim=i), Replicate(), Partial()) across a DeviceMesh. Every 2025 PyTorch distributed feature is built on it:

• FSDP2 = DTensor with Shard(0) on the FSDP mesh.
• Tensor Parallelism (parallelize_module) = DTensor with Shard(0) or Shard(1) on the TP mesh, plus rules for how each operator (matmul, layernorm) handles the sharding.
• 2D parallelism (FSDP × TP) = DTensor on a 2D DeviceMesh: mesh_2d = init_device_mesh("cuda", (dp, tp), mesh_dim_names=("dp", "tp")).
• Pipeline parallelism (torch.distributed.pipelining, stable in 2.5) also uses DTensor for stage boundaries.

The practical impact: for the first time, you can write a single training script that scales from 1 GPU to 8 to 512 without rewriting the parallelization glue. torchtitan (github.com/pytorch/torchtitan, Meta's reference LLM training repo) is the canonical example — read its train.py before your first multi-node run.

10.3 · MosaicML Composer + LLM Foundry — batteries-included alternative

If you'd rather not hand-roll FSDP + AMP + gradient accumulation + checkpointing + wandb, MosaicML's Composer (github.com/mosaicml/composer) wraps it all. Their LLM Foundry (github.com/mosaicml/llm-foundry) is the training script Databricks used for DBRX (132B MoE, 3072 H100s, March 2024). Highlights of what Composer buys you for free:

• Automatic FSDP wrapping with sensible policies.
• Auto-resume from the latest checkpoint on job restart (S3 or local).
• Loss-spike detection with automatic batch-skip recovery (their write-up: mosaicml.com/blog/mpt-7b describes their spike protocol).
• Deterministic dataloader that survives resume mid-epoch — the single most annoying thing to build yourself.
• Native bf16 + fp8 mixed precision (via Transformer Engine on H100).

DBRX's public training receipts (Databricks blog, March 2024): 132B total / 36B active MoE, 12T tokens, 3072 H100s, ~2 months wall-clock, ~$10M in compute. For comparison, their previous MPT-7B run was 8T tokens on 512 A100s in 9.5 days. Composer scaled the same code across both.

Lesson for you: for a 100M run, hand-roll it — you'll learn more. For anything >1B, seriously consider Composer. The hand-rolled version is a career-long tax you don't need to pay.

10.4 · The 2025 comms landscape

A quick reality check on interconnects, mid-2025 pricing:

Rule of thumb the numbers give you: FSDP is bandwidth-bound. If your AllGather + ReduceScatter per step exceeds ~30–50% of step time, you're bottlenecked. Fixes in order of impact: (1) larger micro-batch → more compute per gather; (2) activation checkpointing → fewer gathers per step; (3) upgrade interconnect (💸); (4) switch to 2D parallelism (TP inside the node, FSDP across nodes).

Further reading (10.x)

• FSDP2 announcement — pytorch.org/blog/pytorch2-4
• torchtitan — github.com/pytorch/torchtitan
• DTensor RFC — github.com/pytorch/pytorch/issues/88838
• MosaicML LLM Foundry — github.com/mosaicml/llm-foundry
• DBRX training blog — databricks.com/blog/introducing-dbrx-new-state-art-open-llm
• BLOOM chronicles — github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md
• OPT chronicles — github.com/facebookresearch/metaseq/tree/main/projects/OPT/chronicles.md

11 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (memory accounting) and §6 (decision tree) in 24 hours. Revisit the full session on day 7 focusing on §10 (2025 receipts) and the two chronicles.

Next session (S052): rent an actual GPU — Lambda, RunPod, Modal walkthrough with 2025 dollar amounts.

Sticky note (keep on your desk): "DDP if you fit. FSDP2 if you don't. TP inside the node. Multi-node = IB or don't bother."

Legacy recall drills (kept for spaced-review continuity)

Stretch

Build a memory-estimator. Function estimate(N, world_size, mode) where mode ∈ {"single", "ddp", "zero1", "zero2", "fsdp"} returns per-GPU steady-state bytes. Add optional activation_checkpointing=True that shrinks activation memory by 4×. Use it to answer: "on 8×A100-40GB, what's the biggest model I can train under each mode?"

In your own words

"The three axes of scaling a training run are ____, ____, ____. DDP moves you along the first only; FSDP unlocks the second; TP/PP unlock the third."

Spaced-review pointer

• S018 (mixed precision) — bf16 vs fp32 costs matter to every calculation here.
• S050 (architecture) — the per-param memory is fixed by the choices made there.

Next-session teaser

Enough theory. Session 052 is the actual GPU rental walkthrough — how to spin up an A100 on Runpod/Modal/Vast for under $2/hour, sync your training code, launch, monitor, and kill it before your credit card cries. Screenshot walkthrough, real dollar amounts, real commands.

Bring back tomorrow

• Number: 16 B/param bf16 Adam steady-state.
• Rule: DDP = replicate + all-reduce; FSDP = shard + all-gather.
• Command: torchrun --nproc_per_node=8 train.py.

Previous: ← DL S050 · Next: DL S052 →