3.5 Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP¶

# tiny_ddp.py — a from-scratch DistributedDataParallel.
# Run: torchrun --nproc_per_node=4 tiny_ddp.py
import torch
import torch.nn as nn
import torch.distributed as dist


class TinyDDP(nn.Module):
    """
    A minimal DistributedDataParallel that demonstrates the two core ideas:
      1. Broadcast parameters at init so all replicas start identical.
      2. Bucket gradients and all-reduce each bucket asynchronously as soon
         as it is full during the backward pass, overlapping comm with compute.
    """

    def __init__(self, module: nn.Module, bucket_cap_mb: float = 25.0):
        super().__init__()
        self.module = module
        self.world_size = dist.get_world_size()
        self.bucket_bytes_cap = int(bucket_cap_mb * 1024 * 1024)

        # ── Invariant 1: identical initialization ───────────────────────────────
        # Broadcast params AND buffers (e.g. BN running stats) from rank 0.
        for p in self.module.parameters():
            dist.broadcast(p.data, src=0)
        for b in self.module.buffers():
            dist.broadcast(b.data, src=0)

        # ── Assign parameters to buckets in REVERSE registration order ──────────
        # Reverse order approximates reverse-of-forward = order-of-backward,
        # so the first bucket to fill is the first to be communicable.
        params = [p for p in self.module.parameters() if p.requires_grad]
        params = list(reversed(params))

        self._buckets = []          # list of dicts describing each bucket
        cur, cur_bytes = [], 0
        for p in params:
            pbytes = p.numel() * p.element_size()
            if cur and cur_bytes + pbytes > self.bucket_bytes_cap:
                self._buckets.append(cur)
                cur, cur_bytes = [], 0
            cur.append(p)
            cur_bytes += pbytes
        if cur:
            self._buckets.append(cur)

        # Map each parameter -> (bucket_index) and track readiness counters.
        self._param_to_bucket = {}
        for bidx, bucket in enumerate(self._buckets):
            for p in bucket:
                self._param_to_bucket[p] = bidx

        self._pending_work = []     # async all-reduce handles to wait on
        self._ready_counts = [0] * len(self._buckets)
        self._register_hooks()

    def _register_hooks(self):
        """Attach a post-accumulate-grad hook to every parameter."""
        for p in self.module.parameters():
            if not p.requires_grad:
                continue
            # This hook fires AFTER p.grad has been accumulated in backward.
            p.register_post_accumulate_grad_hook(self._make_hook(p))

    def _make_hook(self, param):
        def hook(p):
            bidx = self._param_to_bucket[p]
            self._ready_counts[bidx] += 1
            # When every parameter in this bucket has a gradient, communicate.
            if self._ready_counts[bidx] == len(self._buckets[bidx]):
                self._all_reduce_bucket(bidx)
        return hook

    def _all_reduce_bucket(self, bidx):
        """Flatten the bucket's grads, async all-reduce, store the handle."""
        grads = [p.grad for p in self._buckets[bidx]]
        flat = torch._utils._flatten_dense_tensors(grads)  # one contiguous buffer
        # async_op=True returns immediately; NCCL runs in the background while
        # the backward pass keeps producing earlier-layer gradients.
        handle = dist.all_reduce(flat, op=dist.ReduceOp.SUM, async_op=True)
        self._pending_work.append((handle, flat, self._buckets[bidx]))

    def forward(self, *args, **kwargs):
        return self.module(*args, **kwargs)

    def finish_gradient_synchronization(self):
        """
        Call once after loss.backward(). Waits on all in-flight all-reduces,
        scatters the reduced flat buffers back into each .grad, and averages.
        """
        for handle, flat, bucket in self._pending_work:
            handle.wait()                      # block until this bucket is reduced
            flat /= self.world_size            # SUM -> MEAN
            # Unflatten the contiguous buffer back into per-parameter grads.
            synced = torch._utils._unflatten_dense_tensors(
                flat, [p.grad for p in bucket]
            )
            for p, g in zip(bucket, synced):
                p.grad.copy_(g)
        # Reset for next iteration.
        self._pending_work.clear()
        self._ready_counts = [0] * len(self._buckets)


