| """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 |
|
|
|
|
| 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) |
| |
| 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) |
| self.held: list[Tensor | None] | None = None |
| self.host: Tensor | None = None |
| self._host_buf: Tensor | None = None |
| self.spills = 0 |
|
|
| |
|
|
| 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 |
|
|
| |
|
|
| @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 |
|
|
| |
|
|
| @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: |
| 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: |
| 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) |
|
|
