Distributed training simulation

Distributed training is how you fit a replica that no longer fits on one GPU, or how you shorten wall time. The words are DDP, FSDP, and ZeRO. The paper title to know for sharded optimiser state is ZeRO. You will simulate the memory story before you pay for a box of H100s.

Intuition, not a vendor tour

• DDP. Each rank holds a full copy of the model. Gradients all-reduce. Memory is easy. Communication is the all-reduce of grads.
• FSDP / ZeRO stage 3 flavour. Parameters, grads, and optimiser state can be sharded. You gather what you need for a layer, then drop it. Memory falls. Communication rises.
• ZeRO stage 1 and 2. Shard optimiser state first, then gradients. Read the idea: the Adam buffers are often the largest residents.

You do not need a 64-GPU job to understand this. Write a simulator that accounts bytes.

from dataclasses import dataclass

@dataclass
class Plan:
    params: int
    dtype_bytes: int
    ranks: int
    adam: bool = True

def ddp_bytes(plan: Plan) -> dict:
    model = plan.params * plan.dtype_bytes
    grads = model
    optim = plan.params * 8 * 2 if plan.adam else 0  # two fp32 moments, rough
    return {"per_rank": model + grads + optim, "mode": "ddp"}

def zero3_bytes(plan: Plan) -> dict:
    model = plan.params * plan.dtype_bytes / plan.ranks
    grads = model
    optim = (plan.params * 8 * 2 / plan.ranks) if plan.adam else 0
    return {"per_rank": model + grads + optim, "mode": "zero3"}

This is a sketch. Real frameworks add buckets, gather buffers, and activation memory. Activations often dominate before parameters do. Write that sentence in your notes so you do not blame ZeRO for an activation OOM.

Assignment

• Estimate per-rank memory for 100M params, fp16, Adam, 1/2/8 ranks, DDP vs ZeRO-3 using the sketch. Show the arithmetic.
• Add a line for activations: batch * seq * width * layers * bytes * a fudge you label as a fudge.
• Write the failure mode: a slow rank, a dropped NCCL, a checkpoint that only one rank wrote.

Opinion checkpoint

Core project 4: a DDP/FSDP/ZeRO intuition simulator.

Next: Supervised fine-tuning.