# ── Usage ───────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")

    model = nn.Sequential(
        nn.Linear(1024, 4096), nn.GELU(),
        nn.Linear(4096, 4096), nn.GELU(),
        nn.Linear(4096, 1024),
    ).to(device)
    ddp = TinyDDP(model, bucket_cap_mb=25.0)
    opt = torch.optim.AdamW(ddp.parameters(), lr=1e-3)

    for step in range(10):
        x = torch.randn(32, 1024, device=device)   # different data per rank
        y = ddp(x).pow(2).mean()                    # toy loss
        opt.zero_grad(set_to_none=True)
        y.backward()                                # hooks fire, comm overlaps
        ddp.finish_gradient_synchronization()       # wait + average
        opt.step()
        if rank == 0:
            print(f"step {step}: loss={y.item():.4f}")

    dist.destroy_process_group()

# verify_tiny_ddp.py -- check TinyDDP == single-process training.
# Run: python verify_tiny_ddp.py   (spawns 4 gloo ranks on CPU; no GPU needed)
import os
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
from tiny_ddp import TinyDDP           # the class defined above

WORLD, LOCAL_B, STEPS, LR = 4, 8, 5, 0.1

def make_model():
    torch.manual_seed(0)               # identical init on every rank
    return nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32))

def make_data():                       # one deterministic global batch per step
    g = torch.Generator().manual_seed(1234)
    return torch.randn(STEPS, WORLD * LOCAL_B, 32, generator=g)

def worker(rank):
    dist.init_process_group("gloo", rank=rank, world_size=WORLD)
    ddp = TinyDDP(make_model(), bucket_cap_mb=1.0)
    opt = torch.optim.SGD(ddp.parameters(), lr=LR)   # SGD: linear, exact
    X = make_data()
    for step in range(STEPS):
        shard = X[step, rank * LOCAL_B:(rank + 1) * LOCAL_B]  # this rank's slice
        loss = ddp(shard).pow(2).mean()
        opt.zero_grad(set_to_none=True)
        loss.backward()                              # hooks fire, buckets reduce
        ddp.finish_gradient_synchronization()        # wait + average
        if step == 0 and rank == 0:                  # save post-sync grad once
            torch.save(torch.cat([p.grad.reshape(-1) for p in ddp.parameters()]),
                       "ddp_grad0.pt")
        opt.step()
    # (Invariant 2) every replica must be bit-identical after each step.
    flat = torch.cat([p.detach().reshape(-1) for p in ddp.parameters()])
    gathered = [torch.empty_like(flat) for _ in range(WORLD)]
    dist.all_gather(gathered, flat)
    for other in gathered:
        assert torch.equal(flat, other), "replicas diverged -- a grad is unsynced"
    if rank == 0:
        torch.save(flat, "ddp_params.pt")
    dist.destroy_process_group()

def reference():                       # single process, FULL batch = ground truth
    model = make_model()
    opt = torch.optim.SGD(model.parameters(), lr=LR)
    X = make_data()
    grad0 = None
    for step in range(STEPS):
        loss = model(X[step]).pow(2).mean()          # mean over the WHOLE batch
        opt.zero_grad(set_to_none=True); loss.backward()
        if step == 0:
            grad0 = torch.cat([p.grad.reshape(-1) for p in model.parameters()])
        opt.step()
    params = torch.cat([p.detach().reshape(-1) for p in model.parameters()])
    return grad0, params

if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "127.0.0.1"; os.environ["MASTER_PORT"] = "29501"
    mp.spawn(worker, nprocs=WORLD, join=True)
    ref_grad0, ref_params = reference()
    ddp_grad0 = torch.load("ddp_grad0.pt")
    ddp_params = torch.load("ddp_params.pt")
    dg = (ref_grad0 - ddp_grad0).abs().max().item()
    dp = (ref_params - ddp_params).abs().max().item()
    print(f"max grad diff (step 0) = {dg:.2e} | max param diff ({STEPS} steps) = {dp:.2e}")
    assert torch.allclose(ref_grad0, ddp_grad0, atol=1e-6), "post-sync grads differ"
    assert torch.allclose(ref_params, ddp_params, atol=1e-5, rtol=1e-4), "DDP != single-process"
    print("OK: TinyDDP reproduces single-process training on the concatenated batch.")

