"""ZeRO-2 sharding of the backbone's fp32 master weights, AdamW moments and gradients across data-parallel ranks (SPEC.md ยง12.9). The trainer keeps bf16 compute weights on every rank (the functional backbone calls its layers' weights directly, so they must be whole on the device) and, without sharding, a full fp32 master copy + two AdamW moments + an fp32 gradient: 16 bytes per trainable backbone parameter, ~22 GB for Qwen3.5-2B, more than an L4 holds. :class:`ZeroMaster` splits those 16 bytes over ``world`` ranks: * The trainable parameters are laid end to end in a flat vector of ``N`` elements cut into buckets (``bucket`` elements, the last one padded to a multiple of ``world``). Rank ``r`` owns the ``r``-th of ``world`` equal pieces of every bucket; its shard is those pieces concatenated (``shard_len`` ~ ``N / world``). * :attr:`master` is ONE fp32 parameter (the shard) that AdamW updates; the optimizer's moments are shard-sized. Element-wise AdamW over a flat tensor is the same arithmetic as over the separate tensors (same decay everywhere, as the unsharded trainer). * A step's bf16 gradients stay on the parameters' ``.grad`` until the step's last part (:meth:`collect`): a one-part step keeps them where autograd put them (no copy); a several-part step accumulates them in a host fp32 vector (no second device copy). :meth:`reduce_scatter` then sums every rank's gradient into the owners' fp32 shard (bucket by bucket, fp32), so :attr:`master` ``.grad`` holds the shard of the sum. * After ``opt.step()``, :meth:`gather_params` rounds each rank's updated shard to bf16 and all-gathers it into every rank's compute weights: bitwise the unsharded trainer's ``p.copy_(master)``. ``sim_world`` (smoke tests on one GPU): shard as rank 0 of a ``sim_world``-rank job with no collectives, so memory and step time are the real per-rank ones; the other ranks' pieces are simply not updated (their weights stay as loaded). Never for real training. """ from __future__ import annotations import bisect import math from typing import Iterator import torch from torch import Tensor, nn BUCKET = 1 << 25 # 32M elements: 128 MB fp32 per collective def _dist(): import torch.distributed as dist return dist class ZeroMaster: def __init__(self, params: list[nn.Parameter], rank: int, world: int, bucket: int = BUCKET, sim_world: int = 0) -> None: if not params: raise ValueError("ZeroMaster needs trainable parameters") self.params = params self.sim = sim_world > 1 self.world = sim_world if self.sim else world self.rank = 0 if self.sim else rank self.comm = world > 1 and not self.sim self.numel = [p.numel() for p in params] self.offsets = [0] for n in self.numel: self.offsets.append(self.offsets[-1] + n) self.N = self.offsets[-1] w = self.world bucket = max(w, bucket - bucket % w) # (start, end, padded length) of every bucket in the flat vector self.buckets: list[tuple[int, int, int]] = [] for s in range(0, self.N, bucket): e = min(self.N, s + bucket) self.buckets.append((s, e, math.ceil((e - s) / w) * w)) self.shard_len = sum(L // w for _, _, L in self.buckets) dev = params[0].device self.device = dev shard = torch.empty(self.shard_len, dtype=torch.float32, device=dev) o = 0 for b, (s, e, L) in enumerate(self.buckets): piece = L // w shard[o:o + piece] = self._bucket_values(b)[self.rank * piece:(self.rank + 1) * piece] o += piece self.master = nn.Parameter(shard) self.grad_buf = torch.zeros_like(shard) # resident: the shard's fp32 gradient self.held: list[Tensor | None] | None = None # one-part step: the bf16 grads self.host: Tensor | None = None # several-part step: fp32 sum on the host self._host_buf: Tensor | None = None # kept once allocated (pinned) self.spills = 0 # ---------------------------------------------------------------- flat layout def _segments(self, s: int, e: int) -> Iterator[tuple[int, int, int, int]]: """(param index, offset in param, offset in [s, e), length) covering [s, e).""" i = bisect.bisect_right(self.offsets, s) - 1 while i < len(self.numel) and self.offsets[i] < e: a = max(s, self.offsets[i]) b = min(e, self.offsets[i + 1]) yield i, a - self.offsets[i], a - s, b - a i += 1 def _bucket_values(self, b: int) -> Tensor: s, e, L = self.buckets[b] out = torch.zeros(L, dtype=torch.float32, device=self.device) for i, po, bo, n in self._segments(s, e): out[bo:bo + n] = self.params[i].detach().reshape(-1)[po:po + n].float() return out def _pieces(self) -> Iterator[tuple[int, int, int, int, int]]: """(bucket start, end, padded length, piece length, shard offset).""" o = 0 for s, e, L in self.buckets: piece = L // self.world yield s, e, L, piece, o o += piece # ---------------------------------------------------------------- gradients @property def has_grad(self) -> bool: return self.held is not None or self.host is not None def collect(self, more: bool) -> None: """Take the parameters' current ``.grad`` (one part's bf16 gradient) into the step's sum and clear them. ``more``: further parts of this step will run, so the sum moves to the host (the next part's backward then has the device memory).""" grads = [p.grad for p in self.params] for p in self.params: p.grad = None if all(g is None for g in grads): return if self.held is None and self.host is None and not more: self.held = grads return if self.host is None: if self._host_buf is None: self._host_buf = torch.zeros(self.N, dtype=torch.float32, pin_memory=self.device.type == "cuda") else: self._host_buf.zero_() self.host = self._host_buf self.spills += 1 if self.held is not None: self._to_host(self.held) self.held = None self._to_host(grads) def _to_host(self, grads: list[Tensor | None]) -> None: assert self.host is not None for i, g in enumerate(grads): if g is not None: o = self.offsets[i] self.host[o:o + self.numel[i]] += g.detach().reshape(-1).float().cpu() def discard(self) -> None: self.held = None self.host = None for p in self.params: p.grad = None self.master.grad = None def reduce_scatter(self) -> None: """Sum the step's gradient over ranks into :attr:`master` ``.grad`` (this rank's shard, fp32). A rank without gradient contributes zeros.""" held, host = self.held, self.host self.held = self.host = None for s, e, L, piece, o in self._pieces(): buf = torch.zeros(L, dtype=torch.float32, device=self.device) if host is not None: buf[:e - s] = host[s:e].to(self.device, non_blocking=True) elif held is not None: for i, po, bo, n in self._segments(s, e): g = held[i] if g is not None: buf[bo:bo + n] = g.reshape(-1)[po:po + n].float() out = self.grad_buf[o:o + piece] if self.comm: _reduce_scatter(out, buf) else: out.copy_(buf[self.rank * piece:(self.rank + 1) * piece]) del buf del held, host self.master.grad = self.grad_buf def grad_sq_norm(self) -> Tensor: """Squared L2 norm of the whole (summed) gradient, identical on every rank.""" sq = self.grad_buf.pow(2).sum() if self.comm: _dist().all_reduce(sq) return sq # ---------------------------------------------------------------- weights @torch.no_grad() def gather_params(self) -> None: """Every rank's updated shard -> bf16 -> the compute weights of every rank.""" dtype = self.params[0].dtype for s, e, L, piece, o in self._pieces(): mine = self.master[o:o + piece].to(dtype) if self.comm: full = torch.empty(L, dtype=dtype, device=self.device) _all_gather(full, mine) lo, hi = 0, e - s else: # one process: only this rank's piece is known full = torch.zeros(L, dtype=dtype, device=self.device) full[self.rank * piece:(self.rank + 1) * piece] = mine lo = self.rank * piece hi = min(e - s, lo + piece) for i, po, bo, n in self._segments(s, e): a, b = max(bo, lo), min(bo + n, hi) if a < b: self.params[i].view(-1)[po + a - bo:po + b - bo] = full[a:b] @torch.no_grad() def full_master(self) -> Tensor: """The whole fp32 master vector (all ranks; tests).""" parts = [] for s, e, L, piece, o in self._pieces(): mine = self.master[o:o + piece].detach() if self.comm: full = torch.empty(L, dtype=mine.dtype, device=self.device) _all_gather(full, mine.contiguous()) else: full = mine parts.append(full[:e - s]) return torch.cat(parts) def _backend() -> str: return str(_dist().get_backend()) def _reduce_scatter(out: Tensor, inp: Tensor) -> None: dist = _dist() if _backend() == "nccl": dist.reduce_scatter_tensor(out, inp) else: # gloo (CPU tests): all-reduce, keep this rank's piece dist.all_reduce(inp) r, n = dist.get_rank(), out.numel() out.copy_(inp[r * n:(r + 1) * n]) def _all_gather(out: Tensor, inp: Tensor) -> None: dist = _dist() if _backend() == "nccl": dist.all_gather_into_tensor(out, inp) else: dist.all_gather(list(out.chunk(dist.get_world_size())), inp) def sharded_static_bytes(trainable_backbone: int, heads: int, weights: int, world: int ) -> float: """:func:`deg.memplan.static_bytes` with ZeRO-2: bf16 weights + bf16 grads whole, fp32 master + AdamW moments + fp32 gradient of the backbone per shard, heads whole.""" return (weights + 2.0 * trainable_backbone + 16.0 * trainable_backbone / world + 16.0 * heads) def broadcast_module(module: nn.Module, src: int = 0) -> None: """Every parameter and buffer of ``module`` from rank ``src`` (fresh heads are drawn from per-rank seeds; replicas must start identical).""" dist = _dist() with torch.no_grad(): for t in [*module.parameters(), *module.buffers()]: dist.broadcast(t.data, src)