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