# tiny_sharded_optim.py -- from-scratch ZeRO-1 (optimizer-state sharding).
# Run: torchrun --nproc_per_node=4 tiny_sharded_optim.py
# (uses the gloo backend so it also runs on a CPU-only laptop)
import torch
import torch.nn as nn
import torch.distributed as dist


class TinyShardedOptimizer:
    """
    ZeRO-1 from scratch. Every rank keeps FULL params and FULL gradients, but
    stores AdamW optimizer state (exp_avg m + exp_avg_sq v, and -- in a bf16
    recipe -- an fp32 master weight) for ONLY the parameters it owns. After the
    step, each owner broadcasts its updated parameters back to all ranks so the
    replicas are bit-identical again for the next forward pass.

    Ownership is greedy-by-size, so per-rank optimizer-state bytes are balanced
    even when the parameter count is NOT divisible by the world size (and some
    rank may legitimately own zero parameters).
    """

    def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), eps=1e-8,
                 weight_decay=0.0):
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        self.params = [p for p in params if p.requires_grad]

        # -- Greedy assignment: biggest params first, each to the least-loaded
        #    rank. Balances owned element counts without any divisibility. -----
        order = sorted(range(len(self.params)),
                       key=lambda i: self.params[i].numel(), reverse=True)
        load = [0] * self.world_size
        self.owner = [0] * len(self.params)          # owner rank per parameter
        for i in order:
            r = min(range(self.world_size), key=lambda r: load[r])
            self.owner[i] = r
            load[r] += self.params[i].numel()

        # Build a LOCAL AdamW over just this rank's owned params. A rank that
        # owns nothing gets no optimizer (guarded in step()).
        self.owned_idx = [i for i, r in enumerate(self.owner) if r == self.rank]
        owned = [self.params[i] for i in self.owned_idx]
        self.local_opt = (torch.optim.AdamW(owned, lr=lr, betas=betas, eps=eps,
                                            weight_decay=weight_decay)
                          if owned else None)

    @torch.no_grad()
    def step(self):
        # 1. Average gradients across ranks (SUM -> MEAN). ZeRO-1 keeps full
        #    gradients resident on every rank; only the optimizer STATE is
        #    sharded. (ZeRO-2 is what additionally frees the non-owned grads.)
        for p in self.params:
            if p.grad is None:
                p.grad = torch.zeros_like(p)
            dist.all_reduce(p.grad, op=dist.ReduceOp.SUM)
            p.grad /= self.world_size

        # 2. Each rank updates ONLY its owned parameters, touching ONLY its
        #    shard of optimizer state -- this is the 12*Psi/N memory saving.
        if self.local_opt is not None:
            self.local_opt.step()

        # 3. Broadcast each updated parameter from its owner to all ranks so
        #    every replica is bit-identical again. The canonical ZeRO-1 fuses
        #    this into one all-gather over a flat buffer; per-owner broadcasts
        #    are the readable equivalent.
        for i, p in enumerate(self.params):
            dist.broadcast(p.data, src=self.owner[i])

    def zero_grad(self, set_to_none=True):
        for p in self.params:
            if set_to_none:
                p.grad = None
            elif p.grad is not None:
                p.grad.zero_()

    def owned_state_bytes(self):
        """Measured bytes of THIS rank's optimizer state (m + v [+ step])."""
        if self.local_opt is None:
            return 0
        total = 0
        for st in self.local_opt.state.values():
            for v in st.values():
                if torch.is_tensor(v):
                    total += v.numel() * v.element_size()
        return total


