From fdad1bd7afea2c253c56b5f888c6a8e845736de4 Mon Sep 17 00:00:00 2001 From: Oliver Clive-Griffin Date: Thu, 28 May 2026 20:37:31 +0000 Subject: [PATCH] three_pool: portals + typestate phases + typed pool-state union Restructure the 3-pool subsystem so invalid states / operations / order of operations are unrepresentable, behavior-preserving. - portals.py: each of the 6 cross-pool DAG edges (+ eval CIOutputs) is one typed class owning payload, routing, group, pack/unpack. Both sender and receiver invoke the same object, so the two sides cannot drift. Replaces the comm methods that were split across layout.py (sender) and the step files (receiver) with the pack format duplicated on each side. - step_{ci,layerwise,ppgd}.py: each step is now a sequence of phase functions returning small frozen bundles; later phases consume earlier phases' typed outputs, so dependency order is a type constraint. - pool_state.py: CIState | LWState | PPGDState discriminated union, matched exhaustively in optimize.py, replacing the optional-attr bag + my_pool string dispatch. - layout.py keeps World (topology + routing helpers) + the 3 in-pool reductions; cross-pool comm methods removed. - DESIGN.md: document the new structure + which invalid orders are now unrepresentable. (scripts/validate_nccl_event_timing.py: incidental pre-commit ruff autofix.) make check clean; 424 non-slow tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- param_decomp_lab/three_pool/DESIGN.md | 61 ++ param_decomp_lab/three_pool/eval_step.py | 15 +- param_decomp_lab/three_pool/layout.py | 568 +-------------- param_decomp_lab/three_pool/optimize.py | 156 ++-- param_decomp_lab/three_pool/pool_state.py | 51 ++ param_decomp_lab/three_pool/portals.py | 666 ++++++++++++++++++ param_decomp_lab/three_pool/step_ci.py | 210 ++++-- param_decomp_lab/three_pool/step_layerwise.py | 242 +++++-- param_decomp_lab/three_pool/step_ppgd.py | 230 ++++-- scripts/validate_nccl_event_timing.py | 1 - 10 files changed, 1378 insertions(+), 822 deletions(-) create mode 100644 param_decomp_lab/three_pool/pool_state.py create mode 100644 param_decomp_lab/three_pool/portals.py diff --git a/param_decomp_lab/three_pool/DESIGN.md b/param_decomp_lab/three_pool/DESIGN.md index f07df7716..e9970639b 100644 --- a/param_decomp_lab/three_pool/DESIGN.md +++ b/param_decomp_lab/three_pool/DESIGN.md @@ -5,6 +5,67 @@ and PPGD — so a **global shared transformer** CI fn is physically realizable: dedicated, replicated CI pool can host a CI fn that spans all sites, while V/U sites are sharded across the LW pool. +## Code structure (how the DAG maps to modules) + +The subsystem is built so the per-step loop reads as close to the dependency +DAG below as the SPMD form allows. Three layered abstractions carry that: + +1. **Cross-pool edges are first-class typed portals** (`portals.py`). Each of + the six DAG edges (plus the eval-only `CiOutputsToPpgd`) is ONE class that + owns its payload type, source/dest pool, batch-position routing, process + group, pack/unpack, and wire dtype. Both the sending and the receiving rank + construct the same portal (from the shared `World`) and call its `send` / + `post_recv` / `recv` side. Because send and recv live on one object, the two + sides' pack/unpack cannot drift — previously each edge was split across + `layout.py` (sender) and a step file (receiver) with the pack layout + duplicated in a docstring on each side. Async sends return a `SendHandle` + (kept alive until `.wait()`); posted recvs return a `PendingPerSiteCi` whose + `.wait()` is the only way to reach the payload. + +2. **Typed phase handoffs within each pool step** (`step_ci.py` / + `step_layerwise.py` / `step_ppgd.py`). Each step is a sequence of phase + functions returning small frozen bundles (`CiForward`, `CiSends`, + `ImpMinLoss`, `GciReceived`, `GciTotal`; `Faith`, `CiLeaves`, `Stoch`; + `ReconSum`, `RawGrads`). A phase can only run once its inputs exist as a + value, so the dependency order is a type constraint — see "Invalid orders + now unrepresentable" below. + +3. **Typed per-pool state union** (`pool_state.py`). `CIState | LWState | + PPGDState`, matched exhaustively in `optimize.py`, replaces the previous + optional-attr bag (`optimizer: ... | None`, `ppgd_state: ... | None`, + `_ci_fn_params`, `_component_params`, `_all_params`) + `match my_pool` string + dispatch. Each pool's state holds exactly the objects its role uses. + +`layout.py` keeps `World` (topology + process groups + batch-position routing +helpers), `ThreePoolLayout` (this rank's identity + per-rank batch-slice +helpers), and the three **in-pool collective reductions** (CI-fn-grad +all-reduce, LW in-block all-reduce, PPGD V/U sum-reduce) — those are within a +single pool, not cross-pool edges, so they aren't portals. + +### Invalid orders now unrepresentable + +* CI: cannot send CI before computing it (`send` takes `CiForward`); cannot + assemble `g_CI_total` before receiving both halves (`GciTotal` takes + `GciReceived`); cannot run the fused backward before both the imp-min loss + and the assembled grads exist (it takes `CiForward` + `ImpMinLoss` + + `GciTotal`); cannot wait the sends before posting them (`wait` is on the + `CiSends` the send phase returned). +* LW: cannot send `g_CI` before the streaming backward populated the re-leafed + CI tensors' `.grad` (the send consumes the `CiLeaves` those grads live on); + cannot read CI values before the posted recv is waited (`CiLeaves` is built + from `PendingPerSiteCi.wait()`). +* PPGD: cannot differentiate before the warmup-refined recon sum exists (grads + take `ReconSum`); cannot scale/send grads before they are produced (sends + take `RawGrads`); cannot step sources before differentiating. +* Pool role: a CI rank with no CI fn, a PPGD rank with an optimizer, etc. are + unrepresentable — each is a distinct dataclass in the `PoolState` union. + +What is NOT captured at the type level: the *relative ordering of independent +overlapped ops* (e.g. posting the async send before the dead-time prefetch, vs +after) is still expressed by statement order, since both are valid orderings +that only trade off overlap, not correctness. The portal `wait()` enforces the +send completes before its buffer is reused, but does not force *when* you wait. + ## Pool roles | Pool | Owns | Sharded? | Notes | diff --git a/param_decomp_lab/three_pool/eval_step.py b/param_decomp_lab/three_pool/eval_step.py index cd9cb82cd..ce45f5b96 100644 --- a/param_decomp_lab/three_pool/eval_step.py +++ b/param_decomp_lab/three_pool/eval_step.py @@ -35,6 +35,7 @@ from param_decomp.run_sink import RunSink from param_decomp.torch_helpers import bf16_autocast from param_decomp_lab.three_pool.layout import ThreePoolLayout +from param_decomp_lab.three_pool.portals import Portals def _slice_batch_dim0(batch: Any, sl: slice) -> tuple[Any, int]: @@ -58,6 +59,7 @@ def _build_metric_context_three_pool( batch: Any, *, layout: ThreePoolLayout, + portals: Portals, step: int, device: str, component_model: ComponentModel, @@ -71,6 +73,7 @@ def _build_metric_context_three_pool( batch = move_batch_to_device(batch, device) match layout.my_pool: case "ci": + assert layout.my_ci_slice_idx is not None batch_local, _ = _slice_batch_dim0(batch, layout.my_batch_slice_ci()) target_output = component_model(batch_local, cache_type="input") ci = component_model.calc_causal_importances( @@ -78,14 +81,18 @@ def _build_metric_context_three_pool( detach_inputs=False, sampling=config.sampling, ) - layout.send_ci_eval_to_ppgd(ci) + portals.ci_outputs_to_ppgd.send(ci, my_ci_slice_idx=layout.my_ci_slice_idx) return None case "ppgd": + assert layout.my_ppgd_slice_idx is not None batch_local, seq_len = _slice_batch_dim0(batch, layout.my_batch_slice_ppgd()) target_output = component_model(batch_local, cache_type="input") weight_deltas = component_model.calc_weight_deltas() - ci = layout.recv_ci_eval_from_ci_pool( - c_per_site, seq_len=seq_len, device=torch.device(device) + ci = portals.ci_outputs_to_ppgd.recv( + my_ppgd_slice_idx=layout.my_ppgd_slice_idx, + site_to_c=c_per_site, + seq_len=seq_len, + device=torch.device(device), ) return MetricContext( model=component_model, @@ -113,6 +120,7 @@ def run_eval_step( slow_step: bool, metrics: list[Metric[Any]], layout: ThreePoolLayout, + portals: Portals, step: int, device: str, component_model: ComponentModel, @@ -155,6 +163,7 @@ def run_eval_step( ctx = _build_metric_context_three_pool( batch, layout=layout, + portals=portals, step=step, device=device, component_model=component_model, diff --git a/param_decomp_lab/three_pool/layout.py b/param_decomp_lab/three_pool/layout.py index 80e286b52..1f4f6c1c6 100644 --- a/param_decomp_lab/three_pool/layout.py +++ b/param_decomp_lab/three_pool/layout.py @@ -1,28 +1,24 @@ -"""World / ThreePoolLayout — the 3-pool topology data model and cross-pool comms. +"""World / ThreePoolLayout — the 3-pool topology data model + in-pool reductions. `World` is purely declarative — identical content on every rank, no per-rank -fields. Built once at startup after `dist.init_process_group`. +fields. Built once at startup after `dist.init_process_group`. It owns the +process groups and the **batch-position routing** (the +``lw_sub_slice_within_ci`` / ``ci_slice_of_*`` bijection) that the cross-pool +portals consume. `ThreePoolLayout` wraps a World, adds this rank's perspective (`my_pool`, `my_owned_sites`, `my_within_block_idx` or `my_ci_slice_idx` or -`my_ppgd_slice_idx`), and hangs the cross-pool comm orchestration methods off -itself. - -Cross-pool exchanges (six total — see ``DESIGN.md`` for the per-step graph): - - CI → LW : CI_T per-site (owned + LW-rank batch slice) - CI → PPGD : CI_T full-model (per-PPGD-rank batch slice) - LW → CI : g_CI_LW per owned site (per-LW-rank batch slice) - PPGD→ CI : g_CI_PPGD full-model (per-PPGD-rank batch slice) - PPGD→ LW : g_VU_PPGD per-owned-site (after in-pool sum-reduce; PPGD-leader-driven) - LW → PPGD : updated V/U per-owned-site (LW-block-leader-driven, broadcast to PPGD pool) - -Plus three collective reductions: +`my_ppgd_slice_idx`), the per-rank batch-slice helpers, and the three in-pool +collective reductions: LW : in-block all-reduce on V/U + faithfulness grads (one per LW block group) CI : in-pool all-reduce on CI fn grads (one collective over the CI pool) PPGD: in-pool sum-reduce on V/U grads (one per site, over the PPGD pool) +The six **cross-pool point-to-point exchanges** live in +``param_decomp_lab.three_pool.portals`` — one typed portal object per DAG edge, +invoked from both the sending and receiving rank so pack/unpack cannot drift. + The defining wrinkle is **3-way batch slicing**: CI/LW/PPGD each shard the global batch on their own axis. The constraint (enforced in ``ThreePoolConfig.validate_topology``) is: @@ -50,14 +46,6 @@ from torch import Tensor from param_decomp._trace import trace -from param_decomp.component_model import CIOutputs - -# All cross-pool tensors are cast to this dtype on the wire (halves bytes vs fp32). -# Downstream pools run inside bf16 autocast already; CI grads and V/U grads -# accumulating into fp32 .grad upcast back to fp32 on receive — standard bf16 -# mixed-precision pattern. -_WIRE_DTYPE: torch.dtype = torch.bfloat16 - # ────────────────────────────────────────────────────────────────────────────── # NCCL-op event timing (PD_NCCL_EVENT_TIMING=1). @@ -121,38 +109,6 @@ def flush_nccl_event_timings() -> None: _NCCL_EVENT_BUFFER.clear() -@dataclass(frozen=True) -class PendingCiRecv: - """One coalesced CI-values irecv, held until ``wait_and_unpack()``. - - The packed buffer carries ``sites`` worth of CI values (in order) as - ``b * seq_len * c_s`` ``_WIRE_DTYPE`` elements each. ``wait_and_unpack`` - blocks on the underlying ``dist.Work`` then materializes per-site - ``[b, seq_len, c_s]`` views into the packed buffer (no copy). - """ - - packed: torch.Tensor - work: "dist.Work" - sites: tuple[str, ...] - site_to_c: dict[str, int] - b: int - seq_len: int - - def wait_and_unpack(self) -> dict[str, torch.Tensor]: - self.work.wait() - out: dict[str, torch.Tensor] = {} - offset = 0 - for s in self.sites: - c_s = self.site_to_c[s] - numel = self.b * self.seq_len * c_s - out[s] = self.packed[offset : offset + numel].view(self.b, self.seq_len, c_s) - offset += numel - assert offset == self.packed.numel(), ( - f"unpack size mismatch: consumed {offset} of {self.packed.numel()}" - ) - return out - - @dataclass(frozen=True) class LayerwiseBlockGroup: """One LW block-DDP group: ranks that replicate V/U for `owned_sites`. @@ -529,207 +485,10 @@ def i_lead_site(self, site: str) -> bool: return self.is_my_site(site) and self.my_is_block_leader # ────────────────────────────────────────────────────────────────────── - # CI-pool comm methods + # In-pool collective reductions (one per pool; not cross-pool edges). + # Cross-pool point-to-point exchanges live in ``portals.py``. # ────────────────────────────────────────────────────────────────────── - def async_send_ci_to_layerwise( - self, ci_full: dict[str, Tensor] - ) -> tuple[list["dist.Work"], list[Tensor]]: - """CI → LW: for each site and each LW rank whose batch shard sits in - my CI slice, isend the corresponding sub-slice. - - ``ci_full`` is keyed by site (CI fn produced CI for ALL sites since the - CI fn is global). Values have shape ``[B_local_ci, S, C_s]``. - Returned buffers must be kept alive until ``work.wait()`` completes. - """ - assert self.my_pool == "ci" and self.my_ci_slice_idx is not None - works: list[dist.Work] = [] - buffers: list[Tensor] = [] - my_lw_block_ranks = self.world.lw_block_ranks_for_ci_slice(self.my_ci_slice_idx) - - with _time_nccl_op("async_send_ci_to_layerwise"): - for bg in self.world.layerwise_block_groups: - for block_rank_idx in my_lw_block_ranks: - target_lw_rank = bg.ranks[block_rank_idx] - sub = self.world.lw_sub_slice_within_ci(block_rank_idx) - # Coalesce all of this block's owned-sites into one packed - # send per (block, block_rank). Layout (must match recv): - # for each site in bg.owned_sites order, b_lw * seq_len * C_s - # contiguous _WIRE_DTYPE elements. - parts = [ - ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() - for site in bg.owned_sites - ] - packed = torch.cat(parts) - works.append( - dist.isend( - packed, dst=target_lw_rank, group=self.world.cross_pool_p2p_group - ) - ) - buffers.append(packed) - return works, buffers - - def async_send_ci_to_ppgd( - self, ci_full: dict[str, Tensor] - ) -> tuple[list["dist.Work"], list[Tensor]]: - """CI → PPGD: for each PPGD rank whose batch shard sits in my CI slice, - isend the full-model CI sub-slice (all sites).""" - assert self.my_pool == "ci" and self.my_ci_slice_idx is not None - works: list[dist.Work] = [] - buffers: list[Tensor] = [] - my_ppgd_slice_idxs = self.world.ppgd_slice_idxs_for_ci_slice(self.my_ci_slice_idx) - - with _time_nccl_op("async_send_ci_to_ppgd"): - for ppgd_slice_idx in my_ppgd_slice_idxs: - target_ppgd_rank = self.world.ppgd_ranks[ppgd_slice_idx] - sub = self.world.ppgd_sub_slice_within_ci(ppgd_slice_idx) - # Coalesce all 96 sites into one packed send per PPGD target. - # Layout (must match recv): for each site in self.world.all_sites - # order, b_pp * seq_len * C_s contiguous _WIRE_DTYPE elements. - parts = [ - ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() - for site in self.world.all_sites - ] - packed = torch.cat(parts) - works.append( - dist.isend(packed, dst=target_ppgd_rank, group=self.world.cross_pool_p2p_group) - ) - buffers.append(packed) - return works, buffers - - def send_ci_eval_to_ppgd(self, ci: CIOutputs) -> None: - """CI → PPGD eval: synchronous send of full CIOutputs (all three dicts — - lower_leaky, upper_leaky, pre_sigmoid) sliced to each PPGD rank within - my CI slice. - - Training-time only ships ``lower_leaky``; eval ships all three so any - metric reading ``ctx.ci`` works without a per-metric audit. Synchronous - because eval is rare and overlap has no value here. - - Pack layout per send (must match ``recv_ci_eval_from_ci_pool``): three - contiguous blocks in order (lower_leaky, upper_leaky, pre_sigmoid). Each - block has, for each site in ``self.world.all_sites`` order, ``b_pp * - seq_len * C_s`` contiguous ``_WIRE_DTYPE`` elements. - """ - assert self.my_pool == "ci" and self.my_ci_slice_idx is not None - my_ppgd_slice_idxs = self.world.ppgd_slice_idxs_for_ci_slice(self.my_ci_slice_idx) - - with _time_nccl_op("send_ci_eval_to_ppgd"): - for ppgd_slice_idx in my_ppgd_slice_idxs: - target = self.world.ppgd_ranks[ppgd_slice_idx] - sub = self.world.ppgd_sub_slice_within_ci(ppgd_slice_idx) - parts: list[Tensor] = [] - for d in (ci.lower_leaky, ci.upper_leaky, ci.pre_sigmoid): - parts.extend( - d[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() - for site in self.world.all_sites - ) - packed = torch.cat(parts) - dist.send(packed, dst=target, group=self.world.cross_pool_p2p_group) - - def recv_g_ci_from_layerwise( - self, - site_to_c: dict[str, int], - seq_len: int, - device: torch.device, - ) -> dict[str, Tensor]: - """CI ← LW: recv per-site CI grads, coalesced per (LW block leader, LW - block rank index) channel. - - Each LW rank coalesces its owned sites into one packed buffer (see - ``send_g_ci_to_ci_pool``); we receive one packed buf per source. Pack - layout (must match sender): for each site in the LW block's owned - sites, ``b_lw * seq_len * c_s`` contiguous ``_WIRE_DTYPE`` elements. - """ - assert self.my_pool == "ci" and self.my_ci_slice_idx is not None - my_lw_block_ranks = self.world.lw_block_ranks_for_ci_slice(self.my_ci_slice_idx) - b_lw = self.world.batch_local_lw - - # Post all irecvs upfront so they pipeline on the NIC. - # Per source: one packed buf containing all of that source's owned sites. - pending: list[tuple[int, int, Tensor, dist.Work, tuple[str, ...]]] = [] - with _time_nccl_op("recv_g_ci_from_layerwise:post_irecvs"): - for bg_idx, bg in enumerate(self.world.layerwise_block_groups): - owned = bg.owned_sites - packed_numel = sum(b_lw * seq_len * site_to_c[s] for s in owned) - for block_rank_idx in my_lw_block_ranks: - src = bg.ranks[block_rank_idx] - buf = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) - w = dist.irecv(buf, src=src, group=self.world.cross_pool_p2p_group) - assert w is not None - pending.append((bg_idx, block_rank_idx, buf, w, owned)) - - # Wait + stitch. Allocate one fp32 dest per site, copy each piece in place. - out: dict[str, Tensor] = {} - b_ci = self.world.batch_local_ci - for site in self.world.all_sites: - c_s = site_to_c[site] - out[site] = torch.empty(b_ci, seq_len, c_s, device=device, dtype=torch.float32) - with _time_nccl_op("recv_g_ci_from_layerwise:wait"): - for _bg_idx, block_rank_idx, buf, w, owned in pending: - w.wait() - sub = self.world.lw_sub_slice_within_ci(block_rank_idx) - offset = 0 - for site in owned: - c_s = site_to_c[site] - n = b_lw * seq_len * c_s - site_view = buf[offset : offset + n].view(b_lw, seq_len, c_s) - out[site][sub].copy_(site_view.to(torch.float32)) - offset += n - return out - - def recv_g_ci_from_ppgd( - self, - site_to_c: dict[str, int], - seq_len: int, - device: torch.device, - ) -> dict[str, Tensor]: - """CI ← PPGD: recv full-model CI grads, coalesced. - - One packed irecv per PPGD source (instead of one per (site, source)), - matching the coalesced ``send_g_ci_to_ci_pool_ppgd``. With 96 sites - and N_PPGD/N_CI=3 sources per CI rank, that's 3 irecvs instead of - 288 — order-of-magnitude NCCL-launch latency cut. - - Pack layout (must match the sender exactly): for each site in - ``self.world.all_sites`` order, ``b_pp * seq_len * c_s`` contiguous - ``_WIRE_DTYPE`` elements. - """ - assert self.my_pool == "ci" and self.my_ci_slice_idx is not None - my_ppgd_slice_idxs = self.world.ppgd_slice_idxs_for_ci_slice(self.my_ci_slice_idx) - b_pp = self.world.batch_local_ppgd - - # Same total numel for every PPGD source (every source sends all sites). - site_numels = {s: b_pp * seq_len * site_to_c[s] for s in self.world.all_sites} - packed_numel = sum(site_numels.values()) - - pending: list[tuple[int, Tensor, dist.Work]] = [] - with _time_nccl_op("recv_g_ci_from_ppgd:post_irecvs"): - for ppgd_slice_idx in my_ppgd_slice_idxs: - src = self.world.ppgd_ranks[ppgd_slice_idx] - packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) - w = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) - assert w is not None - pending.append((ppgd_slice_idx, packed, w)) - - b_ci = self.world.batch_local_ci - out: dict[str, Tensor] = { - s: torch.empty(b_ci, seq_len, site_to_c[s], device=device, dtype=torch.float32) - for s in self.world.all_sites - } - with _time_nccl_op("recv_g_ci_from_ppgd:wait"): - for ppgd_slice_idx, packed, w in pending: - w.wait() - sub = self.world.ppgd_sub_slice_within_ci(ppgd_slice_idx) - offset = 0 - for site in self.world.all_sites: - c_s = site_to_c[site] - n = site_numels[site] - buf = packed[offset : offset + n].view(b_pp, seq_len, c_s) - out[site][sub].copy_(buf.to(torch.float32)) - offset += n - return out - def all_reduce_ci_fn_grads(self, params: Iterable[nn.Parameter]) -> None: """In-pool all-reduce on CI fn grads. Coalesced bucketed reduce — same pattern as ``all_reduce_grads_in_block`` but over the CI pool. @@ -754,139 +513,6 @@ def all_reduce_ci_fn_grads(self, params: Iterable[nn.Parameter]) -> None: ): orig.copy_(reduced) - # ────────────────────────────────────────────────────────────────────── - # Layerwise-pool comm methods - # ────────────────────────────────────────────────────────────────────── - - def async_recv_ci_from_ci_pool( - self, - site_to_c: dict[str, int], - seq_len: int, - device: torch.device, - ) -> PendingCiRecv: - """LW ← CI: irecv one coalesced packet of CI values for all of this - LW rank's owned sites, from the CI rank whose slice contains my LW - batch shard. - - Layout (must match ``async_send_ci_to_layerwise``): for each site in - ``self.my_owned_sites`` order, ``b_lw * seq_len * C_s`` contiguous - ``_WIRE_DTYPE`` elements. Caller calls ``wait_and_unpack()`` to get - per-site ``[b_lw, seq_len, C_s]`` views (no copy). - """ - assert self.my_pool == "layerwise" and self.my_within_block_idx is not None - src_ci_slice = self.world.ci_slice_of_lw_block_rank(self.my_within_block_idx) - src = self.world.ci_ranks[src_ci_slice] - b_lw = self.world.batch_local_lw - - packed_numel = sum(b_lw * seq_len * site_to_c[s] for s in self.my_owned_sites) - packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) - with _time_nccl_op("async_recv_ci_from_ci_pool"): - work = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) - assert work is not None - return PendingCiRecv( - packed=packed, - work=work, - sites=self.my_owned_sites, - site_to_c=site_to_c, - b=b_lw, - seq_len=seq_len, - ) - - def send_g_ci_to_ci_pool(self, g_ci_owned: dict[str, Tensor]) -> None: - """LW → CI: send per-owned-site CI grads (full LW batch slice) to the - CI rank that owns my slice. - - Coalesces this rank's owned sites into one packed send (vs one isend - per site). Smaller win than the PPGD-side coalescing (each LW rank - only owns ~4 sites vs PPGD's 96) but consistent + cuts CI's recv - count from ~96 to ~24 (one per LW block, not per (site, block)). - """ - assert self.my_pool == "layerwise" and self.my_within_block_idx is not None - dst_ci_slice = self.world.ci_slice_of_lw_block_rank(self.my_within_block_idx) - dst = self.world.ci_ranks[dst_ci_slice] - parts = [ - g_ci_owned[s].detach().to(_WIRE_DTYPE).contiguous().flatten() - for s in self.my_owned_sites - ] - packed = torch.cat(parts) - with _time_nccl_op("send_g_ci_to_ci_pool"): - dist.send(packed, dst=dst, group=self.world.cross_pool_p2p_group) - - def recv_g_vu_from_ppgd( - self, - v_templates: dict[str, Tensor], - u_templates: dict[str, Tensor], - ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: - """LW ← PPGD: leader recvs g_VU for owned sites from PPGD leader, then - in-block broadcast. PPGD has already sum-reduced within its pool so a - single recv carries the full-batch grad for our owned sites. - """ - assert self.my_pool == "layerwise" and self.my_block_idx is not None - v_grads: dict[str, Tensor] = {} - u_grads: dict[str, Tensor] = {} - - if self.my_is_block_leader: - my_sites = self.my_owned_sites - packed_numel = sum(v_templates[s].numel() + u_templates[s].numel() for s in my_sites) - sample = v_templates[my_sites[0]] - packed = torch.empty(packed_numel, dtype=_WIRE_DTYPE, device=sample.device) - ppgd_leader = self.world.ppgd_ranks[0] - with _time_nccl_op("recv_g_vu_from_ppgd:recv"): - dist.recv(packed, src=ppgd_leader, group=self.world.cross_pool_p2p_group) - offset = 0 - for s in my_sites: - v_n = v_templates[s].numel() - u_n = u_templates[s].numel() - v_grads[s] = ( - packed[offset : offset + v_n].view_as(v_templates[s]).to(v_templates[s].dtype) - ) - offset += v_n - u_grads[s] = ( - packed[offset : offset + u_n].view_as(u_templates[s]).to(u_templates[s].dtype) - ) - offset += u_n - else: - for s in self.my_owned_sites: - v_grads[s] = torch.empty_like(v_templates[s]) - u_grads[s] = torch.empty_like(u_templates[s]) - - # In-block broadcast leader → other ranks so all replicas see the same g_VU. - block_group = self.world.block_group_groups[self.my_block_idx] - block_leader_rank = self.world.layerwise_block_groups[self.my_block_idx].leader - with _time_nccl_op("recv_g_vu_from_ppgd:in_block_bcast"): - for s in self.my_owned_sites: - v_grads[s] = v_grads[s].contiguous() - u_grads[s] = u_grads[s].contiguous() - dist.broadcast(v_grads[s], src=block_leader_rank, group=block_group) - dist.broadcast(u_grads[s], src=block_leader_rank, group=block_group) - - return v_grads, u_grads - - def async_send_updated_vu_to_ppgd( - self, - v_owned: dict[str, Tensor], - u_owned: dict[str, Tensor], - ) -> tuple[list["dist.Work"], list[Tensor]]: - """LW → PPGD: coalesced leader-rooted broadcast of updated V/U to all - PPGD ranks. Caller must keep the buffer alive until the work handle - completes. - """ - assert self.my_pool == "layerwise" - if not self.my_is_block_leader: - return [], [] - assert self.my_block_idx is not None - my_sites = self.my_owned_sites - parts: list[Tensor] = [] - for s in my_sites: - parts.append(v_owned[s].detach().to(_WIRE_DTYPE).contiguous().flatten()) - parts.append(u_owned[s].detach().to(_WIRE_DTYPE).contiguous().flatten()) - packed = torch.cat(parts) - bcast_group = self.world.cross_pool_bcast_groups[self.my_block_idx] - with _time_nccl_op("async_send_updated_vu_to_ppgd"): - w = dist.broadcast(packed, src=self.my_rank, group=bcast_group, async_op=True) - assert w is not None - return [w], [packed] - def all_reduce_grads_in_block(self, params: Iterable[nn.Parameter]) -> None: """Coalesced in-block DDP all-reduce over V/U + faithfulness grads. @@ -919,131 +545,6 @@ def all_reduce_grads_in_block(self, params: Iterable[nn.Parameter]) -> None: for orig, reduced in zip(bucket, _unflatten_dense_tensors(flat, bucket), strict=True): orig.copy_(reduced) - # ────────────────────────────────────────────────────────────────────── - # PPGD-pool comm methods - # ────────────────────────────────────────────────────────────────────── - - def async_recv_ci_from_ci_pool_ppgd( - self, - site_to_c: dict[str, int], - seq_len: int, - device: torch.device, - ) -> PendingCiRecv: - """PPGD ← CI: irecv one coalesced packet of full-model CI values from - the CI rank that owns my slice. - - Layout (must match ``async_send_ci_to_ppgd``): for each site in - ``self.world.all_sites`` order, ``b_pp * seq_len * C_s`` contiguous - ``_WIRE_DTYPE`` elements. Caller calls ``wait_and_unpack()`` to get - per-site ``[b_pp, seq_len, C_s]`` views (no copy). - """ - assert self.my_pool == "ppgd" and self.my_ppgd_slice_idx is not None - src_ci_slice = self.world.ci_slice_of_ppgd_slice(self.my_ppgd_slice_idx) - src = self.world.ci_ranks[src_ci_slice] - b_pp = self.world.batch_local_ppgd - - packed_numel = sum(b_pp * seq_len * site_to_c[s] for s in self.world.all_sites) - packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) - with _time_nccl_op("async_recv_ci_from_ci_pool_ppgd"): - work = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) - assert work is not None - return PendingCiRecv( - packed=packed, - work=work, - sites=self.world.all_sites, - site_to_c=site_to_c, - b=b_pp, - seq_len=seq_len, - ) - - def recv_ci_eval_from_ci_pool( - self, - site_to_c: dict[str, int], - seq_len: int, - device: torch.device, - ) -> CIOutputs: - """PPGD ← CI eval: synchronous recv of full ``CIOutputs`` from the CI - rank that owns my slice. - - Pack layout (must match ``send_ci_eval_to_ppgd``): three contiguous - blocks in order (lower_leaky, upper_leaky, pre_sigmoid). Each block has, - for each site in ``self.world.all_sites`` order, ``b_pp * seq_len * - C_s`` contiguous ``_WIRE_DTYPE`` elements. Returned dicts are no-copy - views into the packed buffer. - """ - assert self.my_pool == "ppgd" and self.my_ppgd_slice_idx is not None - src_ci_slice = self.world.ci_slice_of_ppgd_slice(self.my_ppgd_slice_idx) - src = self.world.ci_ranks[src_ci_slice] - b_pp = self.world.batch_local_ppgd - - per_block_numel = sum(b_pp * seq_len * site_to_c[s] for s in self.world.all_sites) - packed = torch.empty(3 * per_block_numel, device=device, dtype=_WIRE_DTYPE) - with _time_nccl_op("recv_ci_eval_from_ci_pool"): - dist.recv(packed, src=src, group=self.world.cross_pool_p2p_group) - - out: list[dict[str, Tensor]] = [{}, {}, {}] - offset = 0 - for block_idx in range(3): - for site in self.world.all_sites: - c_s = site_to_c[site] - numel = b_pp * seq_len * c_s - out[block_idx][site] = packed[offset : offset + numel].view(b_pp, seq_len, c_s) - offset += numel - assert offset == packed.numel(), f"unpack mismatch: {offset} of {packed.numel()}" - return CIOutputs(lower_leaky=out[0], upper_leaky=out[1], pre_sigmoid=out[2]) - - def send_g_ci_to_ci_pool_ppgd(self, g_ci_full: dict[str, Tensor]) -> None: - """PPGD → CI: send full-model CI grads (PPGD batch slice) to the CI - rank that owns my slice. - - Coalesces all 96-ish sites into a single packed buffer per - send. Per-site isends launch ~10ms of NCCL overhead each, so at - scale (96 sites × N_PPGD ranks) this phase was ~1 s of pure NCCL - launch latency on every step. Single packed send replaces that with - one NCCL op (NCCL handles big tensors efficiently). - """ - assert self.my_pool == "ppgd" and self.my_ppgd_slice_idx is not None - dst_ci_slice = self.world.ci_slice_of_ppgd_slice(self.my_ppgd_slice_idx) - dst = self.world.ci_ranks[dst_ci_slice] - parts = [ - g_ci_full[s].detach().to(_WIRE_DTYPE).contiguous().flatten() - for s in self.world.all_sites - ] - packed = torch.cat(parts) - with _time_nccl_op("send_g_ci_to_ci_pool_ppgd"): - dist.send(packed, dst=dst, group=self.world.cross_pool_p2p_group) - - def send_g_vu_to_layerwise( - self, - v_grads: dict[str, Tensor], - u_grads: dict[str, Tensor], - ) -> None: - """PPGD-leader-only: send g_VU per-block (coalesced) to each LW block leader. - - Assumes V/U grads have already been sum-reduced within the PPGD pool — - every PPGD rank holds the same values, so only the leader sends. - """ - assert self.my_pool == "ppgd" - if not self.my_is_pool_leader: - return - works: list[dist.Work] = [] - buffers: list[Tensor] = [] - with _time_nccl_op("send_g_vu_to_layerwise:isends"): - for bg in self.world.layerwise_block_groups: - parts: list[Tensor] = [] - for site in bg.owned_sites: - parts.append(v_grads[site].to(_WIRE_DTYPE).contiguous().flatten()) - parts.append(u_grads[site].to(_WIRE_DTYPE).contiguous().flatten()) - packed = torch.cat(parts) - w = dist.isend(packed, dst=bg.leader, group=self.world.cross_pool_p2p_group) - assert w is not None - works.append(w) - buffers.append(packed) - with _time_nccl_op("send_g_vu_to_layerwise:wait"): - for w in works: - w.wait() - del buffers - def sum_reduce_ppgd_grads(self, grads: Iterable[Tensor]) -> None: """In-pool sum-reduce on PPGD V/U grads. Caller passes a flat iterable of tensors; each is all-reduced in place over the PPGD pool group. @@ -1069,46 +570,3 @@ def sum_reduce_ppgd_grads(self, grads: Iterable[Tensor]) -> None: bucket, _unflatten_dense_tensors(flat, bucket), strict=True ): orig.copy_(reduced) - - def recv_updated_vu_from_layerwise( - self, - v_templates: dict[str, Tensor], - u_templates: dict[str, Tensor], - ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: - """PPGD ← LW: coalesced + pipelined recv of updated V/U from each LW block leader. - - Kicks off one async broadcast per block group (they pipeline across the - per-group NCCL streams), then waits + unpacks each contiguous packet - back into per-site V/U dicts (upcasting to the templates' dtype). - Returns ``(v_new, u_new)`` ready for ``components[s].V.copy_(...)``. - """ - assert self.my_pool == "ppgd" - bufs: list[tuple[LayerwiseBlockGroup, Tensor, dist.Work]] = [] - with _time_nccl_op("recv_updated_vu_from_layerwise"): - for bg_idx, bg in enumerate(self.world.layerwise_block_groups): - owned = bg.owned_sites - packed_numel = sum(v_templates[s].numel() + u_templates[s].numel() for s in owned) - sample = v_templates[owned[0]] - packed = torch.empty(packed_numel, dtype=_WIRE_DTYPE, device=sample.device) - bcast_group = self.world.cross_pool_bcast_groups[bg_idx] - w = dist.broadcast(packed, src=bg.leader, group=bcast_group, async_op=True) - assert w is not None - bufs.append((bg, packed, w)) - - v_new: dict[str, Tensor] = {} - u_new: dict[str, Tensor] = {} - for bg, packed, w in bufs: - w.wait() - offset = 0 - for s in bg.owned_sites: - v_n = v_templates[s].numel() - u_n = u_templates[s].numel() - v_new[s] = ( - packed[offset : offset + v_n].view_as(v_templates[s]).to(v_templates[s].dtype) - ) - offset += v_n - u_new[s] = ( - packed[offset : offset + u_n].view_as(u_templates[s]).to(u_templates[s].dtype) - ) - offset += u_n - return v_new, u_new diff --git a/param_decomp_lab/three_pool/optimize.py b/param_decomp_lab/three_pool/optimize.py index c43a3223b..5fb9f0a87 100644 --- a/param_decomp_lab/three_pool/optimize.py +++ b/param_decomp_lab/three_pool/optimize.py @@ -86,6 +86,8 @@ class boundary. flush_nccl_event_timings, ) from param_decomp_lab.three_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp_lab.three_pool.pool_state import CIState, LWState, PoolState, PPGDState +from param_decomp_lab.three_pool.portals import Portals from param_decomp_lab.three_pool.reductions import ( aggregate_losses_to_rank0, aggregate_max_memory_to_rank0, @@ -153,9 +155,9 @@ class ThreePoolTrainer: reconstruction_loss: ReconstructionLoss component_model: ComponentModel layout: ThreePoolLayout + portals: Portals strategy: LayerwiseLossStrategy - optimizer: torch.optim.Optimizer | None - ppgd_state: PersistentPGDState | None + pool_state: PoolState step: int def __init__( @@ -235,6 +237,7 @@ def __init__( ) trace("ThreePoolTrainer.__init__: build_world: done") self.layout = ThreePoolLayout.from_world(world, dist.get_rank()) + self.portals = Portals.from_world(world) decomposition_targets = _decomposition_targets_for_pool( self.layout, self.runtime.c_per_site ) @@ -310,49 +313,43 @@ def __init__( ) trace("ThreePoolTrainer.__init__: LayerwiseLossStrategy: done") - self.optimizer = None - self._all_params: list[nn.Parameter] = [] - self._ci_fn_params: list[nn.Parameter] = [] - self._component_params: list[nn.Parameter] = [] - self.ppgd_state = None - self._pending_ppgd_resume_state: dict[str, Any] | None = None - - trace(f"ThreePoolTrainer.__init__: optimizer build: enter (pool={self.layout.my_pool})") + trace(f"ThreePoolTrainer.__init__: pool state build: enter (pool={self.layout.my_pool})") match self.layout.my_pool: case "ci": assert self.component_model.ci_fn is not None, "CI pool must keep its CI fn" - self._ci_fn_params = list(self.component_model.ci_fn.parameters()) - n_params = sum(p.numel() for p in self._ci_fn_params) + ci_fn_params = list(self.component_model.ci_fn.parameters()) + n_params = sum(p.numel() for p in ci_fn_params) trace(f"ThreePoolTrainer.__init__: CI fn params={n_params / 1e9:.3f}B") - self.optimizer = torch.optim.AdamW( + ci_optimizer = torch.optim.AdamW( [ { - "params": self._ci_fn_params, + "params": ci_fn_params, "lr": pd_config.ci_fn_optimizer.lr_schedule.start_val, } ], weight_decay=0.0, fused=True, ) + self.pool_state = CIState(optimizer=ci_optimizer, ci_fn_params=ci_fn_params) case "layerwise": + component_params: list[nn.Parameter] = [] for name in self.layout.my_owned_sites: - self._component_params.extend( - self.component_model.components[name].parameters() - ) - self._all_params = self._component_params - self.optimizer = torch.optim.AdamW( + component_params.extend(self.component_model.components[name].parameters()) + lw_optimizer = torch.optim.AdamW( [ { - "params": self._component_params, + "params": component_params, "lr": pd_config.components_optimizer.lr_schedule.start_val, } ], weight_decay=0.0, fused=True, ) + self.pool_state = LWState(optimizer=lw_optimizer, component_params=component_params) case "ppgd": - pass # ppgd_state constructed lazily from first batch in run() - trace("ThreePoolTrainer.__init__: optimizer build: done") + # ppgd_state constructed lazily from first batch in run(). + self.pool_state = PPGDState() + trace("ThreePoolTrainer.__init__: pool state build: done") dump_memory_stats("after optimizer build") trace("ThreePoolTrainer.__init__: exit") @@ -418,19 +415,21 @@ def snapshot(self, scratch_dir: Path) -> ThreePoolTrainingState | None: trace("snapshot: gather_full_state_dict_to_rank0 done") my_named_params = self._named_params_for_my_optimizer() - my_optimizer_by_name: dict[str, dict[str, Any]] = ( - optimizer_state_by_name(self.optimizer, my_named_params) - if self.optimizer is not None - else {} - ) my_contribution: dict[str, Any] = { "pool": self.layout.my_pool, - "optimizer_by_name": my_optimizer_by_name, + "optimizer_by_name": {}, } - if self.ppgd_state is not None: - my_contribution["ppgd"] = self.ppgd_state.state_dict() - elif self._pending_ppgd_resume_state is not None: - my_contribution["ppgd"] = self._pending_ppgd_resume_state + match self.pool_state: + case CIState() | LWState(): + my_contribution["optimizer_by_name"] = optimizer_state_by_name( + self.pool_state.optimizer, my_named_params + ) + case PPGDState(ppgd_state=ppgd) if ppgd is not None: + my_contribution["ppgd"] = ppgd.state_dict() + case PPGDState(pending_resume_state=pending) if pending is not None: + my_contribution["ppgd"] = pending + case PPGDState(): + pass # File-based gather rather than dist.gather_object: at XL the aggregate # pickle payload (LW optimizer state + PPGD sources) is ~8 GB and @@ -545,18 +544,20 @@ def _load_canonical_state(self, state: ThreePoolTrainingState) -> None: local_slice = {k: v for k, v in state.component_model.items() if k in local_model_keys} self.component_model.load_state_dict(local_slice, strict=False) - if self.optimizer is not None: - named_params = self._named_params_for_my_optimizer() - match self.layout.my_pool: - case "layerwise": - by_name = state.components_optimizer - case "ci": - by_name = state.ci_fn_optimizer - case _: - by_name = {} - load_optimizer_state_by_name(self.optimizer, named_params, by_name) - if self.layout.my_pool == "ppgd": - self._pending_ppgd_resume_state = state.ppgd_state_by_rank.get(self.layout.my_rank) + named_params = self._named_params_for_my_optimizer() + match self.pool_state: + case CIState(): + load_optimizer_state_by_name( + self.pool_state.optimizer, named_params, state.ci_fn_optimizer + ) + case LWState(): + load_optimizer_state_by_name( + self.pool_state.optimizer, named_params, state.components_optimizer + ) + case PPGDState(): + self.pool_state.pending_resume_state = state.ppgd_state_by_rank.get( + self.layout.my_rank + ) # ============================ Training loop ============================ @@ -611,7 +612,7 @@ def run( train_iterator = itertools.chain([first_batch], train_iterator) _assert_full_global_batch(first_batch, runtime.batch_global) - if layout.my_pool == "ppgd" and self.ppgd_state is None: + if isinstance(self.pool_state, PPGDState) and self.pool_state.ppgd_state is None: trace("Trainer.run: PPGDState ctor: enter") ppgd_cfg = runtime.ppgd_cfg # The 3-pool currently only supports per-batch-per-position sources: @@ -624,7 +625,7 @@ def run( f"{type(ppgd_cfg.scope).__name__}. Replicated scopes need cross-pool " f"source-replica sync, not implemented in the 3-pool." ) - self.ppgd_state = PersistentPGDState( + self.pool_state.ppgd_state = PersistentPGDState( module_to_c=runtime.c_per_site, batch_dims=(layout.world.batch_local_ppgd, *_seq_dims_from_batch(first_batch)), device=device, @@ -637,14 +638,14 @@ def run( router=AllLayersRouter(), reconstruction_loss=self.strategy.recon_loss, ) - if self._pending_ppgd_resume_state is not None: - self.ppgd_state.load_state_dict(self._pending_ppgd_resume_state) - self._pending_ppgd_resume_state = None + if self.pool_state.pending_resume_state is not None: + self.pool_state.ppgd_state.load_state_dict(self.pool_state.pending_resume_state) + self.pool_state.pending_resume_state = None trace("Trainer.run: PPGDState ctor: done") if ( self.step == 0 - and layout.my_pool == "layerwise" + and isinstance(self.pool_state, LWState) and pd_config.faithfulness_warmup_steps > 0 ): trace( @@ -652,7 +653,7 @@ def run( ) run_faithfulness_warmup_layerwise( component_model=self.component_model, - component_params=self._component_params, + component_params=self.pool_state.component_params, n_steps=pd_config.faithfulness_warmup_steps, lr=pd_config.faithfulness_warmup_lr, weight_decay=pd_config.faithfulness_warmup_weight_decay, @@ -695,17 +696,19 @@ def _to_device(b: Any) -> Any: _assert_full_global_batch(batch_T, runtime.batch_global) trace(f"Trainer.run: step {step}: start (pool={layout.my_pool})") - if self.optimizer is not None: - # CI pool: one param group (CI fn); LW pool: one (components). - # PPGD pool has no optimizer. - if layout.my_pool == "ci": - self.optimizer.param_groups[0]["lr"] = get_scheduled_value( + # CI pool: one param group (CI fn); LW pool: one (components). + # PPGD pool has no optimizer to schedule. + match self.pool_state: + case CIState(): + self.pool_state.optimizer.param_groups[0]["lr"] = get_scheduled_value( step, n_steps, ci_fn_lr_schedule ) - elif layout.my_pool == "layerwise": - self.optimizer.param_groups[0]["lr"] = get_scheduled_value( + case LWState(): + self.pool_state.optimizer.param_groups[0]["lr"] = get_scheduled_value( step, n_steps, components_lr_schedule ) + case PPGDState(): + pass step_start = time.perf_counter() should_log = step % cadence.train_log_every == 0 @@ -715,20 +718,18 @@ def _to_device(b: Any) -> Any: assert batch_T.device == device, ( f"3-pool batch device drift at step {step}: {batch_T.device} vs {device}" ) - match layout.my_pool: - case "ci": - assert self.optimizer is not None, ( - f"CI rank {layout.my_rank} missing optimizer" - ) - assert len(self._ci_fn_params) > 0, ( + match self.pool_state: + case CIState(optimizer=opt, ci_fn_params=ci_fn_params): + assert len(ci_fn_params) > 0, ( f"CI rank {layout.my_rank} has no ci_fn params to optimize" ) next_batch_for_prefetch = batch_T_plus_1 if step < n_steps - 1 else None metrics, h_cache_ci = step_ci( layout, + self.portals, self.component_model, - self.optimizer, - self._ci_fn_params, + opt, + ci_fn_params, batch_T=batch_T, batch_T_plus_1=next_batch_for_prefetch, h_cache_T=h_cache_ci, @@ -736,31 +737,30 @@ def _to_device(b: Any) -> Any: current_frac_of_training=step / n_steps if n_steps > 0 else 0.0, should_log=should_log, ) - case "layerwise": - assert self.optimizer is not None, ( - f"LW rank {layout.my_rank} missing optimizer" - ) + case LWState(optimizer=opt, component_params=component_params): assert layout.my_owned_sites, ( f"LW rank {layout.my_rank} has no owned_sites — empty block" ) metrics = step_layerwise( layout, + self.portals, self.component_model, - self.optimizer, - self._all_params, + opt, + component_params, batch_T, runtime, self.strategy, should_log=should_log, ) - case "ppgd": - assert self.ppgd_state is not None, ( + case PPGDState(ppgd_state=ppgd): + assert ppgd is not None, ( f"PPGD rank {layout.my_rank} has no ppgd_state — lazy init failed" ) metrics = step_ppgd( layout, + self.portals, self.component_model, - self.ppgd_state, + ppgd, batch_T, runtime, self.strategy, @@ -790,6 +790,9 @@ def _to_device(b: Any) -> Any: dump_memory_stats(f"step {step} done") if step % cadence.train_log_every == 0: + lw_optimizer = ( + self.pool_state.optimizer if isinstance(self.pool_state, LWState) else None + ) _log_train_metrics( metrics=metrics, layout=layout, @@ -797,7 +800,7 @@ def _to_device(b: Any) -> Any: step=step, step_ms=step_ms, runtime=runtime, - optimizer=self.optimizer, + optimizer=lw_optimizer, sink=sink, ) @@ -811,6 +814,7 @@ def _to_device(b: Any) -> Any: slow_step=eval_loop.should_run_slow_eval(step), metrics=list(eval_loop.metrics), layout=layout, + portals=self.portals, step=step, device=str(device), component_model=self.component_model, diff --git a/param_decomp_lab/three_pool/pool_state.py b/param_decomp_lab/three_pool/pool_state.py new file mode 100644 index 000000000..5ce5a8a7f --- /dev/null +++ b/param_decomp_lab/three_pool/pool_state.py @@ -0,0 +1,51 @@ +"""Typed per-pool training state — replaces the optional-attr bag + string dispatch. + +Each rank plays exactly one of three pool roles, and each role holds a +genuinely different set of mutable training objects: + + * ``CIState`` — the CI fn's optimizer + its parameter list. No V/U, no PPGD. + * ``LWState`` — the components' optimizer + the owned-site parameter list. + * ``PPGDState`` — neither optimizer nor V/U params of its own; it carries the + persistent adversarial-source state (built lazily on the first batch). + +Modelling these as a discriminated union (rather than a single object with +``optimizer: Optimizer | None``, ``ppgd_state: ... | None``, ``ci_fn_params``, +``component_params`` all hanging off it) makes the per-pool variation explicit: +``match pool_state`` is exhaustive, and a phase can't reach for an attribute the +current pool doesn't have. "CI pool with no ci_fn", "PPGD pool with an +optimizer", etc. become unrepresentable. +""" + +from dataclasses import dataclass, field +from typing import Any + +import torch.nn as nn +from torch.optim import Optimizer + +from param_decomp.metrics.persistent_pgd_state import PersistentPGDState + + +@dataclass +class CIState: + optimizer: Optimizer + ci_fn_params: list[nn.Parameter] + + +@dataclass +class LWState: + optimizer: Optimizer + component_params: list[nn.Parameter] + + +@dataclass +class PPGDState: + """PPGD has no optimizer and no owned V/U params. ``ppgd_state`` is built + lazily on the first batch (its source shapes depend on the data's seq dims); + ``pending_resume_state`` carries a resumed source state dict until then. + """ + + ppgd_state: PersistentPGDState | None = None + pending_resume_state: dict[str, Any] | None = field(default=None) + + +PoolState = CIState | LWState | PPGDState diff --git a/param_decomp_lab/three_pool/portals.py b/param_decomp_lab/three_pool/portals.py new file mode 100644 index 000000000..a73f17fc1 --- /dev/null +++ b/param_decomp_lab/three_pool/portals.py @@ -0,0 +1,666 @@ +"""Cross-pool exchanges as first-class typed portal objects. + +Each of the six per-step cross-pool edges in the 3-pool dependency graph +(see ``DESIGN.md``) is defined here exactly ONCE. A portal owns everything +the edge needs — its payload shape, source/dest pool, the batch-position +routing (the ``lw_sub_slice_within_ci`` / ``ci_slice_of_*`` bijection), its +process group, its pack/unpack, and its bf16 wire dtype. Both the sending +rank and the receiving rank construct the SAME portal object (from the +shared ``World``) and invoke it from their respective sides: + + handle = portal.send(payload) # sender side; later handle.wait() + pending = portal.post_recv(...) # receiver side; later pending.wait() -> T + +Because send and recv live on one object, the two sides' pack/unpack layout +cannot drift — the previous design split each edge across ``layout.py`` +(sender) and the receiving step file, with the pack format duplicated in a +docstring on each side. + +The six edges (sender → receiver): + + ``CiValuesToLayerwise`` CI → LW : CI_T per-site (owned + LW-rank slice) + ``CiValuesToPpgd`` CI → PPGD : CI_T full-model (per-PPGD-rank slice) + ``GradCiFromLayerwise`` LW → CI : g_CI_LW per owned site (per-LW-rank slice) + ``GradCiFromPpgd`` PPGD → CI : g_CI_PPGD full-model (per-PPGD-rank slice) + ``GradVuFromPpgd`` PPGD → LW : g_VU_PPGD per owned site (post in-pool reduce) + ``UpdatedVuToPpgd`` LW → PPGD : updated V/U per owned site (leader broadcast) + +Plus the eval-only ``CiOutputsToPpgd`` (CI → PPGD; full ``CIOutputs``). + +The three in-pool collective reductions (CI-fn-grad all-reduce, LW in-block +all-reduce, PPGD V/U sum-reduce) are NOT cross-pool edges; they stay as +methods on ``ThreePoolLayout``. + +All cross-pool tensors are cast to ``_WIRE_DTYPE`` (bf16) on the wire — half +the bytes vs fp32. Downstream pools run inside bf16 autocast; received grads +upcast to fp32 on receive (standard mixed-precision pattern). +""" + +# pyright: reportIndexIssue=false, reportArgumentType=false, reportOperatorIssue=false + +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import Tensor + +from param_decomp.component_model import CIOutputs +from param_decomp_lab.three_pool.layout import World, _time_nccl_op + +# All cross-pool tensors are cast to this dtype on the wire (halves bytes vs fp32). +_WIRE_DTYPE: torch.dtype = torch.bfloat16 + + +# ────────────────────────────────────────────────────────────────────────────── +# Handles — typed deferral wrappers returned by portal send/recv. +# +# A ``SendHandle`` keeps the packed send buffers alive until ``wait()``; a +# ``Pending[T]`` blocks on its work then unpacks the wire buffer into the +# portal's typed payload. Receiver code cannot reach the payload without +# calling ``wait()``, so "use before the transfer completes" is unrepresentable. +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class SendHandle: + """In-flight cross-pool sends + the buffers backing them. + + The buffers must stay referenced until every send completes, so they ride + along on the handle. ``wait()`` blocks on all of them; a handle with no + works (e.g. a non-leader rank that sends nothing) is a no-op wait. + """ + + works: list["dist.Work"] + buffers: list[Tensor] + + def wait(self) -> None: + for w in self.works: + w.wait() + + +@dataclass(frozen=True) +class PendingPerSiteCi: + """One coalesced per-site CI-values irecv, held until ``wait()``. + + The packed buffer carries ``sites`` worth of CI values (in order) as + ``b * seq_len * c_s`` ``_WIRE_DTYPE`` elements each. ``wait`` blocks on the + underlying ``dist.Work`` then materializes per-site ``[b, seq_len, c_s]`` + views into the packed buffer (no copy). + """ + + packed: Tensor + work: "dist.Work" + sites: tuple[str, ...] + site_to_c: dict[str, int] + b: int + seq_len: int + + def wait(self) -> dict[str, Tensor]: + self.work.wait() + out: dict[str, Tensor] = {} + offset = 0 + for s in self.sites: + c_s = self.site_to_c[s] + numel = self.b * self.seq_len * c_s + out[s] = self.packed[offset : offset + numel].view(self.b, self.seq_len, c_s) + offset += numel + assert offset == self.packed.numel(), ( + f"unpack size mismatch: consumed {offset} of {self.packed.numel()}" + ) + return out + + +# ────────────────────────────────────────────────────────────────────────────── +# Edge 1: CI → LW. Per-site CI values, sub-sliced to each LW rank's batch shard. +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class CiValuesToLayerwise: + """CI → LW per-site CI values. + + Sender (CI rank): for each LW block + each LW rank whose batch shard sits + in my CI slice, isend that block's owned-sites packet sub-sliced to the + LW rank. Receiver (LW rank): irecv one coalesced packet for my owned sites + from the CI rank that owns my batch shard. + + Pack layout (one packet per (block, block-rank)): for each site in the + block's owned-sites order, ``b_lw * seq_len * C_s`` contiguous + ``_WIRE_DTYPE`` elements. + """ + + world: World + + def send(self, ci_full: dict[str, Tensor], *, my_ci_slice_idx: int) -> SendHandle: + """``ci_full`` keyed by site, shape ``[B_local_ci, S, C_s]`` (CI fn is + global so it has every site). Returns a handle held alive until wait.""" + w = self.world + works: list[dist.Work] = [] + buffers: list[Tensor] = [] + my_lw_block_ranks = w.lw_block_ranks_for_ci_slice(my_ci_slice_idx) + with _time_nccl_op("CiValuesToLayerwise.send"): + for bg in w.layerwise_block_groups: + for block_rank_idx in my_lw_block_ranks: + target = bg.ranks[block_rank_idx] + sub = w.lw_sub_slice_within_ci(block_rank_idx) + parts = [ + ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() + for site in bg.owned_sites + ] + packed = torch.cat(parts) + works.append(dist.isend(packed, dst=target, group=w.cross_pool_p2p_group)) + buffers.append(packed) + return SendHandle(works=works, buffers=buffers) + + def post_recv( + self, + *, + my_within_block_idx: int, + my_owned_sites: tuple[str, ...], + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + ) -> PendingPerSiteCi: + w = self.world + src_ci_slice = w.ci_slice_of_lw_block_rank(my_within_block_idx) + src = w.ci_ranks[src_ci_slice] + b_lw = w.batch_local_lw + packed_numel = sum(b_lw * seq_len * site_to_c[s] for s in my_owned_sites) + packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + with _time_nccl_op("CiValuesToLayerwise.post_recv"): + work = dist.irecv(packed, src=src, group=w.cross_pool_p2p_group) + assert work is not None + return PendingPerSiteCi( + packed=packed, + work=work, + sites=my_owned_sites, + site_to_c=site_to_c, + b=b_lw, + seq_len=seq_len, + ) + + +# ────────────────────────────────────────────────────────────────────────────── +# Edge 2: CI → PPGD. Full-model CI values, sub-sliced to each PPGD rank. +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class CiValuesToPpgd: + """CI → PPGD full-model CI values. + + Sender (CI rank): for each PPGD rank whose batch shard sits in my CI slice, + isend one packet of all sites sub-sliced to that PPGD rank. Receiver (PPGD + rank): irecv one coalesced full-model packet from the CI rank that owns my + batch shard. + + Pack layout: for each site in ``world.all_sites`` order, ``b_pp * seq_len * + C_s`` contiguous ``_WIRE_DTYPE`` elements. + """ + + world: World + + def send(self, ci_full: dict[str, Tensor], *, my_ci_slice_idx: int) -> SendHandle: + w = self.world + works: list[dist.Work] = [] + buffers: list[Tensor] = [] + my_ppgd_slice_idxs = w.ppgd_slice_idxs_for_ci_slice(my_ci_slice_idx) + with _time_nccl_op("CiValuesToPpgd.send"): + for ppgd_slice_idx in my_ppgd_slice_idxs: + target = w.ppgd_ranks[ppgd_slice_idx] + sub = w.ppgd_sub_slice_within_ci(ppgd_slice_idx) + parts = [ + ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() + for site in w.all_sites + ] + packed = torch.cat(parts) + works.append(dist.isend(packed, dst=target, group=w.cross_pool_p2p_group)) + buffers.append(packed) + return SendHandle(works=works, buffers=buffers) + + def post_recv( + self, + *, + my_ppgd_slice_idx: int, + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + ) -> PendingPerSiteCi: + w = self.world + src_ci_slice = w.ci_slice_of_ppgd_slice(my_ppgd_slice_idx) + src = w.ci_ranks[src_ci_slice] + b_pp = w.batch_local_ppgd + packed_numel = sum(b_pp * seq_len * site_to_c[s] for s in w.all_sites) + packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + with _time_nccl_op("CiValuesToPpgd.post_recv"): + work = dist.irecv(packed, src=src, group=w.cross_pool_p2p_group) + assert work is not None + return PendingPerSiteCi( + packed=packed, + work=work, + sites=w.all_sites, + site_to_c=site_to_c, + b=b_pp, + seq_len=seq_len, + ) + + +# ────────────────────────────────────────────────────────────────────────────── +# Edge 3: LW → CI. Per-owned-site CI grads, stitched into the CI rank's slice. +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class GradCiFromLayerwise: + """LW → CI per-owned-site CI grads. + + Sender (LW rank): coalesce my owned sites' grads (full LW batch slice) into + one packed send to the CI rank that owns my slice. Receiver (CI rank): recv + one packet per LW source whose batch shard sits in my CI slice and stitch + each into a per-site fp32 ``[B_local_ci, S, C_s]`` dest. + + Pack layout (one packet per LW source): for each site in the source block's + owned-sites order, ``b_lw * seq_len * c_s`` contiguous ``_WIRE_DTYPE`` elements. + """ + + world: World + + def send( + self, + g_ci_owned: dict[str, Tensor], + *, + my_within_block_idx: int, + my_owned_sites: tuple[str, ...], + ) -> None: + w = self.world + dst_ci_slice = w.ci_slice_of_lw_block_rank(my_within_block_idx) + dst = w.ci_ranks[dst_ci_slice] + parts = [ + g_ci_owned[s].detach().to(_WIRE_DTYPE).contiguous().flatten() for s in my_owned_sites + ] + packed = torch.cat(parts) + with _time_nccl_op("GradCiFromLayerwise.send"): + dist.send(packed, dst=dst, group=w.cross_pool_p2p_group) + + def recv( + self, + *, + my_ci_slice_idx: int, + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + ) -> dict[str, Tensor]: + w = self.world + my_lw_block_ranks = w.lw_block_ranks_for_ci_slice(my_ci_slice_idx) + b_lw = w.batch_local_lw + + pending: list[tuple[int, Tensor, dist.Work, tuple[str, ...]]] = [] + with _time_nccl_op("GradCiFromLayerwise.recv:post_irecvs"): + for bg in w.layerwise_block_groups: + owned = bg.owned_sites + packed_numel = sum(b_lw * seq_len * site_to_c[s] for s in owned) + for block_rank_idx in my_lw_block_ranks: + src = bg.ranks[block_rank_idx] + buf = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + work = dist.irecv(buf, src=src, group=w.cross_pool_p2p_group) + assert work is not None + pending.append((block_rank_idx, buf, work, owned)) + + b_ci = w.batch_local_ci + out: dict[str, Tensor] = { + s: torch.empty(b_ci, seq_len, site_to_c[s], device=device, dtype=torch.float32) + for s in w.all_sites + } + with _time_nccl_op("GradCiFromLayerwise.recv:wait"): + for block_rank_idx, buf, work, owned in pending: + work.wait() + sub = w.lw_sub_slice_within_ci(block_rank_idx) + offset = 0 + for site in owned: + c_s = site_to_c[site] + n = b_lw * seq_len * c_s + site_view = buf[offset : offset + n].view(b_lw, seq_len, c_s) + out[site][sub].copy_(site_view.to(torch.float32)) + offset += n + return out + + +# ────────────────────────────────────────────────────────────────────────────── +# Edge 4: PPGD → CI. Full-model CI grads, stitched into the CI rank's slice. +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class GradCiFromPpgd: + """PPGD → CI full-model CI grads. + + Sender (PPGD rank): coalesce all sites' grads (PPGD batch slice) into one + packed send to the CI rank that owns my slice. Receiver (CI rank): recv one + packet per PPGD source in my CI slice and stitch each into a per-site fp32 + ``[B_local_ci, S, C_s]`` dest. + + Pack layout: for each site in ``world.all_sites`` order, ``b_pp * seq_len * + c_s`` contiguous ``_WIRE_DTYPE`` elements. + """ + + world: World + + def send(self, g_ci_full: dict[str, Tensor], *, my_ppgd_slice_idx: int) -> None: + w = self.world + dst_ci_slice = w.ci_slice_of_ppgd_slice(my_ppgd_slice_idx) + dst = w.ci_ranks[dst_ci_slice] + parts = [g_ci_full[s].detach().to(_WIRE_DTYPE).contiguous().flatten() for s in w.all_sites] + packed = torch.cat(parts) + with _time_nccl_op("GradCiFromPpgd.send"): + dist.send(packed, dst=dst, group=w.cross_pool_p2p_group) + + def recv( + self, + *, + my_ci_slice_idx: int, + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + ) -> dict[str, Tensor]: + w = self.world + my_ppgd_slice_idxs = w.ppgd_slice_idxs_for_ci_slice(my_ci_slice_idx) + b_pp = w.batch_local_ppgd + + site_numels = {s: b_pp * seq_len * site_to_c[s] for s in w.all_sites} + packed_numel = sum(site_numels.values()) + + pending: list[tuple[int, Tensor, dist.Work]] = [] + with _time_nccl_op("GradCiFromPpgd.recv:post_irecvs"): + for ppgd_slice_idx in my_ppgd_slice_idxs: + src = w.ppgd_ranks[ppgd_slice_idx] + packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + work = dist.irecv(packed, src=src, group=w.cross_pool_p2p_group) + assert work is not None + pending.append((ppgd_slice_idx, packed, work)) + + b_ci = w.batch_local_ci + out: dict[str, Tensor] = { + s: torch.empty(b_ci, seq_len, site_to_c[s], device=device, dtype=torch.float32) + for s in w.all_sites + } + with _time_nccl_op("GradCiFromPpgd.recv:wait"): + for ppgd_slice_idx, packed, work in pending: + work.wait() + sub = w.ppgd_sub_slice_within_ci(ppgd_slice_idx) + offset = 0 + for site in w.all_sites: + c_s = site_to_c[site] + n = site_numels[site] + buf = packed[offset : offset + n].view(b_pp, seq_len, c_s) + out[site][sub].copy_(buf.to(torch.float32)) + offset += n + return out + + +# ────────────────────────────────────────────────────────────────────────────── +# Edge 5: PPGD → LW. Per-owned-site V/U grads (post in-pool sum-reduce). +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class GradVuFromPpgd: + """PPGD → LW per-owned-site V/U grads. + + Sender (PPGD leader): one coalesced isend per LW block to its leader (V/U + grads already sum-reduced within the PPGD pool, so the leader's copy is the + full-batch grad). Receiver (LW): block leader recvs its owned sites, then + in-block broadcasts so every replica sees the same grad. + + Pack layout (per LW block): for each site in the block's owned-sites order, + ``V.numel()`` then ``U.numel()`` contiguous ``_WIRE_DTYPE`` elements. + """ + + world: World + + def send( + self, v_grads: dict[str, Tensor], u_grads: dict[str, Tensor], *, is_pool_leader: bool + ) -> None: + if not is_pool_leader: + return + w = self.world + works: list[dist.Work] = [] + buffers: list[Tensor] = [] + with _time_nccl_op("GradVuFromPpgd.send:isends"): + for bg in w.layerwise_block_groups: + parts: list[Tensor] = [] + for site in bg.owned_sites: + parts.append(v_grads[site].to(_WIRE_DTYPE).contiguous().flatten()) + parts.append(u_grads[site].to(_WIRE_DTYPE).contiguous().flatten()) + packed = torch.cat(parts) + work = dist.isend(packed, dst=bg.leader, group=w.cross_pool_p2p_group) + assert work is not None + works.append(work) + buffers.append(packed) + with _time_nccl_op("GradVuFromPpgd.send:wait"): + for work in works: + work.wait() + del buffers + + def recv( + self, + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + *, + my_block_idx: int, + my_owned_sites: tuple[str, ...], + my_is_block_leader: bool, + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + w = self.world + v_grads: dict[str, Tensor] = {} + u_grads: dict[str, Tensor] = {} + + if my_is_block_leader: + packed_numel = sum( + v_templates[s].numel() + u_templates[s].numel() for s in my_owned_sites + ) + sample = v_templates[my_owned_sites[0]] + packed = torch.empty(packed_numel, dtype=_WIRE_DTYPE, device=sample.device) + ppgd_leader = w.ppgd_ranks[0] + with _time_nccl_op("GradVuFromPpgd.recv:recv"): + dist.recv(packed, src=ppgd_leader, group=w.cross_pool_p2p_group) + offset = 0 + for s in my_owned_sites: + v_n = v_templates[s].numel() + u_n = u_templates[s].numel() + v_grads[s] = ( + packed[offset : offset + v_n].view_as(v_templates[s]).to(v_templates[s].dtype) + ) + offset += v_n + u_grads[s] = ( + packed[offset : offset + u_n].view_as(u_templates[s]).to(u_templates[s].dtype) + ) + offset += u_n + else: + for s in my_owned_sites: + v_grads[s] = torch.empty_like(v_templates[s]) + u_grads[s] = torch.empty_like(u_templates[s]) + + block_group = w.block_group_groups[my_block_idx] + block_leader_rank = w.layerwise_block_groups[my_block_idx].leader + with _time_nccl_op("GradVuFromPpgd.recv:in_block_bcast"): + for s in my_owned_sites: + v_grads[s] = v_grads[s].contiguous() + u_grads[s] = u_grads[s].contiguous() + dist.broadcast(v_grads[s], src=block_leader_rank, group=block_group) + dist.broadcast(u_grads[s], src=block_leader_rank, group=block_group) + + return v_grads, u_grads + + +# ────────────────────────────────────────────────────────────────────────────── +# Edge 6: LW → PPGD. Updated V/U, leader-rooted broadcast to the PPGD pool. +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class UpdatedVuToPpgd: + """LW → PPGD updated V/U. + + Sender (LW block leader): one coalesced leader-rooted broadcast of updated + V/U over the {block-leader} ∪ {ppgd_ranks} group. Receiver (PPGD): one async + broadcast recv per LW block (they pipeline across the per-group NCCL + streams), waited + unpacked into per-site V/U. + + Pack layout (per LW block): for each site in the block's owned-sites order, + ``V.numel()`` then ``U.numel()`` contiguous ``_WIRE_DTYPE`` elements. + """ + + world: World + + def send( + self, + v_owned: dict[str, Tensor], + u_owned: dict[str, Tensor], + *, + my_rank: int, + my_block_idx: int, + my_owned_sites: tuple[str, ...], + my_is_block_leader: bool, + ) -> SendHandle: + if not my_is_block_leader: + return SendHandle(works=[], buffers=[]) + w = self.world + parts: list[Tensor] = [] + for s in my_owned_sites: + parts.append(v_owned[s].detach().to(_WIRE_DTYPE).contiguous().flatten()) + parts.append(u_owned[s].detach().to(_WIRE_DTYPE).contiguous().flatten()) + packed = torch.cat(parts) + bcast_group = w.cross_pool_bcast_groups[my_block_idx] + with _time_nccl_op("UpdatedVuToPpgd.send"): + work = dist.broadcast(packed, src=my_rank, group=bcast_group, async_op=True) + assert work is not None + return SendHandle(works=[work], buffers=[packed]) + + def recv( + self, v_templates: dict[str, Tensor], u_templates: dict[str, Tensor] + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + w = self.world + bufs: list[tuple[tuple[str, ...], Tensor, dist.Work]] = [] + with _time_nccl_op("UpdatedVuToPpgd.recv"): + for bg_idx, bg in enumerate(w.layerwise_block_groups): + owned = bg.owned_sites + packed_numel = sum(v_templates[s].numel() + u_templates[s].numel() for s in owned) + sample = v_templates[owned[0]] + packed = torch.empty(packed_numel, dtype=_WIRE_DTYPE, device=sample.device) + bcast_group = w.cross_pool_bcast_groups[bg_idx] + work = dist.broadcast(packed, src=bg.leader, group=bcast_group, async_op=True) + assert work is not None + bufs.append((owned, packed, work)) + + v_new: dict[str, Tensor] = {} + u_new: dict[str, Tensor] = {} + for owned, packed, work in bufs: + work.wait() + offset = 0 + for s in owned: + v_n = v_templates[s].numel() + u_n = u_templates[s].numel() + v_new[s] = ( + packed[offset : offset + v_n].view_as(v_templates[s]).to(v_templates[s].dtype) + ) + offset += v_n + u_new[s] = ( + packed[offset : offset + u_n].view_as(u_templates[s]).to(u_templates[s].dtype) + ) + offset += u_n + return v_new, u_new + + +# ────────────────────────────────────────────────────────────────────────────── +# Eval-only edge: CI → PPGD. Full CIOutputs (lower/upper/pre_sigmoid). +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class CiOutputsToPpgd: + """CI → PPGD full ``CIOutputs`` (eval only). + + Training ships only ``lower_leaky``; eval ships all three dicts so any metric + reading ``ctx.ci`` works without a per-metric audit. Synchronous because eval + is rare and overlap has no value. + + Pack layout per send: three contiguous blocks (lower_leaky, upper_leaky, + pre_sigmoid). Each block has, for each site in ``world.all_sites`` order, + ``b_pp * seq_len * C_s`` contiguous ``_WIRE_DTYPE`` elements. + """ + + world: World + + def send(self, ci: CIOutputs, *, my_ci_slice_idx: int) -> None: + w = self.world + my_ppgd_slice_idxs = w.ppgd_slice_idxs_for_ci_slice(my_ci_slice_idx) + with _time_nccl_op("CiOutputsToPpgd.send"): + for ppgd_slice_idx in my_ppgd_slice_idxs: + target = w.ppgd_ranks[ppgd_slice_idx] + sub = w.ppgd_sub_slice_within_ci(ppgd_slice_idx) + parts: list[Tensor] = [] + for d in (ci.lower_leaky, ci.upper_leaky, ci.pre_sigmoid): + parts.extend( + d[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() + for site in w.all_sites + ) + packed = torch.cat(parts) + dist.send(packed, dst=target, group=w.cross_pool_p2p_group) + + def recv( + self, + *, + my_ppgd_slice_idx: int, + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + ) -> CIOutputs: + w = self.world + src_ci_slice = w.ci_slice_of_ppgd_slice(my_ppgd_slice_idx) + src = w.ci_ranks[src_ci_slice] + b_pp = w.batch_local_ppgd + + per_block_numel = sum(b_pp * seq_len * site_to_c[s] for s in w.all_sites) + packed = torch.empty(3 * per_block_numel, device=device, dtype=_WIRE_DTYPE) + with _time_nccl_op("CiOutputsToPpgd.recv"): + dist.recv(packed, src=src, group=w.cross_pool_p2p_group) + + out: list[dict[str, Tensor]] = [{}, {}, {}] + offset = 0 + for block_idx in range(3): + for site in w.all_sites: + c_s = site_to_c[site] + numel = b_pp * seq_len * c_s + out[block_idx][site] = packed[offset : offset + numel].view(b_pp, seq_len, c_s) + offset += numel + assert offset == packed.numel(), f"unpack mismatch: {offset} of {packed.numel()}" + return CIOutputs(lower_leaky=out[0], upper_leaky=out[1], pre_sigmoid=out[2]) + + +@dataclass(frozen=True) +class Portals: + """The full set of cross-pool exchange portals, built once per rank from + the shared ``World``. Every rank holds the same set; each pool invokes only + the sides its role plays. Threaded into the step functions so the per-step + flow reads as portal invocations against the dependency DAG. + """ + + ci_values_to_lw: CiValuesToLayerwise + ci_values_to_ppgd: CiValuesToPpgd + grad_ci_from_lw: GradCiFromLayerwise + grad_ci_from_ppgd: GradCiFromPpgd + grad_vu_from_ppgd: GradVuFromPpgd + updated_vu_to_ppgd: UpdatedVuToPpgd + ci_outputs_to_ppgd: CiOutputsToPpgd + + @classmethod + def from_world(cls, world: World) -> "Portals": + return cls( + ci_values_to_lw=CiValuesToLayerwise(world), + ci_values_to_ppgd=CiValuesToPpgd(world), + grad_ci_from_lw=GradCiFromLayerwise(world), + grad_ci_from_ppgd=GradCiFromPpgd(world), + grad_vu_from_ppgd=GradVuFromPpgd(world), + updated_vu_to_ppgd=UpdatedVuToPpgd(world), + ci_outputs_to_ppgd=CiOutputsToPpgd(world), + ) diff --git a/param_decomp_lab/three_pool/step_ci.py b/param_decomp_lab/three_pool/step_ci.py index 966e3de21..50341dff4 100644 --- a/param_decomp_lab/three_pool/step_ci.py +++ b/param_decomp_lab/three_pool/step_ci.py @@ -3,6 +3,15 @@ CI pool is new under 3-pool. Each CI rank holds the full CI fn (replicated) and processes one DP shard of the batch. +The step is a sequence of typed phases. Each phase consumes the previous +phase's typed bundle, so the dependency order is a type constraint: you cannot +send CI values before computing them (``send_ci`` takes ``CiForward``), cannot +assemble the total CI grad before receiving both halves (``GciTotal`` takes +``GciReceived``), cannot run the fused backward before both the imp-min loss +and the assembled grads exist (it takes ``CiForward`` + ``ImpMinLoss`` + +``GciTotal``), and cannot wait the sends before posting them (``wait`` takes +``CiSends``). + Phases (numbered to match ``DESIGN.md`` ``ci/N``): 0. Step 0 only: target_fwd to build H_T (subsequent steps reuse the prev @@ -32,6 +41,7 @@ # pyright: reportArgumentType=false import os +from dataclasses import dataclass from typing import Any import torch @@ -50,11 +60,59 @@ ) from param_decomp.torch_helpers import bf16_autocast from param_decomp_lab.three_pool.layout import ThreePoolLayout +from param_decomp_lab.three_pool.portals import Portals, SendHandle from param_decomp_lab.three_pool.runtime import _ThreePoolRuntime +@dataclass(frozen=True) +class CiForward: + """Phase ci/1 output: the CI fn forward graph + its derived seq_len. + + Holds the live autograd graph (``ci``) that phases 2 (send), 3 (imp_min), + and 8 (fused backward) all attach to. Constructing this is the only way to + obtain a ``CIOutputs`` the rest of the step can act on. + """ + + ci: CIOutputs + seq_len: int + + +@dataclass(frozen=True) +class CiSends: + """Phase ci/2 output: the two in-flight CI-value sends (LW + PPGD).""" + + to_lw: SendHandle + to_ppgd: SendHandle + + +@dataclass(frozen=True) +class ImpMinLoss: + """Phase ci/3 output: the (CI-pool-global) importance-minimality scalar, + still attached to the CI fn graph via ``ci.upper_leaky``.""" + + loss: Tensor + + +@dataclass(frozen=True) +class GciReceived: + """Phases ci/5 + ci/6 output: per-site CI grads from LW and from PPGD, + each stitched onto this CI rank's batch slice.""" + + from_lw: dict[str, Tensor] + from_ppgd: dict[str, Tensor] + + +@dataclass(frozen=True) +class GciTotal: + """Phase ci/7 output: ``g_CI_LW + g_CI_PPGD`` per site, ready to seed the + fused backward on ``ci.lower_leaky``.""" + + per_site: dict[str, Tensor] + + def step_ci( layout: ThreePoolLayout, + portals: Portals, component_model: ComponentModel, optimizer: torch.optim.Optimizer, ci_fn_params: list[nn.Parameter], @@ -72,6 +130,7 @@ def step_ci( ``batch_T_plus_1`` is ``None`` and the prefetch is skipped. """ assert layout.my_pool == "ci" + assert layout.my_ci_slice_idx is not None device = next(component_model.parameters()).device batch_T_local, batch_T_plus_1_local = _slice_batches_for_ci(batch_T, batch_T_plus_1, layout) @@ -79,33 +138,12 @@ def step_ci( if h_cache_T is None: h_cache_T = _target_fwd_and_cache(component_model, batch_T_local, cfg.bf16_autocast) - with bf16_autocast(cfg.bf16_autocast): - ci = component_model.calc_causal_importances( - pre_weight_acts=h_cache_T, sampling="continuous", detach_inputs=False - ) - seq_len = _seq_len_from_ci(ci.lower_leaky) - _assert_ci_shapes(ci.lower_leaky, layout, seq_len, cfg) - - send_works_lw, send_bufs_lw = layout.async_send_ci_to_layerwise(ci.lower_leaky) - send_works_pgd, send_bufs_pgd = layout.async_send_ci_to_ppgd(ci.lower_leaky) - - loss_imp = _importance_minimality_loss( - ci.upper_leaky, - current_frac_of_training, - cfg, - ci_pool_group=layout.world.ci_pool_group, - n_ci_pool=layout.world.n_ci, - ) - - h_cache_T_plus_1: dict[str, Tensor] | None = None - if batch_T_plus_1_local is not None: - h_cache_T_plus_1 = _target_fwd_and_cache( - component_model, batch_T_plus_1_local, cfg.bf16_autocast - ) - - g_ci_lw = layout.recv_g_ci_from_layerwise(cfg.c_per_site, seq_len, device) - g_ci_pgd = layout.recv_g_ci_from_ppgd(cfg.c_per_site, seq_len, device) - g_ci_total = _assemble_g_ci_total(g_ci_lw, g_ci_pgd, layout, cfg, seq_len) + fwd = _ci_fn_forward(component_model, h_cache_T, layout, cfg) + sends = _send_ci_values(portals, fwd, layout.my_ci_slice_idx) + imp = _imp_min_phase(fwd, current_frac_of_training, cfg, layout) + h_cache_T_plus_1 = _prefetch_next_h(component_model, batch_T_plus_1_local, cfg) + received = _recv_g_ci(portals, layout, cfg, fwd.seq_len, device) + total = _assemble_g_ci_total(received, layout, cfg, fwd.seq_len) optimizer.zero_grad(set_to_none=True) # Diagnostic: sync before the bwd so phase("ci/8a") measures only the bwd @@ -114,7 +152,7 @@ def step_ci( # wall was dominated by pending stream work, not by the bwd. Remove after diagnosis. if os.environ.get("PD_SYNC_BEFORE_8A", "").strip() in ("1", "true", "yes"): torch.cuda.synchronize() - _fused_backward_through_ci_fn(loss_imp, ci, g_ci_total, layout, cfg) + _fused_backward_through_ci_fn(imp, fwd, total, layout, cfg) _maybe_emit_ci_fn_bwd_breakdown(component_model) layout.all_reduce_ci_fn_grads(ci_fn_params) @@ -127,24 +165,98 @@ def step_ci( ) optimizer.step() - for w in [*send_works_lw, *send_works_pgd]: - w.wait() - del send_bufs_lw, send_bufs_pgd + sends.to_lw.wait() + sends.to_ppgd.wait() - # imp is already globally aggregated inside ``_importance_minimality_loss`` - # (per_component_sums + n_examples SUM-reduced across CI pool), so every CI - # rank holds the same scalar. Divide by ``n_ci`` so the logger's cross-pool - # SUM all-reduce gives back the global value exactly once. + # imp is already globally aggregated inside ``_imp_min_phase`` (per_component_sums + # + n_examples SUM-reduced across CI pool), so every CI rank holds the same + # scalar. Divide by ``n_ci`` so the logger's cross-pool SUM all-reduce gives back + # the global value exactly once. # # ``.item()`` is a CPU↔GPU sync — only pay it on steps we actually log to. if should_log: - imp_value = loss_imp.item() + imp_value = imp.loss.item() metrics = {"loss/imp": imp_value, "_raw/imp_num": imp_value / layout.world.n_ci} else: metrics = {} return metrics, h_cache_T_plus_1 +def _ci_fn_forward( + component_model: ComponentModel, + h_cache_T: dict[str, Tensor], + layout: ThreePoolLayout, + cfg: _ThreePoolRuntime, +) -> CiForward: + """Phase ci/1. CI fn forward on H_T → ``CIOutputs`` (graph retained for ci/8).""" + with bf16_autocast(cfg.bf16_autocast): + ci = component_model.calc_causal_importances( + pre_weight_acts=h_cache_T, sampling="continuous", detach_inputs=False + ) + seq_len = _seq_len_from_ci(ci.lower_leaky) + _assert_ci_shapes(ci.lower_leaky, layout, seq_len, cfg) + return CiForward(ci=ci, seq_len=seq_len) + + +def _send_ci_values(portals: Portals, fwd: CiForward, my_ci_slice_idx: int) -> CiSends: + """Phase ci/2. Async-ship CI_T to LW (per-site) and PPGD (full-model).""" + to_lw = portals.ci_values_to_lw.send(fwd.ci.lower_leaky, my_ci_slice_idx=my_ci_slice_idx) + to_ppgd = portals.ci_values_to_ppgd.send(fwd.ci.lower_leaky, my_ci_slice_idx=my_ci_slice_idx) + return CiSends(to_lw=to_lw, to_ppgd=to_ppgd) + + +def _imp_min_phase( + fwd: CiForward, + current_frac_of_training: float, + cfg: _ThreePoolRuntime, + layout: ThreePoolLayout, +) -> ImpMinLoss: + """Phase ci/3. Importance-minimality loss on ``ci.upper_leaky``.""" + loss = _importance_minimality_loss( + fwd.ci.upper_leaky, + current_frac_of_training, + cfg, + ci_pool_group=layout.world.ci_pool_group, + n_ci_pool=layout.world.n_ci, + ) + return ImpMinLoss(loss=loss) + + +def _prefetch_next_h( + component_model: ComponentModel, + batch_T_plus_1_local: Any | None, + cfg: _ThreePoolRuntime, +) -> dict[str, Tensor] | None: + """Phase ci/4. Dead-time prefetch of H_{T+1} (skipped on the last step).""" + if batch_T_plus_1_local is None: + return None + return _target_fwd_and_cache(component_model, batch_T_plus_1_local, cfg.bf16_autocast) + + +def _recv_g_ci( + portals: Portals, + layout: ThreePoolLayout, + cfg: _ThreePoolRuntime, + seq_len: int, + device: torch.device, +) -> GciReceived: + """Phases ci/5 + ci/6. Recv CI grads from LW and PPGD.""" + assert layout.my_ci_slice_idx is not None + from_lw = portals.grad_ci_from_lw.recv( + my_ci_slice_idx=layout.my_ci_slice_idx, + site_to_c=cfg.c_per_site, + seq_len=seq_len, + device=device, + ) + from_ppgd = portals.grad_ci_from_ppgd.recv( + my_ci_slice_idx=layout.my_ci_slice_idx, + site_to_c=cfg.c_per_site, + seq_len=seq_len, + device=device, + ) + return GciReceived(from_lw=from_lw, from_ppgd=from_ppgd) + + def _slice_batches_for_ci( batch_T: Any, batch_T_plus_1: Any | None, layout: ThreePoolLayout ) -> tuple[Any, Any | None]: @@ -178,7 +290,7 @@ def _assert_ci_shapes( ) -> None: """Sanity-check CI fn outputs match [B_local_ci, seq_len, C_s] per site. - Catches misconfigured ``c_per_site`` or a wrong per-rank batch slice fast. + Catches a misconfigured ``c_per_site`` or a wrong per-rank batch slice fast. """ batch_local_ci = layout.world.batch_local_ci for s, c in cfg.c_per_site.items(): @@ -190,22 +302,21 @@ def _assert_ci_shapes( def _assemble_g_ci_total( - g_ci_lw: dict[str, Tensor], - g_ci_pgd: dict[str, Tensor], + received: GciReceived, layout: ThreePoolLayout, cfg: _ThreePoolRuntime, seq_len: int, -) -> dict[str, Tensor]: +) -> GciTotal: """Phase ci/7. ``g_CI_total[s] = g_CI_LW[s] + g_CI_PPGD[s]``. Both summands live on this CI rank's batch slice [B_local_ci, S, C_s]. Loss coefficients were already baked in by LW/PPGD before they bwd'd. """ batch_local_ci = layout.world.batch_local_ci - g_ci_total: dict[str, Tensor] = {} + per_site: dict[str, Tensor] = {} for s in layout.world.all_sites: c = cfg.c_per_site[s] - lw, pgd = g_ci_lw[s], g_ci_pgd[s] + lw, pgd = received.from_lw[s], received.from_ppgd[s] assert lw.shape == (batch_local_ci, seq_len, c), ( f"g_ci_lw[{s!r}] shape {tuple(lw.shape)} != expected ({batch_local_ci}, {seq_len}, {c})" ) @@ -213,8 +324,8 @@ def _assemble_g_ci_total( f"g_ci_pgd[{s!r}] shape {tuple(pgd.shape)} != " f"expected ({batch_local_ci}, {seq_len}, {c})" ) - g_ci_total[s] = lw + pgd - return g_ci_total + per_site[s] = lw + pgd + return GciTotal(per_site=per_site) def _maybe_emit_ci_fn_bwd_breakdown(component_model: ComponentModel) -> None: @@ -242,9 +353,9 @@ def _maybe_emit_ci_fn_bwd_breakdown(component_model: ComponentModel) -> None: def _fused_backward_through_ci_fn( - loss_imp: Tensor, - ci: CIOutputs, - g_ci_total: dict[str, Tensor], + imp: ImpMinLoss, + fwd: CiForward, + total: GciTotal, layout: ThreePoolLayout, cfg: _ThreePoolRuntime, ) -> None: @@ -264,10 +375,11 @@ def _fused_backward_through_ci_fn( between the two backward paths — to find out which one dominates and where to optimize next. """ + loss_imp = imp.loss assert loss_imp.dim() == 0, f"loss_imp must be scalar; got {loss_imp.shape}" scaled_imp = cfg.coeff_imp * loss_imp - lower_leaky_tensors = [ci.lower_leaky[s] for s in layout.world.all_sites] - g_ci_total_seeds = [g_ci_total[s] for s in layout.world.all_sites] + lower_leaky_tensors = [fwd.ci.lower_leaky[s] for s in layout.world.all_sites] + g_ci_total_seeds = [total.per_site[s] for s in layout.world.all_sites] torch.autograd.backward( tensors=lower_leaky_tensors, grad_tensors=g_ci_total_seeds, diff --git a/param_decomp_lab/three_pool/step_layerwise.py b/param_decomp_lab/three_pool/step_layerwise.py index 0fea5523a..8d6488030 100644 --- a/param_decomp_lab/three_pool/step_layerwise.py +++ b/param_decomp_lab/three_pool/step_layerwise.py @@ -2,6 +2,13 @@ Trains V/U on the LW pool with the CI fn living on the CI pool. +The step is a sequence of typed phases threaded through ``strategy.context``. +The handoff types make the dependency order a type constraint: the CI grads can +only be sent once the per-site streaming backward has populated the re-leafed CI +tensors' ``.grad`` (``send_g_ci`` consumes the ``CiLeaves`` those grads live +on), and the CI values can only be consumed after the posted recv is waited +(``CiLeaves`` is built from a ``PendingPerSiteCi.wait()``). + Phases (numbered to match ``DESIGN.md`` ``lw/N``): A1. Post async irecv for CI_T from the owning CI rank (overlaps with A2). @@ -24,6 +31,7 @@ # pyright: reportArgumentType=false, reportOperatorIssue=false, reportAttributeAccessIssue=false +from dataclasses import dataclass from typing import Any import torch @@ -37,11 +45,41 @@ from param_decomp.torch_helpers import bf16_autocast from param_decomp_lab.three_pool.layout import ThreePoolLayout from param_decomp_lab.three_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp_lab.three_pool.portals import PendingPerSiteCi, Portals from param_decomp_lab.three_pool.runtime import _ThreePoolRuntime +@dataclass(frozen=True) +class Faith: + """Phase lw/D1 output: faithfulness loss (already backward'd into V/U .grad) + plus the raw sum-sq + numel the logger needs for the global ratio.""" + + loss: Tensor + sum_sq: Tensor + numel: int + + +@dataclass(frozen=True) +class CiLeaves: + """Phase lw/D2 output: re-leafed fp32 CI values (requires_grad=True) per + owned site. The layerwise backward populates ``leaf.grad``; phase D4 reads + that grad off these exact leaves to ship back to the CI pool.""" + + per_site: dict[str, Tensor] + + +@dataclass(frozen=True) +class Stoch: + """Phase lw/D3 output: accumulated stochastic-recon display value (GPU + tensor) + the count of owned sites it averages over.""" + + total: Tensor + n_owned: int + + def step_layerwise( layout: ThreePoolLayout, + portals: Portals, component_model: ComponentModel, optimizer: torch.optim.Optimizer, all_params: list[nn.Parameter], @@ -54,75 +92,144 @@ def step_layerwise( """One LW step: A → D → tail (all_reduce, clip, opt, async send V/U).""" assert layout.my_pool == "layerwise" device = next(component_model.parameters()).device - n_sites_total = len(cfg.c_per_site) batch_local, seq_len = _slice_batch_for_layerwise(batch, layout) with strategy.context(component_model.target_model): - ci_recv_pending = layout.async_recv_ci_from_ci_pool( - {s: cfg.c_per_site[s] for s in layout.my_owned_sites}, - seq_len=seq_len, - device=device, - ) - with torch.no_grad(), bf16_autocast(cfg.bf16_autocast): - target_local = component_model(batch_local).detach() + ci_recv_pending = _post_ci_recv(portals, layout, cfg, seq_len, device) + target_local = _target_fwd(component_model, batch_local, cfg) for param in all_params: param.grad = None with strategy.context(component_model.target_model): - loss_faith, faith_sum_sq_t, faith_numel = _faithfulness_loss( - component_model, device, cfg.numel_global - ) - (cfg.coeff_faith * loss_faith).backward() - - ci_recv = ci_recv_pending.wait_and_unpack() - ci_recv_leaves = _releaf_ci_fp32_for_grads(ci_recv, layout.my_owned_sites) - _assert_ci_recv_shapes(ci_recv_leaves, layout, seq_len, cfg) - - with bf16_autocast(cfg.bf16_autocast): - # Accumulate the display value as a GPU tensor (not a Python float) so - # the per-site ``.item()`` doesn't force a CPU↔GPU sync that serializes - # each site's bwd against the next. ``loss_s.detach()`` so accumulator - # doesn't retain autograd graph. - stoch_total_t = torch.zeros((), device=device) - for i, s in enumerate(layout.my_owned_sites): - if phase_trace_enabled(): - trace(f"lw/D3 site {i + 1}/{len(layout.my_owned_sites)}: {s} fwd+bwd") - loss_s, n_positions = _layerwise_one_site( - component_model, batch_local, target_local, ci_recv_leaves, s, strategy - ) - assert loss_s.dim() == 0, f"layerwise loss for site {s!r} must be scalar" - (cfg.coeff_stoch * loss_s / (n_positions * n_sites_total)).backward() - stoch_total_t = stoch_total_t + (loss_s.detach() / n_positions) - stoch_n_owned = len(layout.my_owned_sites) - - g_ci_owned = {s: ci_recv_leaves[s].grad for s in layout.my_owned_sites} - assert all(g is not None for g in g_ci_owned.values()), ( - "layerwise backward should have populated ci_recv_leaves[s].grad" + faith = _faithfulness_phase(component_model, device, cfg) + + ci_leaves = _wait_ci_and_releaf(ci_recv_pending, layout, seq_len, cfg) + stoch = _layerwise_streaming_phase( + component_model, batch_local, target_local, ci_leaves, layout, cfg, strategy ) - layout.send_g_ci_to_ci_pool(g_ci_owned) - v_grads_pgd, u_grads_pgd = _recv_g_vu_from_ppgd(layout, component_model) - _combine_vu_grads_in_place(component_model, layout, v_grads_pgd, u_grads_pgd) + _send_g_ci(portals, layout, ci_leaves) + _recv_and_combine_g_vu(portals, component_model, layout) if should_log: - stoch_total_value = stoch_total_t.item() + stoch_total_value = stoch.total.item() metrics = { - "loss/faith": loss_faith.item(), - "loss/stoch": stoch_total_value / stoch_n_owned, - "_raw/faith_num": faith_sum_sq_t.item(), - "_raw/faith_den": float(faith_numel), + "loss/faith": faith.loss.item(), + "loss/stoch": stoch_total_value / stoch.n_owned, + "_raw/faith_num": faith.sum_sq.item(), + "_raw/faith_den": float(faith.numel), "_raw/stoch_num": stoch_total_value, - "_raw/stoch_den": float(stoch_n_owned), + "_raw/stoch_den": float(stoch.n_owned), } else: metrics = {} - _sync_tail(layout, component_model, optimizer, all_params, cfg) + _sync_tail(portals, layout, component_model, optimizer, all_params, cfg) return metrics +def _post_ci_recv( + portals: Portals, + layout: ThreePoolLayout, + cfg: _ThreePoolRuntime, + seq_len: int, + device: torch.device, +) -> PendingPerSiteCi: + """Phase lw/A1. Post the async CI-values irecv (waited at D2).""" + assert layout.my_within_block_idx is not None + return portals.ci_values_to_lw.post_recv( + my_within_block_idx=layout.my_within_block_idx, + my_owned_sites=layout.my_owned_sites, + site_to_c={s: cfg.c_per_site[s] for s in layout.my_owned_sites}, + seq_len=seq_len, + device=device, + ) + + +def _target_fwd( + component_model: ComponentModel, batch_local: Any, cfg: _ThreePoolRuntime +) -> Tensor: + """Phase lw/A2. Detached target forward on this rank's batch slice.""" + with torch.no_grad(), bf16_autocast(cfg.bf16_autocast): + return component_model(batch_local).detach() + + +def _faithfulness_phase( + component_model: ComponentModel, device: torch.device, cfg: _ThreePoolRuntime +) -> Faith: + """Phase lw/D1. Faithfulness loss + backward into V/U .grad.""" + loss, sum_sq, numel = _faithfulness_loss(component_model, device, cfg.numel_global) + (cfg.coeff_faith * loss).backward() + return Faith(loss=loss, sum_sq=sum_sq, numel=numel) + + +def _wait_ci_and_releaf( + pending: PendingPerSiteCi, + layout: ThreePoolLayout, + seq_len: int, + cfg: _ThreePoolRuntime, +) -> CiLeaves: + """Phase lw/D2. Wait the CI recv, re-leaf fp32 with grad for the bwd.""" + ci_recv = pending.wait() + per_site = _releaf_ci_fp32_for_grads(ci_recv, layout.my_owned_sites) + _assert_ci_recv_shapes(per_site, layout, seq_len, cfg) + return CiLeaves(per_site=per_site) + + +def _layerwise_streaming_phase( + component_model: ComponentModel, + batch_local: Any, + target_local: Tensor, + ci_leaves: CiLeaves, + layout: ThreePoolLayout, + cfg: _ThreePoolRuntime, + strategy: LayerwiseLossStrategy, +) -> Stoch: + """Phase lw/D3. Per-owned-site stochastic recon, streaming fwd+bwd.""" + n_sites_total = len(cfg.c_per_site) + device = target_local.device + with bf16_autocast(cfg.bf16_autocast): + # Accumulate the display value as a GPU tensor (not a Python float) so + # the per-site ``.item()`` doesn't force a CPU↔GPU sync that serializes + # each site's bwd against the next. ``loss_s.detach()`` so accumulator + # doesn't retain autograd graph. + stoch_total_t = torch.zeros((), device=device) + for i, s in enumerate(layout.my_owned_sites): + if phase_trace_enabled(): + trace(f"lw/D3 site {i + 1}/{len(layout.my_owned_sites)}: {s} fwd+bwd") + loss_s, n_positions = _layerwise_one_site( + component_model, batch_local, target_local, ci_leaves.per_site, s, strategy + ) + assert loss_s.dim() == 0, f"layerwise loss for site {s!r} must be scalar" + (cfg.coeff_stoch * loss_s / (n_positions * n_sites_total)).backward() + stoch_total_t = stoch_total_t + (loss_s.detach() / n_positions) + return Stoch(total=stoch_total_t, n_owned=len(layout.my_owned_sites)) + + +def _send_g_ci(portals: Portals, layout: ThreePoolLayout, ci_leaves: CiLeaves) -> None: + """Phase lw/D4. Ship per-owned-site CI grads back to the CI pool.""" + assert layout.my_within_block_idx is not None + g_ci_owned = {s: ci_leaves.per_site[s].grad for s in layout.my_owned_sites} + assert all(g is not None for g in g_ci_owned.values()), ( + "layerwise backward should have populated ci_leaves[s].grad" + ) + portals.grad_ci_from_lw.send( + g_ci_owned, + my_within_block_idx=layout.my_within_block_idx, + my_owned_sites=layout.my_owned_sites, + ) + + +def _recv_and_combine_g_vu( + portals: Portals, component_model: ComponentModel, layout: ThreePoolLayout +) -> None: + """Phases lw/D5 + lw/D6. Recv PPGD's V/U grads, add to existing .grad.""" + v_grads_pgd, u_grads_pgd = _recv_g_vu_from_ppgd(portals, layout, component_model) + _combine_vu_grads_in_place(component_model, layout, v_grads_pgd, u_grads_pgd) + + def run_faithfulness_warmup_layerwise( *, component_model: ComponentModel, @@ -222,13 +329,21 @@ def _layerwise_one_site( def _recv_g_vu_from_ppgd( + portals: Portals, layout: ThreePoolLayout, component_model: ComponentModel, ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: """Phase lw/D5. Recv V/U grads from PPGD pool (leader recvs, in-block bcast).""" + assert layout.my_block_idx is not None v_templates = {s: component_model.components[s].V for s in layout.my_owned_sites} u_templates = {s: component_model.components[s].U for s in layout.my_owned_sites} - v_grads_pgd, u_grads_pgd = layout.recv_g_vu_from_ppgd(v_templates, u_templates) + v_grads_pgd, u_grads_pgd = portals.grad_vu_from_ppgd.recv( + v_templates, + u_templates, + my_block_idx=layout.my_block_idx, + my_owned_sites=layout.my_owned_sites, + my_is_block_leader=layout.my_is_block_leader, + ) for s in layout.my_owned_sites: assert v_grads_pgd[s].shape == component_model.components[s].V.shape, ( f"v_grads_pgd[{s!r}] shape mismatch from PPGD send" @@ -256,6 +371,7 @@ def _combine_vu_grads_in_place( def _sync_tail( + portals: Portals, layout: ThreePoolLayout, component_model: ComponentModel, optimizer: torch.optim.Optimizer, @@ -276,18 +392,25 @@ def _sync_tail( n_replicas=layout.world.n_per_block, ) optimizer.step() - _async_send_owned_vu_to_ppgd(component_model, layout) + _async_send_owned_vu_to_ppgd(portals, component_model, layout) -def _async_send_owned_vu_to_ppgd(component_model: ComponentModel, layout: ThreePoolLayout) -> None: - """Kickoff async ship of updated V/U → PPGD. Stash handles on the model.""" +def _async_send_owned_vu_to_ppgd( + portals: Portals, component_model: ComponentModel, layout: ThreePoolLayout +) -> None: + """Kickoff async ship of updated V/U → PPGD. Stash handle on the model.""" + assert layout.my_block_idx is not None v_owned = {s: component_model.components[s].V for s in layout.my_owned_sites} u_owned = {s: component_model.components[s].U for s in layout.my_owned_sites} - weight_send_works, weight_send_buffers = layout.async_send_updated_vu_to_ppgd(v_owned, u_owned) - component_model._pending_weight_sends = ( # type: ignore[attr-defined] - weight_send_works, - weight_send_buffers, + handle = portals.updated_vu_to_ppgd.send( + v_owned, + u_owned, + my_rank=layout.my_rank, + my_block_idx=layout.my_block_idx, + my_owned_sites=layout.my_owned_sites, + my_is_block_leader=layout.my_is_block_leader, ) + component_model._pending_weight_send = handle # type: ignore[attr-defined] def _wait_pending_weight_send(component_model: ComponentModel) -> None: @@ -296,11 +419,10 @@ def _wait_pending_weight_send(component_model: ComponentModel) -> None: Defense against the opt step mutating V/U while the previous async send still reads it. """ - pending = getattr(component_model, "_pending_weight_sends", None) - if pending is not None: - for w in pending[0]: - w.wait() - component_model._pending_weight_sends = None # type: ignore[attr-defined] + handle = getattr(component_model, "_pending_weight_send", None) + if handle is not None: + handle.wait() + component_model._pending_weight_send = None # type: ignore[attr-defined] def _faithfulness_loss( diff --git a/param_decomp_lab/three_pool/step_ppgd.py b/param_decomp_lab/three_pool/step_ppgd.py index e2aff157e..92b98bbed 100644 --- a/param_decomp_lab/three_pool/step_ppgd.py +++ b/param_decomp_lab/three_pool/step_ppgd.py @@ -5,6 +5,13 @@ multi-target ``autograd.grad`` and used to apply one more PGD source step. Total source updates per training step = ``n_warmup_steps + 1``. +The step is a sequence of typed phases. The handoff types make the order a type +constraint: the recon sum (``ReconSum``) requires the refined sources from +warmup, the grads require the recon sum, the scaled grads require the raw grads, +and the cross-pool sends consume the scaled grads — so e.g. "send g_VU before +the in-pool reduce" or "step sources before differentiating" become +unrepresentable. + Phases (numbered to match ``DESIGN.md`` ``ppgd/N``): A1. Post async irecv for CI_T from the owning CI rank (concurrent with A2). @@ -39,6 +46,7 @@ # pyright: reportArgumentType=false +from dataclasses import dataclass from typing import Any import torch @@ -49,11 +57,33 @@ from param_decomp.torch_helpers import bf16_autocast from param_decomp_lab.three_pool.layout import ThreePoolLayout from param_decomp_lab.three_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp_lab.three_pool.portals import PendingPerSiteCi, Portals from param_decomp_lab.three_pool.runtime import _ThreePoolRuntime +@dataclass(frozen=True) +class ReconSum: + """Phases ppgd/D4 (after D3 warmup) output: the UNSCALED Σ-over-this-rank's- + examples recon loss + the example count. The grads phase requires this.""" + + sum_loss: Tensor + n_examples: int + + +@dataclass(frozen=True) +class RawGrads: + """Phase ppgd/D5 output: raw ∂sum_loss/∂· grads, UNSCALED, keyed by site + (V/U/CI) or source key (sources). Each consumer normalizes in place next.""" + + v: dict[str, Tensor] + u: dict[str, Tensor] + ci: dict[str, Tensor] + sources: dict[str, Tensor] + + def step_ppgd( layout: ThreePoolLayout, + portals: Portals, component_model: ComponentModel, ppgd_state: PersistentPGDState, batch: Any, @@ -66,6 +96,7 @@ def step_ppgd( ) -> dict[str, float]: """One PPGD step: A → D → blocking recv of updated V/U from LW → copy in.""" assert layout.my_pool == "ppgd" + assert layout.my_ppgd_slice_idx is not None device = next(component_model.parameters()).device all_sites = list(layout.world.all_sites) @@ -74,89 +105,34 @@ def step_ppgd( v_templates, u_templates = _vu_templates(component_model, all_sites) with strategy.context(component_model.target_model): - ci_recv_pending = layout.async_recv_ci_from_ci_pool_ppgd( - cfg.c_per_site, seq_len=seq_len, device=device + ci_recv_pending = _post_ci_recv(portals, layout, cfg, seq_len, device) + target_out = _target_fwd(component_model, batch_local, cfg) + weight_deltas = component_model.calc_weight_deltas() + ci_scratch = _wait_ci_and_releaf(ci_recv_pending, layout, seq_len, cfg) + recon = _warmup_and_recon( + ppgd_state, component_model, batch_local, target_out, ci_scratch, weight_deltas, cfg ) - with torch.no_grad(), bf16_autocast(cfg.bf16_autocast): - target_out = component_model(batch_local).detach() - weight_deltas = component_model.calc_weight_deltas() - ci_recv = ci_recv_pending.wait_and_unpack() - ci_scratch = _releaf_ci_fp32_for_grads(ci_recv) - _assert_ci_scratch_shapes(ci_scratch, layout, seq_len, cfg) - - # No reduce hook: per-batch-per-position sources are independent per batch - # element, so each PPGD rank's slice is self-contained (asserted at state - # construction in optimize.py). Warmup steps on this rank's own grads. - with bf16_autocast(cfg.bf16_autocast): - ppgd_state.warmup( - model=component_model, - batch=batch_local, - target_out=target_out, - ci=ci_scratch, - weight_deltas=weight_deltas, - ) - with bf16_autocast(cfg.bf16_autocast): - sum_loss, n_examples = ppgd_state.compute_recon_sum_and_n( - model=component_model, - batch=batch_local, - target_out=target_out, - ci=ci_scratch, - weight_deltas=weight_deltas, - ) - assert sum_loss.dim() == 0, f"sum_loss should be scalar; got {sum_loss.shape}" - assert n_examples > 0, f"n_examples must be positive; got {n_examples}" - - # One backward over the UNSCALED recon sum, then each consumer applies its - # OWN normalization explicitly below. The three consumers (V/U, CI, sources) - # need genuinely different scalings; folding any one of them into the - # differentiated scalar — as an earlier version did, dividing by n_ppgd for - # the V/U reduce — silently mis-scales the others (it gave the source step a - # spurious 1/n_ppgd it should never have had). - # These come back UNSCALED (raw ∂sum_loss/∂·); each is normalized in place - # just below, at the point of use. - v_grads, u_grads, ci_grads, source_grads = _autograd_grads_wrt_vu_ci_and_sources( - sum_loss, - component_model, - ci_scratch, - all_sites, - ppgd_state.sources, + raw = _autograd_grads_wrt_vu_ci_and_sources( + recon.sum_loss, component_model, ci_scratch, all_sites, ppgd_state.sources ) - - n_examples_local = n_examples - n_ppgd_ranks = layout.world.n_ppgd - - # V/U and CI reproduce the serial gradient of the canonical PPGD loss - # coeff_ppgd * recon_sum_loss / n_examples_global, - # where n_examples_global = n_examples_local * n_ppgd_ranks. Each rank holds - # only its batch-slice's contribution, so the per-rank scale carries the - # 1/n_ppgd_ranks and the in-pool SUM-reduce reassembles the full-batch grad. - vu_and_ci_grad_scale = cfg.coeff_ppgd / (n_examples_local * n_ppgd_ranks) - # Sources are per-rank-local adversary state optimized against THIS rank's own - # recon mean (recon_sum_loss / n_examples_local) — no coeff, no 1/n_ppgd - # (those are V/U-reduction artifacts). Identical to the warmup source grad. - source_grad_scale = 1.0 / n_examples_local - - _scale_grads_in_place(v_grads, vu_and_ci_grad_scale) - _scale_grads_in_place(u_grads, vu_and_ci_grad_scale) - _scale_grads_in_place(ci_grads, vu_and_ci_grad_scale) - _scale_grads_in_place(source_grads, source_grad_scale) + _scale_grads(raw, recon.n_examples, layout, cfg) # CI grad send first: peer-to-peer (each PPGD rank → its paired CI rank), no # in-pool reduce needed, and CI's recv is on the critical path. Sequencing it # behind the V/U reduce wasted ~110 ms of CI wait time. - layout.send_g_ci_to_ci_pool_ppgd(ci_grads) + portals.grad_ci_from_ppgd.send(raw.ci, my_ppgd_slice_idx=layout.my_ppgd_slice_idx) # V/U grads: SUM-reduce across the pool to reassemble the full-batch gradient # (the per-rank scale already carries 1/n_ppgd_ranks). Sources are NOT bundled # here: their cross-rank reduction is scope-dependent (per_batch_per_position # is per-rank-independent and must not be reduced; a blind SUM would mix # unrelated per-position sources). - layout.sum_reduce_ppgd_grads([*v_grads.values(), *u_grads.values()]) - layout.send_g_vu_to_layerwise(v_grads, u_grads) + layout.sum_reduce_ppgd_grads([*raw.v.values(), *raw.u.values()]) + portals.grad_vu_from_ppgd.send(raw.v, raw.u, is_pool_leader=layout.my_is_pool_leader) # Final (N+1)'th source step. per_batch_per_position sources are per-rank # independent, so no cross-rank reduce — step on this rank's own grads, # exactly as warmup does. - ppgd_state.step(source_grads) + ppgd_state.step(raw.sources) # ``.item()`` calls force CPU↔GPU sync. With async NCCL ops in D5b/D6/D7 # still in flight on side streams, syncing here pulls forward the wait @@ -164,18 +140,119 @@ def step_ppgd( # actually is. Defer these to log steps only. if should_log: metrics = { - "loss/ppgd": (sum_loss / n_examples).item(), - "_raw/ppgd_num": sum_loss.item(), - "_raw/ppgd_den": float(n_examples), + "loss/ppgd": (recon.sum_loss / recon.n_examples).item(), + "_raw/ppgd_num": recon.sum_loss.item(), + "_raw/ppgd_den": float(recon.n_examples), } else: metrics = {} - v_new, u_new = layout.recv_updated_vu_from_layerwise(v_templates, u_templates) + v_new, u_new = portals.updated_vu_to_ppgd.recv(v_templates, u_templates) _copy_vu_into_model_in_place(component_model, v_new, u_new, all_sites) return metrics +def _post_ci_recv( + portals: Portals, + layout: ThreePoolLayout, + cfg: _ThreePoolRuntime, + seq_len: int, + device: torch.device, +) -> PendingPerSiteCi: + """Phase ppgd/A1. Post the async full-model CI-values irecv (waited at D2).""" + assert layout.my_ppgd_slice_idx is not None + return portals.ci_values_to_ppgd.post_recv( + my_ppgd_slice_idx=layout.my_ppgd_slice_idx, + site_to_c=cfg.c_per_site, + seq_len=seq_len, + device=device, + ) + + +def _target_fwd( + component_model: ComponentModel, batch_local: Any, cfg: _ThreePoolRuntime +) -> Tensor: + """Phase ppgd/A2. Detached target forward on this rank's batch slice.""" + with torch.no_grad(), bf16_autocast(cfg.bf16_autocast): + return component_model(batch_local).detach() + + +def _wait_ci_and_releaf( + pending: PendingPerSiteCi, + layout: ThreePoolLayout, + seq_len: int, + cfg: _ThreePoolRuntime, +) -> dict[str, Tensor]: + """Phase ppgd/D2. Wait the CI recv, re-leaf fp32 with grad for D5.""" + ci_recv = pending.wait() + ci_scratch = _releaf_ci_fp32_for_grads(ci_recv) + _assert_ci_scratch_shapes(ci_scratch, layout, seq_len, cfg) + return ci_scratch + + +def _warmup_and_recon( + ppgd_state: PersistentPGDState, + component_model: ComponentModel, + batch_local: Any, + target_out: Tensor, + ci_scratch: dict[str, Tensor], + weight_deltas: dict[str, Tensor], + cfg: _ThreePoolRuntime, +) -> ReconSum: + """Phases ppgd/D3 + ppgd/D4. Warmup refines sources, then the (N+1)'th + recon forward over the refined sources yields the UNSCALED recon sum.""" + # No reduce hook: per-batch-per-position sources are independent per batch + # element, so each PPGD rank's slice is self-contained (asserted at state + # construction in optimize.py). Warmup steps on this rank's own grads. + with bf16_autocast(cfg.bf16_autocast): + ppgd_state.warmup( + model=component_model, + batch=batch_local, + target_out=target_out, + ci=ci_scratch, + weight_deltas=weight_deltas, + ) + with bf16_autocast(cfg.bf16_autocast): + sum_loss, n_examples = ppgd_state.compute_recon_sum_and_n( + model=component_model, + batch=batch_local, + target_out=target_out, + ci=ci_scratch, + weight_deltas=weight_deltas, + ) + assert sum_loss.dim() == 0, f"sum_loss should be scalar; got {sum_loss.shape}" + assert n_examples > 0, f"n_examples must be positive; got {n_examples}" + return ReconSum(sum_loss=sum_loss, n_examples=n_examples) + + +def _scale_grads( + raw: RawGrads, n_examples_local: int, layout: ThreePoolLayout, cfg: _ThreePoolRuntime +) -> None: + """Phase ppgd/D5 (post). Apply each consumer's own normalization in place. + + V/U and CI reproduce the serial gradient of the canonical PPGD loss + coeff_ppgd * recon_sum_loss / n_examples_global, + where n_examples_global = n_examples_local * n_ppgd_ranks. Each rank holds + only its batch-slice's contribution, so the per-rank scale carries the + 1/n_ppgd_ranks and the in-pool SUM-reduce reassembles the full-batch grad. + + Sources are per-rank-local adversary state optimized against THIS rank's own + recon mean (recon_sum_loss / n_examples_local) — no coeff, no 1/n_ppgd + (those are V/U-reduction artifacts). Identical to the warmup source grad. + + Folding any one consumer's scaling into the differentiated scalar — as an + earlier version did, dividing by n_ppgd for the V/U reduce — silently + mis-scales the others (it gave the source step a spurious 1/n_ppgd). + """ + n_ppgd_ranks = layout.world.n_ppgd + vu_and_ci_grad_scale = cfg.coeff_ppgd / (n_examples_local * n_ppgd_ranks) + source_grad_scale = 1.0 / n_examples_local + _scale_grads_in_place(raw.v, vu_and_ci_grad_scale) + _scale_grads_in_place(raw.u, vu_and_ci_grad_scale) + _scale_grads_in_place(raw.ci, vu_and_ci_grad_scale) + _scale_grads_in_place(raw.sources, source_grad_scale) + + def _slice_batch_for_ppgd(batch: Any, layout: ThreePoolLayout) -> tuple[Any, int]: """Pull this PPGD rank's batch slice + extract its seq_len.""" sl = layout.my_batch_slice_ppgd() @@ -238,21 +315,18 @@ def _autograd_grads_wrt_vu_ci_and_sources( ci_scratch: dict[str, Tensor], all_sites: list[str], sources: dict[str, Tensor], -) -> tuple[dict[str, Tensor], dict[str, Tensor], dict[str, Tensor], dict[str, Tensor]]: +) -> RawGrads: """Phase ppgd/D5. One ``torch.autograd.grad`` over the UNSCALED recon sum. Differentiates ``recon_sum_loss`` — the raw Σ-over-this-rank's-examples recon loss, NOT yet divided by any example count or multiplied by any coefficient — once, w.r.t. V/U, the received-CI scratch leaves, and the PPGD sources. The - caller (``step_ppgd``) applies each consumer's own normalization afterward, + caller applies each consumer's own normalization afterward (see ``_scale_grads``), because the three consumers need different scalings and folding any of them into this scalar would mis-scale the others. Fusing all four gradient sets into one backward avoids a separate source-only forward+backward. ``retain_graph=False`` — last use of the graph. - - Returns ``(v_grads, u_grads, ci_grads, source_grads)``, each keyed by site - (sources keyed by source key). """ source_keys = list(sources.keys()) v_params = [component_model.components[s].V for s in all_sites] @@ -283,7 +357,7 @@ def _autograd_grads_wrt_vu_ci_and_sources( assert ci_grads[s].shape == ci_scratch[s].shape, f"ci_grad[{s!r}] shape mismatch" for k in source_keys: assert source_grads[k].shape == sources[k].shape, f"source_grad[{k!r}] shape mismatch" - return v_grads, u_grads, ci_grads, source_grads + return RawGrads(v=v_grads, u=u_grads, ci=ci_grads, sources=source_grads) def _copy_vu_into_model_in_place( diff --git a/scripts/validate_nccl_event_timing.py b/scripts/validate_nccl_event_timing.py index 23b5cd530..08548dd30 100644 --- a/scripts/validate_nccl_event_timing.py +++ b/scripts/validate_nccl_event_timing.py @@ -16,7 +16,6 @@ import torch import torch.distributed as dist - from param_decomp.three_pool import layout as L PAYLOAD_NUMEL = 64 * 1024 * 1024 # 256 MB fp32 — big enough to time transfer