# -- Verification: bit-identical to unsharded AdamW, and a memory measurement --
if __name__ == "__main__":
    dist.init_process_group(backend="gloo")     # CPU-friendly; use nccl on GPUs
    rank, world = dist.get_rank(), dist.get_world_size()

    def build():
        torch.manual_seed(0)                    # identical init on every rank
        # 3 Linears -> 6 parameter tensors; 6 is NOT divisible by world=4.
        # dim=257 is deliberately not a multiple of the world size either.
        return nn.Sequential(nn.Linear(257, 512), nn.GELU(),
                             nn.Linear(512, 512), nn.GELU(),
                             nn.Linear(512, 257))
    m_shard, m_ref = build(), build()           # two identical copies

    sharded = TinyShardedOptimizer(m_shard.parameters(), lr=1e-3,
                                   betas=(0.9, 0.95), weight_decay=0.0)
    ref = torch.optim.AdamW(m_ref.parameters(), lr=1e-3,
                            betas=(0.9, 0.95), weight_decay=0.0)  # explicit wd!

    torch.manual_seed(100 + rank)               # different data per rank
    for step in range(20):
        x = torch.randn(16, 257)
        m_shard(x).pow(2).mean().backward()
        sharded.step(); sharded.zero_grad()     # ZeRO-1 update + re-broadcast
        # Unsharded reference: same data, manual all-reduce(mean), full AdamW.
        loss = m_ref(x).pow(2).mean()
        ref.zero_grad(set_to_none=True); loss.backward()
        for p in m_ref.parameters():
            dist.all_reduce(p.grad, op=dist.ReduceOp.SUM); p.grad /= world
        ref.step()

    # Bit-identical trajectories: the per-parameter AdamW update is elementwise
    # and both paths feed it the identical all-reduced gradient, so sharding
    # WHERE the update runs cannot change the result -- expect exact 0.0.
    max_diff = max((a - b).abs().max().item()
                   for a, b in zip(m_shard.parameters(), m_ref.parameters()))
    if rank == 0:
        print(f"max|param_sharded - param_unsharded| = {max_diff:.3e}")
        assert max_diff == 0.0, "ZeRO-1 must match unsharded AdamW exactly"

    # Per-rank memory: optimizer state for OWNED params only.
    owned = sharded.owned_state_bytes()
    full = sum(8 * p.numel() for p in m_ref.parameters())   # m + v, fp32
    print(f"rank {rank}: owns {owned/1e6:.3f} MB optim state vs "
          f"{full/1e6:.3f} MB unsharded (~{world}x less across the group)")
    dist.destroy_process_group()

# fsdp_train.py — end-to-end FSDP on a toy transformer.
# Run: torchrun --nproc_per_node=4 fsdp_train.py
import functools
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
    StateDictType,
    FullStateDictConfig,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper, CheckpointImpl, apply_activation_checkpointing,
)


# ── A tiny transformer block — the unit FSDP will shard around ───────────────────
class Block(nn.Module):
    def __init__(self, d=1024, h=16):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, h, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x), need_weights=False)
        x = x + a
        x = x + self.mlp(self.ln2(x))
        return x


class ToyTransformer(nn.Module):
    def __init__(self, vocab=50_000, d=1024, n_layers=12):
        super().__init__()
        self.emb = nn.Embedding(vocab, d)
        self.blocks = nn.ModuleList([Block(d) for _ in range(n_layers)])
        self.head = nn.Linear(d, vocab, bias=False)

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


def main():
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")

    # Build the model on meta/CPU; FSDP will shard it onto the GPUs.
    model = ToyTransformer()

    # ── Auto-wrap policy: make every `Block` its own FSDP unit ──────────────────
    # Each Block becomes a FlatParameter sharded across ranks. Params are
    # all-gathered just-in-time per block and freed right after.
    auto_wrap = functools.partial(
        transformer_auto_wrap_policy, transformer_layer_cls={Block}
    )

    # ── Mixed precision: bf16 compute, fp32 reductions for numerical safety ─────
    mp = MixedPrecision(
        param_dtype=torch.bfloat16,      # params gathered/used in bf16
        reduce_dtype=torch.float32,      # grad reduce-scatter accumulates in fp32
        buffer_dtype=torch.bfloat16,
    )

    model = FSDP(
        model,
        auto_wrap_policy=auto_wrap,
        mixed_precision=mp,
        sharding_strategy=ShardingStrategy.FULL_SHARD,   # ZeRO-3 semantics
        device_id=device,
        use_orig_params=True,            # play nicely with torch.compile & optims
    )

    # ── Activation checkpointing: recompute block activations in backward to
    #    save activation memory (orthogonal to FSDP param sharding) ─────────────
    apply_activation_checkpointing(
        model,
        checkpoint_wrapper_fn=functools.partial(
            checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT
        ),
        check_fn=lambda m: isinstance(m, Block),
    )

    # The optimizer sees only this rank's parameter shards (sharded optim state).
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95))

    model.train()
    for step in range(20):
        idx = torch.randint(0, 50_000, (8, 512), device=device)   # local batch
        logits = model(idx)
        loss = nn.functional.cross_entropy(
            logits.view(-1, logits.size(-1)), idx.view(-1)
        )
        opt.zero_grad(set_to_none=True)
        loss.backward()        # all-gathers + reduce-scatters happen inside, overlapped
        # Gradient clipping across shards: FSDP provides a sharded-aware clip.
        model.clip_grad_norm_(1.0)
        opt.step()
        if rank == 0:
            print(f"step {step:3d} | loss {loss.item():.4f}")

    # ── Checkpointing: no single rank holds the full model under FULL_SHARD.
    #    Gather a full (unsharded) state dict ONLY on rank 0 to write to disk. ──
    save_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_cfg):
        full_sd = model.state_dict()       # all-gathers params, materialized on rank 0
        if rank == 0:
            torch.save(full_sd, "ckpt_full.pt")
    # For large models prefer SHARDED_STATE_DICT (each rank writes its own shard),
    # which avoids materializing the whole model on one GPU/host.

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

# fsdp2_train.py — the same ZeRO-3 recipe with the FSDP2 (`fully_shard`) API.
# Run: torchrun --nproc_per_node=4 fsdp2_train.py
import os
import torch
import torch.nn as nn
import torch.distributed as dist
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_model_state_dict, StateDictOptions,
)
from fsdp_train import ToyTransformer          # the model defined above

dist.init_process_group(backend="nccl")
local_rank = int(os.environ.get("LOCAL_RANK", 0))
torch.cuda.set_device(local_rank)

# A 1-D mesh over all ranks == FULL_SHARD. For HYBRID_SHARD (HSDP) use a 2-D mesh
#   init_device_mesh("cuda", (n_nodes, gpus_per_node),
#                    mesh_dim_names=("replicate", "shard"))
# and pass that mesh instead — FSDP2 reads the shard/replicate roles off the mesh.
mesh = init_device_mesh("cuda", (dist.get_world_size(),))
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16,   # gather/compute in bf16
                          reduce_dtype=torch.float32)   # reduce grads in fp32

model = ToyTransformer().cuda()
for blk in model.blocks:                        # one call per FSDP2 unit ...
    fully_shard(blk, mesh=mesh, mp_policy=mp)
fully_shard(model, mesh=mesh, mp_policy=mp)     # ... then the root: shards emb+head

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95))
GRAD_ACCUM = 2
for step in range(20):
    opt.zero_grad(set_to_none=True)
    for micro in range(GRAD_ACCUM):
        last = micro == GRAD_ACCUM - 1
        # FSDP2's no_sync(): skip the gradient reduce-scatter until the last
        # micro-batch (grads accumulate unsharded on each rank until then).
        model.set_requires_gradient_sync(last)
        idx = torch.randint(0, 50_000, (8, 512), device="cuda")
        loss = nn.functional.cross_entropy(
            model(idx).flatten(0, 1), idx.flatten()) / GRAD_ACCUM
        loss.backward()
    # Params and grads are DTensors, so the STOCK clip is already shard-aware:
    # it reduces the squared norm across the mesh. No FSDP-specific call needed
    # (contrast FSDP1, which requires model.clip_grad_norm_).
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

# Checkpointing goes through torch.distributed.checkpoint's state-dict helpers,
# which understand DTensor. full_state_dict=True gathers to rank 0 (small models);
# drop it to get each rank's shard for torch.distributed.checkpoint.save (big ones).
sd = get_model_state_dict(model, options=StateDictOptions(
    full_state_dict=True, cpu_offload=True))
if dist.get_rank() == 0:
    torch.save(sd, "ckpt_fsdp2.pt")
dist.destroy_process_group()

# accelerate_fsdp.yaml — launch with: accelerate launch --config_file accelerate_fsdp.yaml train.py
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
mixed_precision: bf16
num_processes: 8                          # GPUs on this node
fsdp_config:
  fsdp_version: 2                         # select FSDP2 (recent Accelerate versions;
                                          # older ones expose fsdp_sharding_strategy instead)
  fsdp_reshard_after_forward: true        # ZeRO-3 semantics (false => ZeRO-2-ish)
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_transformer_layer_cls_to_wrap: Block
  fsdp_activation_checkpointing: true

{
  "train_micro_batch_size_per_gpu": 8,
  "gradient_accumulation_steps": 4,
  "bf16": { "enabled": true },
  "gradient_clipping": 1.0,
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "offload_optimizer": { "device": "cpu", "pin_memory": true }
  }
}

import deepspeed
engine, opt, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(), config="ds_config.json")
for x, y in loader:                    # engine handles accumulation + clipping
    loss = nn.functional.cross_entropy(engine(x).flatten(0, 1), y.flatten())
    engine.backward(loss)              # NOT loss.backward()
    engine.step()                      # NOT opt.step(); also steps the LR schedule

# train.py -- one torchrun-launchable pretraining loop.
# 1 GPU:   torchrun --nproc_per_node=1 train.py --data ./tokens --parallel ddp
# 8 GPUs:  torchrun --nproc_per_node=8 train.py --data ./tokens --parallel fsdp
# 2 nodes: torchrun --nnodes=2 --node_rank=$RANK --nproc_per_node=8 \
#                   --rdzv_endpoint=$HEAD:29500 train.py --data ./tokens --parallel fsdp
import argparse, contextlib, functools, glob, math, os, time
import numpy as np
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import (FullyShardedDataParallel as FSDP,
    MixedPrecision, ShardingStrategy, StateDictType, FullStateDictConfig)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy


class Block(nn.Module):                         # one Block == one FSDP unit
    def __init__(self, d, h):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, h, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x, mask):
        a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
                         attn_mask=mask, need_weights=False)   # causal mask
        x = x + a
        return x + self.mlp(self.ln2(x))


class GPT(nn.Module):
    def __init__(self, vocab=50304, d=768, h=12, n_layers=12, ctx=1024):
        super().__init__()
        self.ctx = ctx
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(ctx, d)
        self.blocks = nn.ModuleList([Block(d, h) for _ in range(n_layers)])
        self.lnf = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok.weight       # weight tying (Press & Wolf): head shares the embedding
                                                  # -> 124M at these dims, not ~163M untied

    def forward(self, idx):
        B, T = idx.shape
        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))[None]
        mask = torch.triu(torch.full((T, T), float("-inf"), device=idx.device), 1)
        for blk in self.blocks:
            x = blk(x, mask)
        return self.head(self.lnf(x))


class ShardedTokenLoader:
    """Reads (ctx+1)-token windows from uint16 .bin shards. All ranks share one
    RNG stream, then rank r strides its window start by r*(ctx+1), so the global
    batch is a partition -- no two ranks train on the same token span."""
    def __init__(self, data_dir, ctx, local_bsz, rank, world, split="train"):
        self.files = sorted(glob.glob(os.path.join(data_dir, f"{split}_*.bin")))
        assert self.files, f"no {split}_*.bin shards in {data_dir}"
        self.ctx, self.bsz, self.rank, self.world = ctx, local_bsz, rank, world
        self.mm = [np.memmap(f, dtype=np.uint16, mode="r") for f in self.files]
        self.sizes = [len(m) - (ctx + 1) for m in self.mm]
        self.g = torch.Generator().manual_seed(1234)     # identical on all ranks

    def batch(self, device):
        xs, ys = [], []
        for _ in range(self.bsz):
            fi = int(torch.randint(len(self.mm), (1,), generator=self.g))
            base = int(torch.randint(self.sizes[fi], (1,), generator=self.g))
            off = (base + self.rank * (self.ctx + 1)) % self.sizes[fi]   # disjoint
            w = torch.from_numpy(self.mm[fi][off:off + self.ctx + 1].astype(np.int64))
            xs.append(w[:-1]); ys.append(w[1:])
        x = torch.stack(xs).to(device, non_blocking=True)
        y = torch.stack(ys).to(device, non_blocking=True)
        return x, y


def lr_at(step, warmup, total, base_lr, min_lr):    # linear warmup + cosine decay
    if step < warmup:
        return base_lr * (step + 1) / warmup
    if step >= total:
        return min_lr
    r = (step - warmup) / max(1, total - warmup)
    return min_lr + 0.5 * (base_lr - min_lr) * (1 + math.cos(math.pi * r))


def save_checkpoint(model, step, ckpt_dir, rank):
    """Full (unsharded) state dict gathered on rank 0. For big models prefer a
    SHARDED_STATE_DICT + torch.distributed.checkpoint (see chapter 3.12)."""
    path = os.path.join(ckpt_dir, f"step_{step}.pt")
    if isinstance(model, FSDP):
        cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
        with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
            sd = model.state_dict()             # all-gathers params to rank 0
            if rank == 0:
                torch.save({"model": sd, "step": step}, path)
    elif rank == 0:                             # DDP: rank 0 holds the full model
        torch.save({"model": model.module.state_dict(), "step": step}, path)
    dist.barrier()                              # no rank races ahead of the write


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--parallel", choices=["ddp", "fsdp"], default="fsdp")
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--warmup", type=int, default=100)
    ap.add_argument("--local-bsz", type=int, default=8)
    ap.add_argument("--ctx", type=int, default=1024)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--grad-accum", type=int, default=1)
    ap.add_argument("--backend", default="nccl", help="nccl on GPUs, gloo on CPU")
    ap.add_argument("--ckpt-every", type=int, default=500)
    ap.add_argument("--ckpt-dir", default="./ckpt")
    args = ap.parse_args()

    dist.init_process_group(backend=args.backend)
    rank, world = dist.get_rank(), dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    if args.backend == "nccl":
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
    else:
        device = torch.device("cpu")            # CPU/gloo path for the laptop toy
    torch.manual_seed(0)                        # identical init on every rank

    model = GPT(ctx=args.ctx).to(device)
    if args.parallel == "ddp":
        # CPU/gloo DDP takes device_ids=None; only CUDA/nccl pins a device index.
        model = DDP(model, device_ids=[local_rank] if args.backend == "nccl" else None)
        clip = lambda: nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    else:
        mp = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32,
                            buffer_dtype=torch.bfloat16)
        model = FSDP(model,
            auto_wrap_policy=functools.partial(
                transformer_auto_wrap_policy, transformer_layer_cls={Block}),
            mixed_precision=mp, sharding_strategy=ShardingStrategy.FULL_SHARD,
            device_id=device, use_orig_params=True)
        clip = lambda: model.clip_grad_norm_(1.0)    # sharded-aware global norm

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr,
                            betas=(0.9, 0.95), weight_decay=0.1)
    loader = ShardedTokenLoader(args.data, args.ctx, args.local_bsz, rank, world)
    os.makedirs(args.ckpt_dir, exist_ok=True)

    model.train(); t0 = time.time()
    for step in range(args.steps):
        for grp in opt.param_groups:
            grp["lr"] = lr_at(step, args.warmup, args.steps, args.lr, args.lr * 0.1)
        opt.zero_grad(set_to_none=True)
        step_loss = 0.0
        for micro in range(args.grad_accum):    # gradient accumulation: G micro-steps / opt step
            x, y = loader.batch(device)
            last = micro == args.grad_accum - 1
            # Under DDP, skip the all-reduce on every micro-step but the last.
            sync = contextlib.nullcontext() if last or args.parallel != "ddp" else model.no_sync()
            with sync:
                loss = nn.functional.cross_entropy(
                    model(x).view(-1, 50304), y.view(-1)) / args.grad_accum
                loss.backward()                 # collectives fire on the last micro-step
            step_loss += loss.item()            # sum of G scaled losses = full-batch mean
        clip()
        opt.step()
        if rank == 0 and step % 20 == 0:
            ms = (time.time() - t0) / (step + 1) * 1e3
            print(f"step {step:5d} | loss {step_loss:.4f} "
                  f"| lr {opt.param_groups[0]['lr']:.2e} | {ms:.0f} ms/step")
        if step > 0 and step % args.ckpt_every == 0:
            save_checkpoint(model, step, args.ckpt_dir, rank)

    save_checkpoint(model, args.steps, args.ckpt_dir, rank)
    dist.destroy_process_group()


if __name__ == "__main__":
    main()