From d626960884f9dd659ee26449ad1d0ef09839e077 Mon Sep 17 00:00:00 2001 From: Oliver Clive-Griffin Date: Wed, 27 May 2026 12:34:23 +0000 Subject: [PATCH 001/551] Add 2-pool / 3-pool subsystems with resumption integration --- .github/workflows/checks.yaml | 12 +- Makefile | 4 +- param_decomp/_trace.py | 105 ++ param_decomp/component_model.py | 32 + param_decomp/fused_linear_kl.py | 130 ++ param_decomp/grad_clip.py | 59 + param_decomp/optimize.py | 6 +- param_decomp/run_sink.py | 66 +- param_decomp/sdpa_strict.py | 62 + param_decomp/three_pool/DESIGN.md | 546 +++++++++ param_decomp/three_pool/__init__.py | 33 + param_decomp/three_pool/checkpoint.py | 140 +++ param_decomp/three_pool/config.py | 164 +++ param_decomp/three_pool/layout.py | 947 +++++++++++++++ param_decomp/three_pool/optimize.py | 1058 +++++++++++++++++ param_decomp/three_pool/profiler.py | 120 ++ param_decomp/three_pool/reductions.py | 126 ++ param_decomp/three_pool/runtime.py | 80 ++ param_decomp/three_pool/step_ci.py | 317 +++++ param_decomp/three_pool/step_layerwise.py | 427 +++++++ param_decomp/three_pool/step_ppgd.py | 303 +++++ param_decomp/training_state.py | 44 +- param_decomp/two_pool/__init__.py | 26 + param_decomp/two_pool/config.py | 93 ++ param_decomp/two_pool/layout.py | 774 ++++++++++++ param_decomp/two_pool/loss_strategy.py | 112 ++ param_decomp/two_pool/optimize.py | 681 +++++++++++ param_decomp/two_pool/pool_a.py | 655 ++++++++++ param_decomp/two_pool/pool_b.py | 345 ++++++ param_decomp/two_pool/profiler.py | 86 ++ param_decomp/two_pool/reductions.py | 124 ++ param_decomp/two_pool/runtime.py | 71 ++ param_decomp_lab/experiments/lm/run.py | 152 ++- param_decomp_lab/experiments/resid_mlp/run.py | 3 +- param_decomp_lab/experiments/tms/run.py | 3 +- param_decomp_lab/experiments/utils.py | 17 +- param_decomp_lab/run_sink.py | 94 +- param_decomp_lab/tests/test_gpt2.py | 4 +- param_decomp_lab/tests/test_resid_mlp.py | 4 +- param_decomp_lab/tests/test_resumption.py | 12 +- param_decomp_lab/tests/test_tms.py | 4 +- 41 files changed, 7914 insertions(+), 127 deletions(-) create mode 100644 param_decomp/_trace.py create mode 100644 param_decomp/fused_linear_kl.py create mode 100644 param_decomp/grad_clip.py create mode 100644 param_decomp/sdpa_strict.py create mode 100644 param_decomp/three_pool/DESIGN.md create mode 100644 param_decomp/three_pool/__init__.py create mode 100644 param_decomp/three_pool/checkpoint.py create mode 100644 param_decomp/three_pool/config.py create mode 100644 param_decomp/three_pool/layout.py create mode 100644 param_decomp/three_pool/optimize.py create mode 100644 param_decomp/three_pool/profiler.py create mode 100644 param_decomp/three_pool/reductions.py create mode 100644 param_decomp/three_pool/runtime.py create mode 100644 param_decomp/three_pool/step_ci.py create mode 100644 param_decomp/three_pool/step_layerwise.py create mode 100644 param_decomp/three_pool/step_ppgd.py create mode 100644 param_decomp/two_pool/__init__.py create mode 100644 param_decomp/two_pool/config.py create mode 100644 param_decomp/two_pool/layout.py create mode 100644 param_decomp/two_pool/loss_strategy.py create mode 100644 param_decomp/two_pool/optimize.py create mode 100644 param_decomp/two_pool/pool_a.py create mode 100644 param_decomp/two_pool/pool_b.py create mode 100644 param_decomp/two_pool/profiler.py create mode 100644 param_decomp/two_pool/reductions.py create mode 100644 param_decomp/two_pool/runtime.py diff --git a/.github/workflows/checks.yaml b/.github/workflows/checks.yaml index 29dcfb75c..054f878b0 100644 --- a/.github/workflows/checks.yaml +++ b/.github/workflows/checks.yaml @@ -43,12 +43,17 @@ jobs: - name: Install dependencies run: make install-ci - - name: clear uv cache - run: uv cache clean || true - - name: Print dependencies run: uv pip list + - name: Cache HuggingFace downloads + uses: actions/cache@v4 + with: + path: ${{ github.workspace }}/.hf-cache + key: hf-${{ github.run_id }} + restore-keys: | + hf- + - name: Run basedpyright run: uv run basedpyright @@ -67,6 +72,7 @@ jobs: WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} OMPI_ALLOW_RUN_AS_ROOT: 1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1 + HF_HOME: ${{ github.workspace }}/.hf-cache build-frontend: runs-on: ubuntu-latest defaults: diff --git a/Makefile b/Makefile index 839874a2c..94c30e4b2 100644 --- a/Makefile +++ b/Makefile @@ -32,9 +32,7 @@ check-app: # 1. create a fresh venv with `--clear` -- this is mostly only for local testing of the CI install # 2. install with `uv sync` but with some special options: # > `--frozen` to enforce using the lock file for consistent dependency versions -# > `--link-mode copy` because: -# - symlinks/hardlinks dont work half the time anyway -# - we want to kill the cache after installing before we run the tests +# > `--link-mode copy` because symlinks/hardlinks dont work half the time anyway # > `--extra-index-url` to get cpu-only pytorch wheels. installing with just `uv sync` will download a bunch of cuda stuff we cannot use anyway, since there are no GPUs anyways. takes up a lot of space and makes the install take 5x as long # > `--index-strategy unsafe-best-match` because pytorch index won't have every version of each package we need. markupsafe is a particular pain point # Note: explored the `--compile-bytecode` option for test speedups, nothing came of it. see https://github.com/goodfire-ai/param-decomp/pull/187/commits/740f6a28f4d3378078c917125356b6466f155e71 diff --git a/param_decomp/_trace.py b/param_decomp/_trace.py new file mode 100644 index 000000000..70a4c77e7 --- /dev/null +++ b/param_decomp/_trace.py @@ -0,0 +1,105 @@ +"""Startup + phase tracing for debugging silent stretches. + +The training pipeline has long stretches of pure CPU/CUDA compute that emit no +output for minutes (model build, CI fn build, first-step kernel compile, …). +When the slurm log freezes for 10 min we can't tell a hung job from a slow +one. ``trace(msg)`` is a one-line ``logger.info`` with rank + ms-since-import, +sprinkled at macro boundaries to give a real-time timeline. + +By default every rank logs. Set ``PD_TRACE_RANKS=,,...`` (e.g. +``0,96,100`` for one rank per pool in a 3-pool job) to restrict. + +Phase-level tracing (every ``PhaseProfiler.phase`` enter/exit) is much +noisier so it's opt-in via ``PD_PHASE_TRACE=1``. Combined with +``PD_TRACE_RANKS`` to keep volume sane. +""" + +import os +import sys +import time + +import torch +import torch.distributed as dist + +_TRACE_START = time.perf_counter() + + +def _trace_ranks() -> set[int] | None: + raw = os.environ.get("PD_TRACE_RANKS", "").strip() + if not raw: + return None + return {int(r) for r in raw.split(",") if r.strip()} + + +def _should_log(rank: int) -> bool: + allowed = _trace_ranks() + return allowed is None or rank in allowed + + +def _current_rank() -> int: + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + # torchrun sets RANK before our Python entrypoint runs. Use it so traces + # before init_distributed are still rank-correct (otherwise every rank + # logs as rank=0 and the PD_TRACE_RANKS filter is useless). + env_rank = os.environ.get("RANK") + if env_rank is not None: + return int(env_rank) + return 0 + + +def trace(msg: str) -> None: + """Emit a rank-tagged timeline checkpoint to stdout (force-flushed). + + Uses ``print(..., flush=True)`` rather than ``logger.info`` so the message + bypasses Python's logging stack and forces an immediate fd flush — under + slurm, even with PYTHONUNBUFFERED, logger.info output can get stuck in + Python's logging-layer or torchrun-to-srun pipe buffers for minutes. + + Format: ``[trace rank=R +ELAPSED_MS] MSG`` — easy to grep and tail. + """ + rank = _current_rank() + if not _should_log(rank): + return + elapsed_ms = (time.perf_counter() - _TRACE_START) * 1000.0 + print(f"[trace rank={rank} +{elapsed_ms:9.1f}ms] {msg}", flush=True) + sys.stdout.flush() # belt + braces — slurm log buffering bit us before + + +def phase_trace_enabled() -> bool: + """``PhaseProfiler.phase`` should emit per-phase entry traces.""" + return os.environ.get("PD_PHASE_TRACE", "").strip() in ("1", "true", "yes") + + +def dump_memory_stats(label: str) -> None: + """Emit a single-line summary of ``torch.cuda.memory_stats`` for this rank. + + Includes the headline numbers for capacity planning + fragmentation: + * ``cur``: bytes currently held by tensors + * ``peak``: peak bytes held since the last reset_peak (typically since + start of training, unless ``PhaseProfiler.phase`` is resetting per phase) + * ``reserved``: bytes the CUDA caching allocator is holding (some unused) + * ``free``: caching allocator's free-but-reserved bytes + * ``num_alloc_retries``: count of alloc-then-shrink-cache retries — non-zero + means fragmentation is biting + * ``num_ooms``: count of OOM events the allocator has observed and recovered from + + Gated by ``PD_TRACE_RANKS`` like ``trace()`` so we don't fan out 104 ranks. + """ + rank = _current_rank() + if not _should_log(rank): + return + device = torch.cuda.current_device() + s = torch.cuda.memory_stats(device) + cur = s["allocated_bytes.all.current"] / 1e9 + peak = s["allocated_bytes.all.peak"] / 1e9 + reserved = s["reserved_bytes.all.current"] / 1e9 + free = reserved - cur + elapsed_ms = (time.perf_counter() - _TRACE_START) * 1000.0 + print( + f"[mem rank={rank} +{elapsed_ms:9.1f}ms] {label} " + f"cur={cur:.2f}gb peak={peak:.2f}gb reserved={reserved:.2f}gb free={free:.2f}gb " + f"retries={s.get('num_alloc_retries', 0)} ooms={s.get('num_ooms', 0)}", + flush=True, + ) + sys.stdout.flush() diff --git a/param_decomp/component_model.py b/param_decomp/component_model.py index a2182474d..55c777d0b 100644 --- a/param_decomp/component_model.py +++ b/param_decomp/component_model.py @@ -118,6 +118,36 @@ def __init__( self.lower_leaky_fn = SIGMOID_TYPES[sigmoid_type] self.upper_leaky_fn = SIGMOID_TYPES[sigmoid_type] + def drop_ci_fn(self) -> None: + """Free the CI fn — for pools that only receive CI values via NCCL. + + Call AFTER ``__init__`` (so RNG draws used to init the CI fn stay + identical across pools) but BEFORE ``.to(device)`` (so the unused + params never touch GPU memory). Idempotent. + """ + if self.ci_fn is not None: + del self.ci_fn + self.ci_fn = None + + def drop_components(self) -> None: + """Free the per-site V/U components — for pools that don't run V/U. + + Only safe when no component is an ``EmbeddingComponents`` (the CI fn + wrapper consults those to convert token IDs to acts). Asserted at + runtime. Same lifecycle as ``drop_ci_fn``: post-init, pre-device. + """ + from param_decomp.components import EmbeddingComponents + + assert not any(isinstance(c, EmbeddingComponents) for c in self.components.values()), ( + "drop_components() called on a ComponentModel with embedding components — " + "the CI fn wrapper needs those to convert token IDs to acts." + ) + # Both the in-place clear (for the dict reference held by ci_fn wrapper) + # AND ModuleDict re-init (to unregister the children from nn.Module so + # .to(device) skips them). + self.components.clear() + self._components = nn.ModuleDict() + def target_weight(self, module_name: str) -> Float[Tensor, "rows cols"]: """Weight matrix of a target module in PD's `[d_out, d_in]` row-major convention. @@ -324,6 +354,7 @@ def calc_causal_importances( if detach_inputs: pre_weight_acts = {k: v.detach() for k, v in pre_weight_acts.items()} + assert self.ci_fn is not None, "calc_causal_importances called after drop_ci_fn" ci_fn_outputs = self.ci_fn(pre_weight_acts) return self._apply_sigmoid_to_ci_outputs(ci_fn_outputs, sampling) @@ -404,6 +435,7 @@ def component_grad_norms( ci_fn_grad_norm_sq_sum: Float[Tensor, ""] = torch.zeros((), device=device) missing_ci_fn_grad = False + assert component_model.ci_fn is not None, "component_grad_norms needs the CI fn" for local_param_name, local_param in component_model.ci_fn.named_parameters(): if local_param.grad is None: missing_ci_fn_grad = True diff --git a/param_decomp/fused_linear_kl.py b/param_decomp/fused_linear_kl.py new file mode 100644 index 000000000..2c4017028 --- /dev/null +++ b/param_decomp/fused_linear_kl.py @@ -0,0 +1,130 @@ +"""Memory-efficient fused linear + KL-divergence loss. + +Computes ``KL(softmax(target_logits) || softmax(pred_logits))`` where +``pred_logits = pred_hidden @ lm_head_weight.T`` and ``target_logits = +target_hidden @ lm_head_weight.T``, **without ever materializing the full +[N, vocab] logits tensor.** + +This is the same algorithmic substitute for the unfused pattern +:: + pred_logits = pred_hidden @ lm_head_weight.T # [N, vocab] + target_logits = target_hidden @ lm_head_weight.T # [N, vocab] + loss = recon_loss_kl(pred_logits, target_logits) + +but processes the (N=batch*seq) dimension in chunks so peak memory is +``O(chunk_size * vocab)`` instead of ``O(N * vocab)``. For Qwen vocab=152K and +b=64 s=2048 (N=131K), the unfused path materializes ~40-80 GB of vocab-scale +tensors; this kernel keeps it under ~50 MB at chunk_size=128. + +Design notes: + - Only ``pred_hidden`` requires grad. ``target_hidden`` is the cached target + pre-LM-head activation (frozen target). ``lm_head_weight`` is the target's + frozen LM head — no grad either. + - Backward grad-on-pred-hidden is computed during *forward* and saved for + backward (cheap; just one [N, d_model] tensor). This is the same shape + pattern as the standard autograd path but with chunked vocab-tensor + materialization. + - Inspired by Liger Kernel's ``LigerFusedLinearJSD`` chunked-loss pattern + (Liger has no fused-linear-KL variant as of 2026-01). +""" + +from typing import Any, cast, override + +import torch +import torch.nn.functional as F +from jaxtyping import Float +from torch import Tensor + + +class FusedLinearKLDiv(torch.autograd.Function): + """Fused linear + KL divergence loss with vocab-dimension chunking. + + Forward signature mirrors the unfused pair (matmul + ``recon_loss_kl``). + Backward provides grad w.r.t. ``pred_hidden`` only (the other inputs are + frozen — target's LM head and the cached target hidden states). + """ + + @staticmethod + @override + def forward( + ctx: Any, + pred_hidden: Float[Tensor, "n d_model"], + target_hidden: Float[Tensor, "n d_model"], + lm_head_weight: Float[Tensor, "vocab d_model"], + chunk_size: int = 1024, + ) -> tuple[Float[Tensor, ""], int]: + assert pred_hidden.shape == target_hidden.shape, ( + f"pred_hidden {pred_hidden.shape} vs target_hidden {target_hidden.shape}" + ) + assert pred_hidden.shape[-1] == lm_head_weight.shape[-1], ( + f"d_model mismatch: pred_hidden {pred_hidden.shape[-1]} vs lm_head_weight {lm_head_weight.shape[-1]}" + ) + n, _ = pred_hidden.shape + device = pred_hidden.device + # Accumulator dtype: fp32 for numeric stability when summing many small + # KL contributions. The chunk compute can stay in pred_hidden's dtype + # (typically bf16 under autocast). + total_loss = torch.zeros((), device=device, dtype=torch.float32) + grad_pred_hidden = torch.zeros_like(pred_hidden) + + for start in range(0, n, chunk_size): + end = min(start + chunk_size, n) + ph = pred_hidden[start:end] + th = target_hidden[start:end] + + # Forward: chunk-local logits + KL, no grad inside the chunk (we + # compute the grad-on-pred-hidden contribution analytically below). + with torch.no_grad(): + pred_logits = ph @ lm_head_weight.t() # [chunk, vocab] + target_logits = th @ lm_head_weight.t() # [chunk, vocab] + + log_q = F.log_softmax(pred_logits, dim=-1) + p = F.softmax(target_logits, dim=-1) + # F.kl_div(log_q, p, reduction='none') = p * (log p - log q) + # Summing over vocab gives per-position KL. + kl_per_pos = F.kl_div(log_q, p, reduction="none").sum(dim=-1) + total_loss = total_loss + kl_per_pos.sum().to(torch.float32) + + # Analytical grad of (per-position-KL summed over chunk) w.r.t. + # pred_logits is (softmax(pred_logits) - softmax(target_logits)). + # Chain through the linear: ∂L/∂pred_hidden = (q - p) @ W. + q = log_q.exp() + grad_pred_logits = q - p # [chunk, vocab], same dtype as ph + grad_pred_hidden[start:end] = grad_pred_logits @ lm_head_weight + + ctx.save_for_backward(grad_pred_hidden) + # Return loss + n (mirrors recon_loss_kl's return type). + # Cast back to pred_hidden dtype for downstream sum (matches the + # unfused recon_loss_kl behaviour). + return total_loss.to(pred_hidden.dtype), n + + @staticmethod + @override + def backward( + ctx: Any, + *grad_outputs: Any, + ) -> tuple[Tensor | None, None, None, None]: + grad_loss = grad_outputs[0] + (grad_pred_hidden,) = ctx.saved_tensors + # grad_pred_hidden was computed for loss-coefficient 1.0; scale by the + # incoming gradient on the loss output. + return grad_loss * grad_pred_hidden, None, None, None + + +def fused_linear_kl_div( + pred_hidden: Float[Tensor, "n d_model"], + target_hidden: Float[Tensor, "n d_model"], + lm_head_weight: Float[Tensor, "vocab d_model"], + chunk_size: int = 1024, +) -> tuple[Float[Tensor, ""], int]: + """Functional wrapper around :class:`FusedLinearKLDiv`. + + Returns (sum_of_kl_over_positions, n_positions) — matching the shape of + ``recon_loss_kl`` so the call site is a drop-in replacement once the + upstream code is refactored to expose pred_hidden + target_hidden + + lm_head_weight (instead of pre-materialized logits). + """ + return cast( + "tuple[Float[Tensor, ''] , int]", + FusedLinearKLDiv.apply(pred_hidden, target_hidden, lm_head_weight, chunk_size), + ) diff --git a/param_decomp/grad_clip.py b/param_decomp/grad_clip.py new file mode 100644 index 000000000..3e0280cb9 --- /dev/null +++ b/param_decomp/grad_clip.py @@ -0,0 +1,59 @@ +"""Cross-pool gradient norm clipping. + +Single-pool training uses ``torch.nn.utils.clip_grad_norm_`` directly: it +computes the L2 norm of all the rank's grads and scales them in place if the +norm exceeds the threshold. That's correct when every rank owns the same set +of parameters and the DDP all-reduce has already produced identical grads on +every rank. + +In multi-pool training (2-pool, 3-pool), ranks own DISJOINT parameter subsets +(different sites' V/U + CI fn entries across LW blocks). The "global" L2 norm +that 1-pool clips on is ``sqrt(sum over all ranks of sum_sq_local)``, not a +single rank's local norm. This module provides that reduction. + +Equivalence note: when every rank has the same parameters (DDP within a +block) ``sum_sq_local`` is identical on every replica, so a naive cross-rank +SUM would double-count by ``n_replicas``. Dividing by ``n_replicas`` after +the SUM recovers the global value. +""" + +import torch +import torch.distributed as dist +from torch import Tensor, nn + + +def cross_pool_clip_grad_norm( + params: list[nn.Parameter], + max_norm: float, + *, + group: dist.ProcessGroup, + n_replicas: int, +) -> Tensor: + """L2-norm grad clip across a pool, matching ``clip_grad_norm_``'s semantics. + + Each rank's ``params`` are this rank's owned subset (which may be disjoint + across blocks within the pool). ``n_replicas`` is the number of DDP partners + that hold identical copies of the parameters on this rank — for an LW block + of size ``n_per_block``, this is ``n_per_block``; for a fully-replicated CI + pool, it's ``n_ci``. + + Returns the pre-clip global L2 norm as a 0-dim tensor on the param device, + matching ``torch.nn.utils.clip_grad_norm_``'s return contract. + """ + assert params, "cross_pool_clip_grad_norm called with empty params" + device = params[0].device + sq_local = torch.zeros((), device=device, dtype=torch.float32) + for p in params: + if p.grad is not None: + sq_local = sq_local + (p.grad.detach().to(torch.float32) ** 2).sum() + dist.all_reduce(sq_local, op=dist.ReduceOp.SUM, group=group) + sq_global = sq_local / n_replicas + total_norm = sq_global.sqrt() + # Match torch.nn.utils.clip_grad_norm_'s convention: scale only if norm + # exceeds max_norm; otherwise leave grads untouched. + clip_coef = torch.clamp(max_norm / (total_norm + 1e-6), max=1.0) + if clip_coef.item() < 1.0: + for p in params: + if p.grad is not None: + p.grad.mul_(clip_coef) + return total_norm diff --git a/param_decomp/optimize.py b/param_decomp/optimize.py index 511e8a8ba..e9466a557 100644 --- a/param_decomp/optimize.py +++ b/param_decomp/optimize.py @@ -49,7 +49,7 @@ from param_decomp.metrics.dispatch import instantiate_metrics from param_decomp.metrics.output import collect_metric_outputs from param_decomp.metrics.persistent_pgd_recon import validate_pgd_scope -from param_decomp.run_sink import RunSink +from param_decomp.run_sink import OnePoolRunSink from param_decomp.schedule import get_scheduled_value from param_decomp.torch_helpers import bf16_autocast, loop_dataloader from param_decomp.training_state import TrainingState @@ -491,7 +491,7 @@ def _load_state(self, state: TrainingState) -> None: def run( self, train_loader: DataLoader[Any], - sink: RunSink, + sink: OnePoolRunSink, cadence: Cadence, eval_loop: EvalLoop | None = None, ) -> None: @@ -665,5 +665,3 @@ def run( if is_main_process(): logger.info("Finished training loop.") - - diff --git a/param_decomp/run_sink.py b/param_decomp/run_sink.py index 865e966df..a3c2411c6 100644 --- a/param_decomp/run_sink.py +++ b/param_decomp/run_sink.py @@ -1,28 +1,31 @@ -"""`RunSink` Protocol: where `Trainer.run` sends its output (metrics, console lines, checkpoints).""" +"""Per-pool `RunSink` Protocols. + +Each pool's trainer accepts a typed sink: 1-pool's `Trainer` takes +`OnePoolRunSink`, 3-pool's `ThreePoolTrainer` takes `ThreePoolRunSink`. The +only difference between the two protocols is the `checkpoint` parameter's +state type. `log`, `console`, and `finish` are identical. + +Concrete implementations live in `param_decomp_lab.run_sink` (`OnePoolSink`, +`ThreePoolSink`) — they share a private base for the local-files / wandb / +console plumbing, then add typed `checkpoint` methods that delegate to a +shared `_persist`. + +Timing — when the trainer emits — lives separately: `param_decomp.configs.Cadence` +owns train-log + checkpoint periods, and `param_decomp.optimize.EvalLoop` owns +the eval period. +""" from typing import Any, Protocol, runtime_checkable -from param_decomp.training_state import TrainingState +from param_decomp.training_state import ThreePoolTrainingState, TrainingState @runtime_checkable -class RunSink(Protocol): - """Side-effect sink for a PD training run. - - The trainer treats this object as opaque: it reports what happened, never *where* - the record should go. Callers point the methods at whatever output channels they - want (local files, wandb, S3, a no-op handle on non-main DP ranks, ...). - """ +class OnePoolRunSink(Protocol): + """Side-effect sink for a 1-pool training run.""" def log(self, metrics: dict[str, Any], step: int) -> None: - """Record a flat metrics dict at `step`. - - Args: - metrics: Flat dict whose keys are already namespaced (e.g. - `"train/loss/total"`, `"eval/ci_l0/L0"`) by the trainer. Values may be - scalars, PIL images, or other artefact types the concrete sink supports. - step: Training step at which the values were measured. - """ + """Record a flat metrics dict at ``step``.""" ... def console(self, *lines: str) -> None: @@ -30,15 +33,34 @@ def console(self, *lines: str) -> None: ... def checkpoint(self, snapshot: TrainingState) -> None: - """Persist a training state. + """Persist a 1-pool training state. - `snapshot.step` is the training step; sinks use it for naming. The lab - sink writes `snapshot.component_model` to `model_.pth` and the - whole `snapshot` (cfgs, optimizer state, metric state, etc.) to - `training_.pth`. + The lab sink writes ``snapshot.component_model`` to + ``model_.pth`` (for downstream tools) and the whole ``snapshot`` + to ``training_.pth`` (for resumption). """ ... def finish(self) -> None: """End-of-run cleanup (close handles, finish wandb run, etc.).""" ... + + +@runtime_checkable +class ThreePoolRunSink(Protocol): + """Side-effect sink for a 3-pool training run. + + Identical to `OnePoolRunSink` apart from the checkpoint parameter's state + type. Two separate protocols rather than a union so each trainer's + ``run()`` signature can only accept a sink wired to its own pool's state. + """ + + def log(self, metrics: dict[str, Any], step: int) -> None: ... + def console(self, *lines: str) -> None: ... + def checkpoint(self, snapshot: ThreePoolTrainingState) -> None: ... + def finish(self) -> None: ... + + +# Convenience alias for the few call sites that don't care which pool the sink +# serves (e.g. 2-pool's `run()` while resumption stays unported there). +RunSink = OnePoolRunSink | ThreePoolRunSink diff --git a/param_decomp/sdpa_strict.py b/param_decomp/sdpa_strict.py new file mode 100644 index 000000000..766f80eee --- /dev/null +++ b/param_decomp/sdpa_strict.py @@ -0,0 +1,62 @@ +"""Verify that flash-attention can dispatch on representative SDPA inputs. + +Why not the strict per-call or global toggle? + * Per-call ``sdpa_kernel(FLASH_ATTENTION)`` doesn't compose with + ``torch.compile`` — Dynamo runs the FA dispatch check against + FakeTensors at trace time and fails with ``No available kernel``. + * The global backend toggle (``torch.backends.cuda.enable_*_sdp``) has + the same effect during Dynamo's trace. + +What we do instead: + * Call ``verify_flash_attention_available(...)`` once at trainer init + with the shapes our production SDPA calls use (head_dim, dtype, + is_causal mode). If FA can't dispatch on those shapes, raise loudly. + * After that, leave SDPA's runtime backend selection alone. PyTorch will + pick FA when inputs allow it (and our inputs always will, since they + came from the same config that passed verification). + +The cheat is that PD's SDPA inputs are config-determined and stable across +the run — startup pass therefore implies runtime pass. +""" + +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + + +def verify_flash_attention_available( + *, + head_dim: int, + n_heads: int = 8, + seq_len: int = 64, + is_causal: bool = False, + device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, +) -> None: + """Confirm FA can dispatch on a representative SDPA call. Raises if not. + + Args: + head_dim: per-head dimension. FA caps at 128 (older builds) or 256 + (newer) — pass our largest production head_dim so a too-large + config errors early. + n_heads / seq_len: shape of the test SDPA. Don't matter much for + FA's dispatch decision; defaults keep the test cheap. + is_causal: match the mask mode of the production call. + device: defaults to current CUDA device. + dtype: bf16 by default (FA requires bf16/fp16). + """ + if device is None: + device = torch.device(f"cuda:{torch.cuda.current_device()}") + q = torch.randn(1, n_heads, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + try: + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=is_causal) + except RuntimeError as e: + raise RuntimeError( + f"flash-attention dispatch failed for shape " + f"(n_heads={n_heads}, seq_len={seq_len}, head_dim={head_dim}, " + f"dtype={dtype}, is_causal={is_causal}) — production SDPA calls " + f"will silently fall back to a slower kernel. Underlying error: {e}" + ) from e diff --git a/param_decomp/three_pool/DESIGN.md b/param_decomp/three_pool/DESIGN.md new file mode 100644 index 000000000..2ea932747 --- /dev/null +++ b/param_decomp/three_pool/DESIGN.md @@ -0,0 +1,546 @@ +# Three-pool training — design sketch + +Extension of the 2-pool subsystem (`param_decomp/two_pool/`). Splits the CI +function into its own pool so a **global shared transformer** CI fn becomes +physically realizable again (under 2-pool, sites are sharded across pool-A +ranks, which structurally rules out a CI fn that spans all sites). + +## Pool roles + +| Pool | Owns | Sharded? | Notes | +|------------|---------------------------------|----------|-------| +| **CI** | CI fn + CI-fn optimizer state | DP across batch (replicated CI fn) | Computes canonical CI values + importance-minimality. Multi-rank DP, same pattern as today's PPGD pool. | +| **Layerwise** | V/U + V/U optimizer state | by site (block groups) + DP across batch within a group | Layerwise stoch recon + faithfulness. Same shape as today's Pool A minus the CI fn. | +| **PPGD** | full V/U replica + PPGD sources | DP across batch (replicated V/U) | Full-model PPGD. Same as today's Pool B. Inner-loop warmup owns source updates; final recon backward only seeds V/U + CI grads. | + +## Per-step dependency graph (steady state, step T) + +Each pool's column reads top-to-bottom in time. Solid arrows are within-step +deps; dashed arrows are cross-step. + +```mermaid +flowchart TD + classDef ci fill:#eef6ff,stroke:#3b6db5,color:#000 + classDef lw fill:#f1fbef,stroke:#3e8a4e,color:#000 + classDef pgd fill:#fff4ec,stroke:#b9692a,color:#000 + classDef cross fill:#fff,stroke:#999,color:#000,stroke-dasharray:3 3 + + subgraph CI["CI pool · multi-rank DP across batch"] + direction TB + A0["H_T ready · prefetched in T-1 dead time"]:::ci + A1["A1 · CI fn fwd on H_T → CI_T"]:::ci + A2_LW["A2a · send CI_T per-site → Layerwise
routed by owner + batch slice"]:::ci + A2_PG["A2b · send CI_T full-model → PPGD
routed by batch slice"]:::ci + A3["A3 · imp_min loss on CI_T
backward to leaf grad g_CI_imp"]:::ci + A4["A4 · target_fwd batch T+1 → H_T+1
dead-time fill"]:::ci + A5["A5 · recv g_CI_LW from Layerwise
per-site, per-LW-rank slice"]:::ci + A6["A6 · recv g_CI_PPGD from PPGD
full-model, per-PPGD-rank slice"]:::ci + A7["A7 · assemble g_CI_total per CI rank's batch slice
= g_CI_imp + slice of g_CI_LW + slice of g_CI_PPGD"]:::ci + A8["A8 · backward through CI-fn graph"]:::ci + A9["A9 · in-pool all-reduce on CI fn grads"]:::ci + A10["A10 · AdamW step on CI fn"]:::ci + + A0 --> A1 --> A2_LW + A1 --> A2_PG + A1 --> A3 --> A4 + A4 -.->|"H_T+1 for T+1.A1"| A11_next["T+1 · A1"]:::cross + A5 --> A7 + A6 --> A7 + A3 --> A7 + A7 --> A8 --> A9 --> A10 + A10 -.->|"CI fn weights for T+1.A1"| A11_next + end + + subgraph LW["Layerwise pool · sharded by site · DP within block group"] + direction TB + B0["V/U updated in T-1 dead time"]:::lw + B1["B1 · target_fwd batch T → L_T
per LW rank's batch slice"]:::lw + B2["B2 · wait for CI_T owned sites/slice"]:::lw + B3["B3 · layerwise stoch recon, per owned site, streaming
→ g_VU_LW owned, g_CI_LW owned/slice"]:::lw + B4["B4 · faithfulness loss sharded across owned sites
→ g_VU_faith owned"]:::lw + B5["B5 · send g_CI_LW → CI pool"]:::lw + B6["B6 · recv g_VU_PPGD owned ← PPGD"]:::lw + B7["B7 · combine V/U grads: g_VU_LW + g_VU_faith + g_VU_PPGD"]:::lw + B8["B8 · in-block all-reduce on V/U grads + faithfulness grads"]:::lw + B9["B9 · AdamW step on V/U"]:::lw + B10["B10 · isend updated V/U → PPGD"]:::lw + + B0 --> B1 + B1 --> B3 + B2 --> B3 + B3 --> B5 + B3 --> B7 + B4 --> B7 + B6 --> B7 + B7 --> B8 --> B9 --> B10 + B9 -.->|"V/U for T+1.B3"| B11_next["T+1 · B3"]:::cross + end + + subgraph PG["PPGD pool · DP across batch · replicated V/U"] + direction TB + C0["fresh V/U received in T-1 dead time"]:::pgd + C1["C1 · target_fwd batch T → L_T
per PPGD rank's batch slice"]:::pgd + C2["C2 · wait for CI_T full-model/slice"]:::pgd + C3["C3 · PPGD warmup: refines sources in-place
inner loop owns the source updates"]:::pgd + C4["C4 · PPGD final recon with refined sources"]:::pgd + C5["C5 · backward: g_VU_PPGD, g_CI_PPGD
no source backward at this stage"]:::pgd + C6["C6 · sum-reduce g_VU_PPGD across PPGD ranks"]:::pgd + C7["C7 · send g_VU_PPGD owned → owning LW rank"]:::pgd + C8["C8 · send g_CI_PPGD slice → CI pool"]:::pgd + C9["C9 · recv updated V/U ← Layerwise
completes during T+1's CI window"]:::pgd + + C0 --> C1 + C1 --> C3 + C2 --> C3 + C3 --> C4 --> C5 + C5 --> C6 --> C7 + C5 --> C8 + C9 -.->|"V/U for T+1.C3"| C11_next["T+1 · C3"]:::cross + end + + %% Cross-pool edges within step T + A2_LW -.->|per-site CI values| B2 + A2_PG -.->|full-model CI values| C2 + B5 -.->|g_CI_LW| A5 + C8 -.->|g_CI_PPGD| A6 + C7 -.->|g_VU_PPGD| B6 + B10 -.->|updated V/U| C9 +``` + +## Synchronized timeline (one step T, all three pools share a vertical axis) + +`sequenceDiagram` gives every actor the same vertical time axis, so you can +read across rows to see what each pool is doing at the same moment. Cross-pool +arrows are the actual sends. `Note over` annotations cover work that happens +inside a pool without messaging. + +```mermaid +sequenceDiagram + autonumber + participant CI as CI pool + participant LW as Layerwise pool + participant PG as PPGD pool + + Note over CI: A1 · CI fn fwd on H_T → CI_T + Note over LW: B1 · target_fwd batch T → L_T (LW rank's slice) + Note over PG: C1 · target_fwd batch T → L_T (PPGD rank's slice) + + CI->>LW: A2a · CI_T per-site, owned + LW-rank slice + CI->>PG: A2b · CI_T full-model, PPGD-rank slice + + par CI does imp_min + dead-time prefetch + Note over CI: A3 · imp_min loss → g_CI_imp (leaf grad) + Note over CI: A4 · target_fwd batch T+1 → H_T+1 (prefetch) + and Layerwise recon + Note over LW: B3 · layerwise stoch recon (per owned site, streaming)
→ g_VU_LW owned, g_CI_LW owned/slice + Note over LW: B4 · faithfulness (sharded over owned sites)
→ g_VU_faith owned + and PPGD recon + Note over PG: C3 · PPGD warmup (inner loop owns source updates) + Note over PG: C4-C5 · final recon + bwd
→ g_VU_PPGD full, g_CI_PPGD full + Note over PG: C6 · sum-reduce g_VU_PPGD within PPGD pool + end + + LW->>CI: B5 · g_CI_LW (per-site, per-LW-rank slice) + PG->>CI: C8 · g_CI_PPGD (full-model, per-PPGD-rank slice) + PG->>LW: C7 · g_VU_PPGD (per-owned-site, to owning LW rank) + + Note over CI: A7 · assemble g_CI_total per CI rank's slice + Note over CI: A8 · backward through CI-fn graph + Note over CI: A9 · in-pool all-reduce on CI fn grads + Note over CI: A10 · AdamW step on CI fn + + Note over LW: B7 · combine V/U grads (LW + faith + PPGD) + Note over LW: B8 · in-block all-reduce on V/U grads + + rect rgb(245, 245, 245) + Note over CI,PG: ===== step boundary =====
CI starts T+1.A1 (CI fn fwd) immediately;
LW + PPGD use this window to hide V/U opt + V/U ship-back + end + + Note over CI: T+1.A1 · CI fn fwd on H_T+1 → CI_T+1 + Note over LW: B9 · AdamW step on V/U (hidden behind T+1.A1) + LW-->>PG: B10/C9 · isend updated V/U → PPGD (hidden behind T+1.A1) + + CI->>LW: T+1.A2a · CI_T+1 per-site + CI->>PG: T+1.A2b · CI_T+1 full-model +``` + +The `par` block is where the visible sync shines: all three pools fire in +parallel and you can see at-a-glance that the recon pools' heavy lifting +(B3-B4, C3-C5) happens concurrently with CI's dead-time prefetch (A3-A4). +The `rect` is the step boundary, with the deferred V/U opt + ship-back drawn +as happening *during* T+1's CI-fn-fwd window. + +## Cross-step pipeline (the overlap that hides the V/U opt step) + +```mermaid +gantt + title Three-pool overlap (each row is one rank in that pool) + dateFormat X + axisFormat %s + section CI pool + T.A1 CI fn fwd :ci1, 0, 2 + T.A2 send CI_T :ci2, after ci1, 1 + T.A3 imp_min + leaf grad :ci3, after ci2, 1 + T.A4 prefetch target_fwd T+1 :ci4, after ci3, 4 + T.A7 assemble + bwd CI :ci7, after ci4, 2 + T.A9 all-reduce + opt step :ci9, after ci7, 1 + T+1.A1 CI fn fwd :ci_n, after ci9, 2 + + section Layerwise pool + T.B1 target_fwd T :lw1, 0, 2 + T.B3 layerwise stoch recon :lw3, after lw1, 4 + T.B4 faithfulness :lw4, after lw3, 1 + T.B6/B7 recv g_VU_PPGD + combine :lw6, after lw4, 1 + T.B8/B9 all-reduce + opt step (hides behind T+1.A1) :lw9, after lw6, 2 + T.B10 ship V/U → PPGD :lw10, after lw9, 1 + T+1.B1 target_fwd T+1 :lw_n, after lw10, 2 + + section PPGD pool + T.C1 target_fwd T :pg1, 0, 2 + T.C3 PPGD warmup :pg3, after pg1, 4 + T.C4/C5 final recon + bwd :pg5, after pg3, 1 + T.C6-C8 sum-reduce + sends :pg7, after pg5, 1 + T.C9 recv updated V/U (hides behind T+1.A1) :pg9, after pg7, 2 + T+1.C1 target_fwd T+1 :pg_n, after pg9, 2 +``` + +## Strict cross-step edges + +Only four edges actually force a wait between steps: + +| Edge | Hidden behind | +|---|---| +| `T+1.A1` (CI fn fwd) needs `T.A10` (CI fn AdamW) | — (CI fn fwd kicks off T+1) | +| `T+1.A1` (CI fn fwd) needs `T.A4` (`H_{T+1}` prefetch) | T's recon window | +| `T+1.B3` (Layerwise stoch recon) needs `T.B9` (V/U AdamW) | T+1.A1 (CI fn fwd) | +| `T+1.C3` (PPGD warmup) needs `T.B10` (V/U ship) → `T.C9` recv | T+1.A1 (CI fn fwd) | + +Everything else fits inside step T. + +## Routing complexity vs 2-pool + +The new wrinkle is **3-way batch slicing**. Today both pools either replicate +CI (pool A) or receive a full-model copy (pool B). Under 3-pool: + +- CI rank `i` produces CI values for batch slice `S_ci[i]` and all sites. +- Layerwise rank `j` needs CI values for its owned sites `O[j]` and batch slice `S_lw[j]`. +- PPGD rank `k` needs CI values for all sites and batch slice `S_pgd[k]`. + +In general `S_ci[i]`, `S_lw[j]`, `S_pgd[k]` don't align — so the CI→LW and +CI→PPGD sends become **many-to-many along the batch dim**. Symmetric for grads +coming back. A reasonable simplification for the MVP is to constrain the +batch splits so each LW slice and each PPGD slice fits inside exactly one CI +slice (i.e. choose `N_ci` to divide both `N_lw` and `N_pgd`, or vice versa), +which reduces it to one-to-many fan-out + many-to-one reduction. + +## Open design questions to resolve before coding + +1. **Batch-split divisibility.** Pick an N_ci / N_lw / N_pgd compatibility + rule. Easiest: `N_ci` divides `N_pgd` and `N_lw_per_block`. Validator should + reject anything else loudly. +2. **CI value wire dtype.** Today CI is shipped bf16 and upcast to fp32 on the + Layerwise side. Same pattern likely works for both downstream pools. +3. **Where do imp_min grads enter the CI-fn backward?** Cleanest is: imp_min + is computed inside the CI-fn forward graph and contributes directly to the + single CI-fn backward, with `g_CI_LW + g_CI_PPGD` injected as a `grad_tensors=` + seed on the same backward call (same shape as today's pool-A combined backward). +4. **Validator extensions.** Mirror today's `_validate_pd_config_for_two_pool`: + require `ImportanceMinimalityLoss` lives on CI pool; allow `mode: layerwise` + *or* `mode: global` (with `fn_type: global_shared_transformer`) since CI + ownership is no longer sharded. +5. **Checkpointing.** Still no distributed-aware checkpoint, so `save_every` + stays None for the MVP. + +--- + +# Async pipelining (`defer_vu_opt=True`) + +The MVP runs the LW V/U opt step at end of step T (sync mode). When +`ThreePoolConfig.defer_vu_opt=True`, the LW pool's in-block all_reduce on V/U ++ faith grads is kicked off as **`async_op=True`** at end of step T and waited +at the top of step T+1 — overlapped with target_fwd on the default CUDA +stream. PPGD pool's V/U recv from LW is symmetrically deferred (otherwise +LW's deferred send would deadlock against PPGD's blocking sync recv). + +## Why the async pattern is the actual win + +`optimizer.step()` enqueues AdamW kernels onto the default CUDA stream and +returns to Python immediately — it's not on Python's critical path. The +**blocking** op in the tail is the in-block all_reduce (NCCL collective). +By switching it to `async_op=True` and doing useful compute between +kickoff and wait, we hide most of the all_reduce latency. + +`target_fwd` is the natural overlap candidate: it runs the frozen +target_model (no V/U dependency, no CI dependency), so it can execute on the +default CUDA stream while the all_reduce runs on the NCCL stream. + +**Caveat**: the win is only observable when `N_per_block_lw > 1`. With +`N_per_block_lw == 1` (the example YAML), the in-block all_reduce is a +no-op — the toggle has no observable wall-clock effect. + +## Per-iter execution order in async mode + +```mermaid +flowchart TD + classDef ci fill:#eef6ff,stroke:#3b6db5,color:#000 + classDef lw fill:#f1fbef,stroke:#3e8a4e,color:#000 + classDef pgd fill:#fff4ec,stroke:#b9692a,color:#000 + classDef block fill:#ffe9e0,stroke:#c0392b,color:#000,stroke-width:2px + classDef ovlp fill:#fffacd,stroke:#b8860b,color:#000 + + subgraph LW_T["Layerwise pool — iter T"] + direction TB + LWA1["A1 · post async recv_ci
(NCCL irecv, returns instantly)"]:::lw + LWA2["A2 · target_fwd
(default-stream kernels enqueued)"]:::lw + LWB1["B1 · wait pending V/U send (BLOCKS Python)
— from iter T-1's finalize"]:::block + LWB2["B2 · wait+unflatten async all_reduce (BLOCKS Python)
— from iter T-1's kickoff"]:::block + LWB3["B3 · opt step (T-1's grads, T-1's LR)
mutates V/U; async kernels"]:::lw + LWB4["B4 · async send V/U → PPGD
(NCCL isend, returns instantly)"]:::lw + LWC["C · zero V/U .grad"]:::lw + LWD1["D1 · faith loss + backward
uses fresh V/U"]:::lw + LWD2["D2 · wait recv_ci (BLOCKS Python)"]:::block + LWD3["D3 · layerwise stoch recon (streaming)
uses fresh V/U"]:::lw + LWD4["D4 · send g_CI_LW → CI (blocking isend wait)"]:::lw + LWD5["D5 · recv g_VU from PPGD (BLOCKS Python)"]:::block + LWD6["D6 · combine V/U grads (add PPGD's)"]:::lw + LWE["E · kickoff async all_reduce on combined V/U grads
(returns state for iter T+1's B2)"]:::lw + end + + subgraph OVERLAP["concurrent on LW during B1–B2:"] + direction TB + OV1["GPU default stream: target_fwd kernels running"]:::ovlp + OV2["NCCL stream: prev iter's V/U send completing,
prev iter's all_reduce completing"]:::ovlp + end + + LWA1 --> LWA2 --> LWB1 --> LWB2 --> LWB3 --> LWB4 --> LWC --> LWD1 --> LWD2 --> LWD3 --> LWD4 --> LWD5 --> LWD6 --> LWE + LWA2 -.->|"runs concurrently"| OVERLAP + LWB2 -.->|"target_fwd hides under this wait"| OVERLAP + + %% Cross-step edges + LWE -.->|"state → iter T+1's B2"| LWN["iter T+1's B2"]:::ovlp + LWB4 -.->|"work handle → iter T+1's B1"| LWN +``` + +```mermaid +flowchart TD + classDef ci fill:#eef6ff,stroke:#3b6db5,color:#000 + classDef pgd fill:#fff4ec,stroke:#b9692a,color:#000 + classDef block fill:#ffe9e0,stroke:#c0392b,color:#000,stroke-width:2px + + subgraph PG_T["PPGD pool — iter T"] + direction TB + PGA1["A1 · post async recv_ci"]:::pgd + PGA2["A2 · target_fwd
(kernels enqueued; overlap window)"]:::pgd + PGB["B · wait+unpack prev V/U broadcast (BLOCKS Python)
copy into components — fresh V/U for warmup
(target_fwd runs concurrently on default stream)"]:::block + PGD1["D1 · calc_weight_deltas (fresh V/U)"]:::pgd + PGD2["D2 · wait recv_ci"]:::block + PGD3["D3 · PPGD warmup (inner loop, refines sources)"]:::pgd + PGD4["D4 · final recon loss"]:::pgd + PGD5["D5 · backward → g_VU + g_CI"]:::pgd + PGD6["D6 · in-pool sum-reduce g_VU
(NCCL collective, BLOCKS)"]:::block + PGD7["D7 · send g_VU → LW (leader-only)"]:::pgd + PGD8["D8 · send g_CI → CI (per-rank)"]:::pgd + PGE["E · kickoff async recv V/U from LW
(returns state for iter T+1's B)"]:::pgd + end + + PGA1 --> PGA2 --> PGB --> PGD1 --> PGD2 --> PGD3 --> PGD4 --> PGD5 --> PGD6 --> PGD7 --> PGD8 --> PGE + + PGE -.->|"state → iter T+1's B"| PGN["iter T+1's B"] +``` + +## Synchronized timeline — async mode, all three pools + +Shared vertical time axis. `par` blocks are the genuine concurrency windows. + +```mermaid +sequenceDiagram + autonumber + participant CI as CI pool + participant LW as Layerwise pool + participant PG as PPGD pool + + Note over CI: A1 · CI fn fwd on H_T → CI_T (graph retained) + Note over LW: A1-A2 · post recv_ci + target_fwd
(kernels enqueued on default stream) + Note over PG: A1-A2 · post recv_ci + target_fwd
(kernels enqueued on default stream) + + par CI sends + dead-time prefetch + CI->>LW: A2a · CI_T per-site (sub-sliced) + CI->>PG: A2b · CI_T full-model (sub-sliced) + Note over CI: A4 · target_fwd batch T+1 → H_T+1 + and LW finalizes iter T-1 (waits hide behind target_fwd kernels) + Note over LW: B1-B4 · wait prev send, wait+unflatten all_reduce,
opt step (T-1's grads, T-1's LR), async send V/U + and PG finalizes iter T-1 (wait hides behind target_fwd kernels) + Note over PG: B · wait+unpack prev V/U recv, copy into components + end + + LW-->>PG: async send V/U (kicked off in LW B4) — completes during PG iter T's work + Note over LW: C · zero V/U .grad + Note over PG: D1 · calc_weight_deltas with fresh V/U + + par concurrent recon + imp_min + Note over LW: D1-D3 · faith + layerwise stoch recon (uses fresh V/U) + Note over PG: D2-D4 · wait recv_ci, PPGD warmup, final recon + Note over CI: A3 · imp_min loss → leaf grad + end + + LW->>CI: D4 · g_CI_LW per-owned-site + PG->>CI: D8 · g_CI_PPGD full-model (per-rank slice) + PG->>LW: D7 · g_VU_PPGD per-owned-site (after pool sum-reduce) + + Note over LW: D6 · combine V/U grads + Note over LW: E · kickoff async all_reduce (state for iter T+1's B) + Note over PG: E · kickoff async recv V/U (state for iter T+1's B) + + par CI fused backward + opt + Note over CI: A7-A9 · assemble g_CI, fused bwd, in-pool AVG-reduce, AdamW + and (LW + PG idle — pending async ops on NCCL streams) + Note over LW: (waiting for next iter) + Note over PG: (waiting for next iter) + end + + rect rgb(245, 245, 245) + Note over CI,PG: ===== iter T+1 starts ===== + end + + Note over CI: T+1.A1 · CI fn fwd on H_T+1 → CI_T+1 + Note over LW: T+1.A1-A2 · post recv_ci + target_fwd + Note over PG: T+1.A1-A2 · post recv_ci + target_fwd + Note over LW: T+1.B1-B4 · (now applies iter T's grads with iter T's LR) + Note over PG: T+1.B · (now copies V/U updated by iter T) +``` + +## Where Python actually blocks vs runs async + +Helpful summary for reading the diagrams: what kind of call is each phase? + +| Pool | Phase | Op | Blocks Python? | +|---|---|---|---| +| LW | A1 | `dist.irecv(..., async_op=True)` — `async_recv_ci_from_ci_pool` | No (returns Work handles) | +| LW | A2 | `component_model(batch)` then `.detach()` | No (kernels enqueued, return) | +| LW | B1 | `Work.wait()` on pending V/U send | **Yes** | +| LW | B2 | `Work.wait()` on async all_reduce | **Yes** | +| LW | B3 | `optimizer.step()` | No (kernels enqueued) | +| LW | B4 | `dist.isend` then return | No | +| LW | C | `param.grad = None` | No (Python only) | +| LW | D1 | `(c*loss).backward()` | No (kernels enqueued) | +| LW | D2 | `Work.wait()` on recv_ci | **Yes** | +| LW | D3 | streaming `.backward()` × N sites | No (kernels) | +| LW | D4 | `dist.isend` + `Work.wait()` (sync semantic) | **Yes** | +| LW | D5 | `dist.recv` + `dist.broadcast` (sync) | **Yes** | +| LW | D6 | `.add_()` on grads | No (kernels) | +| LW | E | `dist.all_reduce(async_op=True)` | No (returns Work) | +| CI | A1 | `calc_causal_importances` | No (kernels) | +| CI | A2 | `dist.isend` × N | No | +| CI | A3 | imp_min compute + `dist_fn.all_reduce` (autograd-aware) | **Yes** (the autograd all_reduce blocks) | +| CI | A4 | target_fwd of batch T+1 | No (kernels) | +| CI | A5-6| `recv_g_ci_from_layerwise` + `recv_g_ci_from_ppgd` | **Yes** (sync wait on each `dist.irecv` Work) | +| CI | A7 | tensor `.add` for assembly | No (kernels) | +| CI | A8 | `torch.autograd.backward` | No (kernels) | +| CI | A9 | sync `dist.all_reduce` (in-pool AVG) | **Yes** | +| CI | A10 | `optimizer.step()` | No (kernels) | +| PG | A1 | async `dist.irecv` × N | No | +| PG | A2 | target_fwd | No (kernels) | +| PG | B | `Work.wait()` on prev V/U broadcast | **Yes** | +| PG | D2 | `Work.wait()` on recv_ci | **Yes** | +| PG | D3-D5 | PPGD warmup loop + backward | Mostly kernels; `.item()` calls inside warmup may force CPU sync | +| PG | D6 | sync `dist.all_reduce` (in-pool SUM) | **Yes** | +| PG | D7-D8 | `dist.isend` × N + waits | **Yes** (waits at end) | +| PG | E | async `dist.broadcast` kickoff | No (returns Work) | + +Critical-path blocking points (red boxes in flowchart) are where additional +overlap could win wall-clock. Everything else is already happening +concurrently with whatever other compute/NCCL work has been kicked off. + +# Remaining lack-of-dependency to exploit + +A scan of the above for places where Python blocks on something that could +have been kicked off earlier (or where two blocking ops could be merged). +Listed roughly by leverage. + +### 1. **LW: post async recv_g_vu early, wait at combine** + +`recv_g_vu_from_ppgd` is currently sync at D5. It only blocks because we +haven't posted the irecv. The data dep (`combine` at D6) only needs the +grads at D6 — so we can post the irecv right after the target_fwd kicks off +(A2), letting the recv overlap with everything in B + C + D1-D4. Wait at D5. + +Implementation cost: add async variant in layout.py +(`async_recv_g_vu_from_ppgd_kickoff` + `wait_and_unpack_g_vu`), thread the +state through D5 in step_layerwise. Modest (~50 LOC). + +Expected win: hides recv latency (probably significant if PPGD is the slow +pool — its g_VU send happens after PPGD's full warmup + recon + sum-reduce, +which is the bulk of PPGD's step). Could shave ~30-100 ms per step. + +### 2. **CI: post async recvs of g_CI from BOTH downstream pools concurrently** + +Currently `recv_g_ci_from_layerwise` and `recv_g_ci_from_ppgd` are two +separate sync calls (each itself pipelined internally). Could refactor to +post both pools' irecvs in one call, then wait on all together. The +imp_min + prefetch H_T+1 work happens between A2 and A5 — moving the recv +posts to right after A2 (the sends) lets them race with imp_min compute on +the GPU. + +Implementation cost: similar to (1), ~30 LOC. + +Expected win: hides the recv setup latency. The actual blocking time of +recv is dominated by waiting for LW/PPGD to send their grads; this +refactor doesn't change that, just removes the per-call NCCL setup overhead. +Small but real (~5-10 ms). + +### 3. **CI: async in-pool all_reduce on CI fn grads (mirror the LW pattern)** + +CI's in-pool all_reduce (A9) is currently sync. Same trick as LW: kick off +async at end of iter T, wait at start of iter T+1 — overlapped with CI's +next iter's target_fwd / CI fn fwd kernels. + +Implementation cost: pattern-match what we just did for LW. ~80 LOC. + +Expected win: hides the CI in-pool all_reduce latency. For `N_ci > 1` +(currently `N_ci=1` in the example, so no-op). When multi-rank CI is used, +this is a real win — maybe ~10-30 ms / step. + +### 4. **PPGD: async send_g_vu + send_g_ci can be kicked off in parallel** + +Today (D7, D8) are sequential `dist.isend + wait` blocks. Combining them so +both pools' sends share one wait at end of phase D removes a serialization. + +Implementation cost: ~20 LOC. + +Expected win: small (~5 ms) — both sends are short. + +### 5. **LW: send_g_ci + recv_g_vu overlap** + +D4 (send g_CI to CI) and D5 (recv g_VU from PPGD) currently sequential. +They're independent ops — could post both async then wait on both. + +Implementation cost: ~20 LOC. + +Expected win: small (~5-10 ms), as both are short NCCL ops. + +### 6. **Subsumed: LW + PPGD V/U opt deferred** + +Already implemented as `defer_vu_opt`. This is the canonical example of the +"hide blocking NCCL behind useful compute on a different stream" pattern. + +--- + +The pattern that emerges: **most Python-blocking points in the per-step flow +are NCCL ops, and most can be converted to `async_op=True` kickoff + later +wait with a fixed structural change (split layout method into kickoff + +wait_and_unpack).** The only blocking points that can't be elided this way +are the ones where the data dep is genuinely "I need this immediately to +proceed" — and the obvious example is the `Work.wait()` on recv_ci right +before zero_grad + faith_bwd on LW (we need fresh V/U... wait no, faith +doesn't need CI; only layerwise needs CI). Actually that wait could be +deferred further too — faith bwd could happen before wait_ci, and only +layerwise streaming needs CI. Push wait_ci into D3 (right before layerwise). +Small additional overlap. + +Next session, if profile data backs up these projections, the biggest wins +look like: + +- **(1) async recv_g_vu** — moderate effort, likely significant savings +- **(3) async CI in-pool all_reduce** — only if multi-rank CI + +Defer the rest until profile data shows them on the critical path. diff --git a/param_decomp/three_pool/__init__.py b/param_decomp/three_pool/__init__.py new file mode 100644 index 000000000..128b9ccb5 --- /dev/null +++ b/param_decomp/three_pool/__init__.py @@ -0,0 +1,33 @@ +"""3-pool training strategy for SPD on large frozen target models. + +Splits GPUs into three heterogeneous pools that run in wall-clock parallel: + + CI pool owns the CI fn + AdamW state (replicated across ranks, DP + across batch). Runs target+CI forward, importance-minimality + loss, fused backward through CI fn graph seeded by downstream + pools' CI gradients. Enables global shared CI fns (which + 2-pool's sharded-by-site Pool A cannot). + Layerwise pool owns V/U + AdamW state (sharded by site, block-DDP within + group). Runs target forward, faithfulness loss, per-site + streaming layerwise stoch recon. + PPGD pool stateless full V/U replica + persistent PPGD sources. Runs + target forward, PPGD warmup (inner loop owns source updates), + final recon backward seeding V/U + CI grads only. + +``optimize_three_pool`` mirrors :func:`param_decomp.optimize.optimize`'s call +shape. ``ThreePoolConfig`` declares the topology. + +See ``DESIGN.md`` for the per-step dependency graph + the pipelining tricks. +""" + +from param_decomp.three_pool.config import LayerwiseBlockGroupSpec, ThreePoolConfig +from param_decomp.three_pool.optimize import ThreePoolTrainer, optimize_three_pool +from param_decomp.three_pool.profiler import PhaseProfiler + +__all__ = [ + "LayerwiseBlockGroupSpec", + "PhaseProfiler", + "ThreePoolConfig", + "ThreePoolTrainer", + "optimize_three_pool", +] diff --git a/param_decomp/three_pool/checkpoint.py b/param_decomp/three_pool/checkpoint.py new file mode 100644 index 000000000..c2e4eeea8 --- /dev/null +++ b/param_decomp/three_pool/checkpoint.py @@ -0,0 +1,140 @@ +"""Distributed checkpoint assembly for 3-pool training. + +Rank 0 only holds its own block's V/U (and a random-init CI fn). To produce a +checkpoint that matches the schema ``load_component_model_from_checkpoint`` +expects (one flat ``state_dict`` with every site's V/U + the trained CI fn), +this module gathers state onto rank 0 from: + + * Every LW block leader → its owned sites' V/U. + * The CI pool leader → all CI fn params. + +Rank 0 assembles the gathered tensors into a temporary "full" ``ComponentModel`` +(with all sites in its decomposition targets) and returns that model's +``state_dict()``. The temp model shares ``target_model`` with rank 0's existing +component model — ``ComponentModel.__init__`` doesn't mutate ``target_model``, +so this is safe. + +Non-leader ranks no-op. All ranks must enter ``gather_full_state_dict_to_rank0`` +in sync so the P2P sends/recvs match up. +""" + +# pyright: reportArgumentType=false + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor + +from param_decomp.batch_and_loss_fns import RunBatch +from param_decomp.ci_fns import CiConfig +from param_decomp.ci_sigmoids import SigmoidType +from param_decomp.component_model import ComponentModel +from param_decomp.decomposition_targets import DecompositionTarget +from param_decomp.three_pool.layout import ThreePoolLayout + + +def gather_full_state_dict_to_rank0( + layout: ThreePoolLayout, + component_model: ComponentModel, + target_model: nn.Module, + run_batch: RunBatch, + ci_config: CiConfig, + sigmoid_type: SigmoidType, + c_per_site: dict[str, int], + device: torch.device, +) -> dict[str, Tensor] | None: + """Collect V/U + CI fn onto rank 0 and return the assembled state_dict. + + Returns the state_dict on rank 0; ``None`` everywhere else. All ranks must + call this function in sync (the gather uses ordered P2P sends/recvs). + """ + match layout.my_pool: + case "layerwise" if layout.my_rank == 0: + return _rank0_assemble( + layout=layout, + local_component_model=component_model, + target_model=target_model, + run_batch=run_batch, + ci_config=ci_config, + sigmoid_type=sigmoid_type, + c_per_site=c_per_site, + device=device, + ) + case "layerwise" if layout.my_is_block_leader: + for s in layout.my_owned_sites: + comp = component_model.components[s] + dist.send(comp.V.data.contiguous(), dst=0) + dist.send(comp.U.data.contiguous(), dst=0) + return None + case "ci" if layout.my_is_pool_leader: + assert component_model.ci_fn is not None, "CI pool must keep its CI fn" + for _, p in component_model.ci_fn.named_parameters(): + dist.send(p.data.contiguous(), dst=0) + return None + case _: + # Non-leader LW ranks, non-leader CI ranks, all PPGD ranks: no-op. + return None + + +def _rank0_assemble( + layout: ThreePoolLayout, + local_component_model: ComponentModel, + target_model: nn.Module, + run_batch: RunBatch, + ci_config: CiConfig, + sigmoid_type: SigmoidType, + c_per_site: dict[str, int], + device: torch.device, +) -> dict[str, Tensor]: + """Build a full ComponentModel as an assembly buffer; populate from local + rank's V/U + recvs; return its state_dict.""" + full_targets = [ + DecompositionTarget(module_path=s, C=c_per_site[s]) for s in layout.world.all_sites + ] + # Share target_model with the existing local component_model — ComponentModel + # __init__ doesn't mutate target_model, so two ComponentModels can coexist + # over the same target. + full_cm = ComponentModel( + target_model=target_model, + run_batch=run_batch, + decomposition_targets=full_targets, + ci_config=ci_config, + sigmoid_type=sigmoid_type, + ).to(device) + + # Copy rank 0's own trained V/U into the assembly buffer. + with torch.no_grad(): + for s in layout.my_owned_sites: + full_cm.components[s].V.data.copy_(local_component_model.components[s].V.data) + full_cm.components[s].U.data.copy_(local_component_model.components[s].U.data) + + # Recv V/U from every other LW block leader. Order on both sides must match + # the iteration order of `bg.owned_sites` for each non-rank-0 block leader. + for bg in layout.world.layerwise_block_groups: + if bg.leader == 0: + continue + for s in bg.owned_sites: + V_template = full_cm.components[s].V.data + U_template = full_cm.components[s].U.data + V_buf = torch.empty_like(V_template) + U_buf = torch.empty_like(U_template) + dist.recv(V_buf, src=bg.leader) + dist.recv(U_buf, src=bg.leader) + with torch.no_grad(): + V_template.copy_(V_buf) + U_template.copy_(U_buf) + + # Recv CI fn params from CI pool leader. Same `named_parameters()` iteration + # order on both sides since the CI fn is constructed from the same config. + ci_leader = layout.world.ci_ranks[0] + assert full_cm.ci_fn is not None, "checkpoint reconstruction needs a CI fn" + for _, p in full_cm.ci_fn.named_parameters(): + buf = torch.empty_like(p.data) + dist.recv(buf, src=ci_leader) + with torch.no_grad(): + p.data.copy_(buf) + + state_dict = {k: v.cpu() for k, v in full_cm.state_dict().items()} + del full_cm + torch.cuda.empty_cache() + return state_dict diff --git a/param_decomp/three_pool/config.py b/param_decomp/three_pool/config.py new file mode 100644 index 000000000..e8d986d27 --- /dev/null +++ b/param_decomp/three_pool/config.py @@ -0,0 +1,164 @@ +"""Serializable config for 3-pool training topology. + +``ThreePoolConfig`` carries only what's genuinely 3-pool-specific: which ranks +form the CI pool (replicated CI fn + DP across batch), the Layerwise pool's +block groups (replicated V/U + shared owned sites within a group), and the +PPGD pool (stateless full V/U replica + DP across batch). + +Everything else (batch size, loss coefficients, optimizer LRs, PPGD config, +faithfulness warmup, LR schedules, autocast, ci_config, sigmoid_type, +decomposition_targets) lives on the regular ``PDConfig`` / ``RuntimeConfig`` +that the 3-pool path consumes — so a 3-pool training run is configured like a +normal SPD run plus this topology block. + +See ``DESIGN.md`` for the per-step dependency graph and the rationale for +splitting the CI fn into its own pool (enables global shared transformer CI +fns again — under 2-pool, sites are sharded across pool-A ranks, which +structurally rules out CI fns that span all sites). + +Topology integrity checks (rank disjointness, uniform N_per_block, CI-pool +divisibility into Layerwise/PPGD batch shards) run at validation time; checks +that depend on ``pd.batch_size`` (divisibility) run in +``_validate_pd_config_for_three_pool`` since they need the paired ``PDConfig`` +to evaluate. +""" + +from typing import Self + +from pydantic import Field, model_validator + +from param_decomp.base_config import BaseConfig + + +class LayerwiseBlockGroupSpec(BaseConfig): + """One block-DDP group on the Layerwise pool: ranks that replicate V/U for + a shared set of sites. The first rank is the block leader (canonical actor + for cross-pool sends). Within a group, in-block all-reduce keeps the + replicas in sync after each optimizer step. + + Serializable mirror of ``param_decomp.three_pool.layout.LayerwiseBlockGroup``. + The layout module's dataclass is constructed from this at runtime. + + Identical shape to ``two_pool.config.BlockGroupSpec`` — duplicated rather + than imported so each subsystem owns its own config surface. + """ + + ranks: list[int] = Field(..., description="Ranks that replicate V/U for `owned_sites`.") + owned_sites: list[str] = Field( + ..., description="Module paths in the target that this group owns." + ) + + +class ThreePoolConfig(BaseConfig): + """Topology for 3-pool training. Pairs with a regular ``PDConfig``. + + Every Layerwise-pool rank lives in exactly one ``LayerwiseBlockGroupSpec``. + The CI pool and PPGD pool ranks are listed explicitly. The three pools + must be rank-disjoint. + + Batch-split divisibility (see ``DESIGN.md`` open Q1): + + * ``N_ci`` must divide ``N_per_block_layerwise`` so each Layerwise rank's + batch slice fits inside exactly one CI rank's slice (one-to-many + fan-out for CI values; many-to-one reduction for CI grads). + * ``N_ci`` must divide ``N_ppgd`` for the same reason on the CI↔PPGD + side. + + These two constraints reduce the otherwise many-to-many batch routing to + the simpler one-to-many / many-to-one shape that the comms layer + implements. The validator rejects any other arrangement loudly. + """ + + ci_ranks: list[int] = Field( + ..., + description="Ranks assigned to the CI pool (replicated CI fn, DP across batch). " + "Must divide both N_per_block_layerwise and N_ppgd (see class docstring).", + ) + layerwise_block_groups: list[LayerwiseBlockGroupSpec] = Field( + ..., + description="Layerwise-pool block groups. Each group's ranks replicate V/U for " + "the group's owned sites.", + ) + ppgd_ranks: list[int] = Field( + ..., + description="Ranks assigned to the PPGD pool (stateless full V/U replica, DP " + "across batch).", + ) + use_fused_kl: bool = Field( + default=True, + description="If True (default), Layerwise + PPGD bypass the LM head and " + "compute KL via the chunked fused linear+KL kernel — never materializes " + "[b_local, seq, vocab] tensors. If False, the unfused path is used " + "(materialize logits + standard recon_loss). Same lever as in 2-pool; " + "defaults to fused because the kernel reliably cuts peak memory ~50%% on " + "large-vocab targets with negligible step-time cost.", + ) + defer_vu_opt: bool = Field( + default=False, + description="If True, defer the Layerwise pool's V/U AdamW step + the " + "V/U ship-back to PPGD from end-of-step-T to start-of-step-T+1, so they " + "hide behind T+1's CI fn forward window. Requires symmetric deferral on " + "the PPGD pool (recv_updated_vu also moves to start-of-step-T+1) — " + "otherwise deadlock. Sync and deferred modes are mathematically " + "equivalent (the deferred tail still uses step-T's grads with step-T's " + "LR via `get_scheduled_value(step-1, ...)`); the toggle is purely a " + "wall-clock perf knob. A/B by flipping the flag and comparing " + "`perf/step_ms` traces.", + ) + + @model_validator(mode="after") + def validate_topology(self) -> Self: + assert self.ci_ranks, "ci_ranks must be non-empty" + assert self.layerwise_block_groups, "layerwise_block_groups must be non-empty" + assert self.ppgd_ranks, "ppgd_ranks must be non-empty" + + # Per-block uniformity + non-empty owned sites + n_per_block = len(self.layerwise_block_groups[0].ranks) + assert n_per_block > 0, "layerwise_block_groups[0].ranks must be non-empty" + for bg in self.layerwise_block_groups: + assert len(bg.ranks) == n_per_block, ( + f"all layerwise block groups must have the same N_per_block; got " + f"{n_per_block} vs {len(bg.ranks)} for sites {bg.owned_sites[:2]}..." + ) + assert bg.owned_sites, "layerwise_block_group must own at least one site" + + # Rank disjointness across all three pools + ci_set = set(self.ci_ranks) + assert len(ci_set) == len(self.ci_ranks), "duplicate rank in ci_ranks" + + lw_flat = [r for bg in self.layerwise_block_groups for r in bg.ranks] + lw_set = set(lw_flat) + assert len(lw_set) == len(lw_flat), ( + "the same rank appears in multiple layerwise block groups" + ) + + pgd_set = set(self.ppgd_ranks) + assert len(pgd_set) == len(self.ppgd_ranks), "duplicate rank in ppgd_ranks" + + ci_lw_overlap = sorted(ci_set & lw_set) + ci_pgd_overlap = sorted(ci_set & pgd_set) + lw_pgd_overlap = sorted(lw_set & pgd_set) + assert not ci_lw_overlap, ( + f"CI pool and Layerwise pool must be rank-disjoint; overlap: {ci_lw_overlap}" + ) + assert not ci_pgd_overlap, ( + f"CI pool and PPGD pool must be rank-disjoint; overlap: {ci_pgd_overlap}" + ) + assert not lw_pgd_overlap, ( + f"Layerwise pool and PPGD pool must be rank-disjoint; overlap: {lw_pgd_overlap}" + ) + + # Batch-split divisibility (MVP constraint — see DESIGN.md open Q1). + n_ci = len(self.ci_ranks) + n_ppgd = len(self.ppgd_ranks) + assert n_per_block % n_ci == 0, ( + f"N_ci ({n_ci}) must divide N_per_block_layerwise ({n_per_block}) so each " + f"Layerwise rank's batch slice fits inside exactly one CI rank's slice. " + f"Tip: choose a smaller CI pool, or grow the layerwise block-DDP factor." + ) + assert n_ppgd % n_ci == 0, ( + f"N_ci ({n_ci}) must divide N_ppgd ({n_ppgd}) so each PPGD rank's batch " + f"slice fits inside exactly one CI rank's slice. " + f"Tip: choose a smaller CI pool, or grow the PPGD pool." + ) + return self diff --git a/param_decomp/three_pool/layout.py b/param_decomp/three_pool/layout.py new file mode 100644 index 000000000..e5199f388 --- /dev/null +++ b/param_decomp/three_pool/layout.py @@ -0,0 +1,947 @@ +"""World / ThreePoolLayout — the 3-pool topology data model and cross-pool comms. + +`World` is purely declarative — identical content on every rank, no per-rank +fields. Built once at startup after `dist.init_process_group`. + +`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: + + 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 new wrinkle vs 2-pool is **3-way batch slicing**: CI/LW/PPGD each shard +the global batch on their own axis. The MVP constraint (enforced in +``ThreePoolConfig.validate_topology``) is: + + N_ci | N_per_block_layerwise + N_ci | N_ppgd + +which reduces the otherwise many-to-many CI↔LW and CI↔PPGD routing to a +one-to-many fan-out (CI rank → K downstream ranks) plus a many-to-one +reduction (downstream ranks → one CI rank). See "Batch-split routing" below. +""" + +# pyright: reportIndexIssue=false, reportArgumentType=false, reportOperatorIssue=false + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any, Literal + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor + +# 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 + + +@dataclass(frozen=True) +class LayerwiseBlockGroup: + """One LW block-DDP group: ranks that replicate V/U for `owned_sites`. + + The first rank is the block leader — the canonical actor for cross-pool + sends. Within a group, in-block all-reduce keeps the replicas in sync. + Runtime mirror of ``LayerwiseBlockGroupSpec``. + """ + + ranks: tuple[int, ...] + owned_sites: tuple[str, ...] + + @property + def leader(self) -> int: + return self.ranks[0] + + def __post_init__(self) -> None: + assert len(self.ranks) > 0, "block group must have at least one rank" + assert len(self.ranks) == len(set(self.ranks)), ( + f"duplicate ranks in block group: {self.ranks}" + ) + + +@dataclass(frozen=True) +class World: + """Declarative 3-pool topology — identical content on every rank. + + Three rank-disjoint pools: + + * ``ci_ranks``: replicate the CI fn, DP across batch. + * ``layerwise_block_groups``: per-block replication of V/U (DDP within + block) + per-block batch sharding across ranks-within-block. + * ``ppgd_ranks``: replicate the full V/U, DP across batch. + + Process groups are constructed at world-build time and stored here so the + layout doesn't have to plumb them through call sites. + """ + + world_size: int + ci_ranks: tuple[int, ...] + layerwise_block_groups: tuple[LayerwiseBlockGroup, ...] + ppgd_ranks: tuple[int, ...] + all_sites: tuple[str, ...] + batch_global: int + + ci_pool_group: dist.ProcessGroup + layerwise_pool_group: dist.ProcessGroup + ppgd_pool_group: dist.ProcessGroup + block_group_groups: tuple[dist.ProcessGroup, ...] + # One process group per LW block: {block_leader} ∪ {ppgd_ranks}. Used for + # leader-rooted broadcasts when shipping updated V/U from LW → PPGD pool. + # Matches the cross_pool_bcast_groups pattern in two_pool. + cross_pool_bcast_groups: tuple[dist.ProcessGroup, ...] + + # ── Sizes ── + + @property + def n_ci(self) -> int: + return len(self.ci_ranks) + + @property + def n_ppgd(self) -> int: + return len(self.ppgd_ranks) + + @property + def n_blocks(self) -> int: + return len(self.layerwise_block_groups) + + @property + def n_per_block(self) -> int: + size = len(self.layerwise_block_groups[0].ranks) + assert all(len(bg.ranks) == size for bg in self.layerwise_block_groups) + return size + + @property + def n_layerwise(self) -> int: + return sum(len(bg.ranks) for bg in self.layerwise_block_groups) + + @property + def layerwise_ranks(self) -> tuple[int, ...]: + return tuple(r for bg in self.layerwise_block_groups for r in bg.ranks) + + @property + def batch_local_ci(self) -> int: + assert self.batch_global % self.n_ci == 0 + return self.batch_global // self.n_ci + + @property + def batch_local_lw(self) -> int: + assert self.batch_global % self.n_per_block == 0 + return self.batch_global // self.n_per_block + + @property + def batch_local_ppgd(self) -> int: + assert self.batch_global % self.n_ppgd == 0 + return self.batch_global // self.n_ppgd + + # ── Site routing ── + + def block_idx_of_site(self, site: str) -> int: + for i, bg in enumerate(self.layerwise_block_groups): + if site in bg.owned_sites: + return i + raise KeyError(site) + + def block_leader_of_site(self, site: str) -> int: + return self.layerwise_block_groups[self.block_idx_of_site(site)].leader + + # ── Batch-split routing (the new wrinkle) ── + + @property + def k_lw_per_ci(self) -> int: + """How many LW batch shards (per block) fit inside one CI batch shard. + + Validator guarantees N_ci | N_per_block_layerwise, so this is an integer. + """ + assert self.n_per_block % self.n_ci == 0 + return self.n_per_block // self.n_ci + + @property + def k_ppgd_per_ci(self) -> int: + """How many PPGD batch shards fit inside one CI batch shard. + + Validator guarantees N_ci | N_ppgd, so this is an integer. + """ + assert self.n_ppgd % self.n_ci == 0 + return self.n_ppgd // self.n_ci + + def ci_slice_of_lw_block_rank(self, block_rank_idx: int) -> int: + """Which CI rank's slice contains LW rank `block_rank_idx`'s batch shard.""" + return block_rank_idx // self.k_lw_per_ci + + def ci_slice_of_ppgd_slice(self, ppgd_slice_idx: int) -> int: + """Which CI rank's slice contains PPGD rank `ppgd_slice_idx`'s batch shard.""" + return ppgd_slice_idx // self.k_ppgd_per_ci + + def lw_sub_slice_within_ci(self, block_rank_idx: int) -> slice: + """Within a CI rank's local batch tensor [B_ci, ...], the sub-slice + that corresponds to LW rank `block_rank_idx`. + + Assumes the caller is the CI rank with + `ci_slice_idx == ci_slice_of_lw_block_rank(block_rank_idx)`. + """ + b_lw = self.batch_local_lw + ci_slice_idx = self.ci_slice_of_lw_block_rank(block_rank_idx) + global_start = block_rank_idx * b_lw + ci_global_start = ci_slice_idx * self.batch_local_ci + local_start = global_start - ci_global_start + return slice(local_start, local_start + b_lw) + + def ppgd_sub_slice_within_ci(self, ppgd_slice_idx: int) -> slice: + """Within a CI rank's local batch tensor [B_ci, ...], the sub-slice + that corresponds to PPGD rank `ppgd_slice_idx`.""" + b_pp = self.batch_local_ppgd + ci_slice_idx = self.ci_slice_of_ppgd_slice(ppgd_slice_idx) + global_start = ppgd_slice_idx * b_pp + ci_global_start = ci_slice_idx * self.batch_local_ci + local_start = global_start - ci_global_start + return slice(local_start, local_start + b_pp) + + def lw_block_ranks_for_ci_slice(self, ci_slice_idx: int) -> tuple[int, ...]: + """LW block_rank_idxs whose batch shards sit inside CI rank `ci_slice_idx`'s slice.""" + k = self.k_lw_per_ci + return tuple(range(ci_slice_idx * k, (ci_slice_idx + 1) * k)) + + def ppgd_slice_idxs_for_ci_slice(self, ci_slice_idx: int) -> tuple[int, ...]: + """PPGD slice idxs whose batch shards sit inside CI rank `ci_slice_idx`'s slice.""" + k = self.k_ppgd_per_ci + return tuple(range(ci_slice_idx * k, (ci_slice_idx + 1) * k)) + + +def build_world( + ci_ranks: list[int], + layerwise_block_groups: list[LayerwiseBlockGroup], + ppgd_ranks: list[int], + batch_global: int, + device: torch.device | None = None, +) -> World: + """Construct the World + all process groups. Must be called on every rank + after ``dist.init_process_group``. + + Pass ``device`` (this rank's GPU) so we can pre-warm the cross-pool NCCL + broadcast groups before they're first used inside the training loop. See + ``_prewarm_cross_pool_bcast_groups`` for why pre-warming is needed. + """ + world_size = dist.get_world_size() + layerwise_ranks = [r for bg in layerwise_block_groups for r in bg.ranks] + assert len(ci_ranks) + len(layerwise_ranks) + len(ppgd_ranks) == world_size, ( + f"rank count mismatch: ci={len(ci_ranks)} + lw={len(layerwise_ranks)} + " + f"ppgd={len(ppgd_ranks)} != world_size={world_size}" + ) + assert set(ci_ranks).isdisjoint(set(layerwise_ranks)) + assert set(ci_ranks).isdisjoint(set(ppgd_ranks)) + assert set(layerwise_ranks).isdisjoint(set(ppgd_ranks)) + assert len(set(layerwise_ranks)) == len(layerwise_ranks) + + all_sites = tuple(s for bg in layerwise_block_groups for s in bg.owned_sites) + assert len(set(all_sites)) == len(all_sites), "a site is owned by more than one block group" + + my_rank = dist.get_rank() + + def _make_group(name: str, ranks: list[int]) -> Any: + # Per-rank trace before/after each collective new_group. Lets us + # localize precisely where the world wedges if NCCL deadlocks. + print(f"[build_world rank={my_rank}] before {name} ranks={ranks}", flush=True) + g = dist.new_group(ranks=ranks) + print(f"[build_world rank={my_rank}] after {name} ranks={ranks}", flush=True) + return g + + ci_pool_group = _make_group("ci_pool_group", ci_ranks) + layerwise_pool_group = _make_group("layerwise_pool_group", layerwise_ranks) + ppgd_pool_group = _make_group("ppgd_pool_group", ppgd_ranks) + block_group_groups = tuple( + _make_group(f"block_group_groups[{i}]", list(bg.ranks)) + for i, bg in enumerate(layerwise_block_groups) + ) + cross_pool_bcast_groups = tuple( + _make_group(f"cross_pool_bcast_groups[{i}]", [bg.leader, *ppgd_ranks]) + for i, bg in enumerate(layerwise_block_groups) + ) + + if device is not None: + _prewarm_cross_pool_bcast_groups( + cross_pool_bcast_groups=cross_pool_bcast_groups, + layerwise_block_groups=layerwise_block_groups, + ppgd_ranks=ppgd_ranks, + my_rank=my_rank, + device=device, + ) + + return World( + world_size=world_size, + ci_ranks=tuple(ci_ranks), + layerwise_block_groups=tuple(layerwise_block_groups), + ppgd_ranks=tuple(ppgd_ranks), + all_sites=all_sites, + batch_global=batch_global, + ci_pool_group=ci_pool_group, + layerwise_pool_group=layerwise_pool_group, + ppgd_pool_group=ppgd_pool_group, + block_group_groups=block_group_groups, + cross_pool_bcast_groups=cross_pool_bcast_groups, + ) + + +def _prewarm_cross_pool_bcast_groups( + *, + cross_pool_bcast_groups: tuple[Any, ...], + layerwise_block_groups: list[LayerwiseBlockGroup], + ppgd_ranks: list[int], + my_rank: int, + device: torch.device, +) -> None: + """Trigger NCCL communicator init on each cross-pool bcast group. + + First use of a new NCCL process group blocks on a synchronous global + communicator init across all participating ranks — even when the user + passes ``async_op=True``. In the training loop, PPGD's + ``E_kickoff_async_recv_vu`` is the first user of these groups (it does + irecv-side broadcasts for the V/U-from-LW pipeline that defer_vu_opt + introduces). The matching send-side broadcast doesn't fire until LW + step N+1 phase B4 — which means on log steps (``train_log_every`` is up), + we end up in a deadlock: + + * LW rank 0 stuck in ``dist.recv`` from PPGD leader inside + ``_log_train_metrics`` (PPGD leader hasn't sent yet). + * PPGD leader stuck inside the first ``dist.broadcast`` of E_kickoff + (communicator init waiting for LW block leaders to call into NCCL). + * LW step N+1 (and hence phase B4) can't start until + ``_log_train_metrics`` returns. + + Pre-warming each group here with a 1-element dummy broadcast does the + NCCL init once at setup time, while every participant is still in + lockstep at ``build_world``. After this, the first real broadcast in the + training loop is a normal communicator op that can interleave with other + work. + """ + dummy = torch.zeros(1, device=device) + ppgd_set = set(ppgd_ranks) + for bg, group in zip(layerwise_block_groups, cross_pool_bcast_groups, strict=True): + if my_rank == bg.leader or my_rank in ppgd_set: + dist.broadcast(dummy, src=bg.leader, group=group) + + +@dataclass(frozen=True) +class ThreePoolLayout: + """This rank's view of the 3-pool world + cross-pool comm methods. + + Single class with ``my_pool`` switch (mirrors ``two_pool.BlockDDPLayout``). + Each comm method asserts its caller pool up-front so misuse is loud. + """ + + world: World + my_rank: int + my_pool: Literal["ci", "layerwise", "ppgd"] + + # LW-only fields + my_block_idx: int | None + my_within_block_idx: int | None + my_is_block_leader: bool + my_owned_sites: tuple[str, ...] + + # CI-only / PPGD-only fields + my_ci_slice_idx: int | None + my_ppgd_slice_idx: int | None + my_is_pool_leader: bool + + @classmethod + def from_world(cls, world: World, my_rank: int) -> "ThreePoolLayout": + if my_rank in world.ci_ranks: + slice_idx = world.ci_ranks.index(my_rank) + return cls( + world=world, + my_rank=my_rank, + my_pool="ci", + my_block_idx=None, + my_within_block_idx=None, + my_is_block_leader=False, + my_owned_sites=(), + my_ci_slice_idx=slice_idx, + my_ppgd_slice_idx=None, + my_is_pool_leader=(my_rank == world.ci_ranks[0]), + ) + for bg_idx, bg in enumerate(world.layerwise_block_groups): + if my_rank in bg.ranks: + within = bg.ranks.index(my_rank) + return cls( + world=world, + my_rank=my_rank, + my_pool="layerwise", + my_block_idx=bg_idx, + my_within_block_idx=within, + my_is_block_leader=(within == 0), + my_owned_sites=bg.owned_sites, + my_ci_slice_idx=None, + my_ppgd_slice_idx=None, + my_is_pool_leader=False, + ) + if my_rank in world.ppgd_ranks: + slice_idx = world.ppgd_ranks.index(my_rank) + return cls( + world=world, + my_rank=my_rank, + my_pool="ppgd", + my_block_idx=None, + my_within_block_idx=None, + my_is_block_leader=False, + my_owned_sites=(), + my_ci_slice_idx=None, + my_ppgd_slice_idx=slice_idx, + my_is_pool_leader=(my_rank == world.ppgd_ranks[0]), + ) + raise ValueError(f"rank {my_rank} not in any pool") + + # ── Per-rank slice helpers ── + + def my_batch_slice_ci(self) -> slice: + assert self.my_pool == "ci" and self.my_ci_slice_idx is not None + b = self.world.batch_local_ci + return slice(self.my_ci_slice_idx * b, (self.my_ci_slice_idx + 1) * b) + + def my_batch_slice_lw(self) -> slice: + assert self.my_pool == "layerwise" and self.my_within_block_idx is not None + b = self.world.batch_local_lw + return slice(self.my_within_block_idx * b, (self.my_within_block_idx + 1) * b) + + def my_batch_slice_ppgd(self) -> slice: + assert self.my_pool == "ppgd" and self.my_ppgd_slice_idx is not None + b = self.world.batch_local_ppgd + return slice(self.my_ppgd_slice_idx * b, (self.my_ppgd_slice_idx + 1) * b) + + def is_my_site(self, site: str) -> bool: + return self.my_pool == "layerwise" and site in self.my_owned_sites + + def i_lead_site(self, site: str) -> bool: + return self.is_my_site(site) and self.my_is_block_leader + + # ────────────────────────────────────────────────────────────────────── + # CI-pool comm methods + # ────────────────────────────────────────────────────────────────────── + + 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) + + 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) + for site in bg.owned_sites: + buf = ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous() + works.append(dist.isend(buf, dst=target_lw_rank)) + buffers.append(buf) + 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) + + 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) + for site in self.world.all_sites: + buf = ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous() + works.append(dist.isend(buf, dst=target_ppgd_rank)) + buffers.append(buf) + return works, buffers + + 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, ...]]] = [] + 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) + 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) + 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]] = [] + 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) + 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 + } + 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. + """ + assert self.my_pool == "ci" + if dist.get_world_size(self.world.ci_pool_group) <= 1: + return + grads: list[Tensor] = [p.grad for p in params if p.grad is not None] + if not grads: + return + buckets: dict[tuple[torch.dtype, torch.device], list[Tensor]] = {} + for g in grads: + buckets.setdefault((g.dtype, g.device), []).append(g) + from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + + for bucket in buckets.values(): + flat = _flatten_dense_tensors(bucket) + dist.all_reduce(flat, op=dist.ReduceOp.AVG, group=self.world.ci_pool_group) + for orig, reduced in zip(bucket, _unflatten_dense_tensors(flat, bucket), strict=True): + 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, + ) -> tuple[dict[str, Tensor], list["dist.Work"]]: + """LW ← CI: irecv per-owned-site CI values from the CI rank whose + slice contains my LW batch shard. Returns raw bf16 buffers + work + handles; caller waits + casts to fp32.""" + 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 + + out: dict[str, Tensor] = {} + works: list[dist.Work] = [] + for site in self.my_owned_sites: + C = site_to_c[site] + buf = torch.empty(b_lw, seq_len, C, device=device, dtype=_WIRE_DTYPE) + w = dist.irecv(buf, src=src) + assert w is not None + out[site] = buf + works.append(w) + return out, works + + 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) + dist.send(packed, dst=dst) + + 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] + dist.recv(packed, src=ppgd_leader) + 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 + 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. Mirrors two_pool's async_send_updated_weights_to_pool_b. + 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] + 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. + + Synchronous variant — Python blocks until the collective completes. + Identical pattern to ``two_pool.BlockDDPLayout.all_reduce_grads_in_block``. + Used by the sync path in ``step_layerwise_tail``. + """ + states = self.async_all_reduce_grads_in_block_kickoff(params) + self.wait_and_unflatten_all_reduce(states) + + def async_all_reduce_grads_in_block_kickoff( + self, params: Iterable[nn.Parameter] + ) -> list[tuple[list[Tensor], Tensor, "dist.Work"]]: + """Kick off async coalesced in-block all-reduce. Returns one + ``(bucket, flat, work)`` tuple per (dtype, device) bucket. Caller MUST + keep these alive (storing them across iteration boundaries) until + ``wait_and_unflatten_all_reduce`` runs — the ``flat`` tensors are the + NCCL buffers, and freeing them while NCCL is still operating would be + undefined. + + Empty return when the block group is 1-rank (no-op) or no grads. + + The async variant lets the caller overlap the all-reduce with other + compute on the default CUDA stream (e.g. the next iteration's + target_fwd, which is V/U-independent and runs on a different stream). + """ + assert self.my_pool == "layerwise" and self.my_block_idx is not None + block_group = self.world.block_group_groups[self.my_block_idx] + if dist.get_world_size(block_group) <= 1: + return [] + grads: list[Tensor] = [p.grad for p in params if p.grad is not None] + if not grads: + return [] + buckets: dict[tuple[torch.dtype, torch.device], list[Tensor]] = {} + for g in grads: + buckets.setdefault((g.dtype, g.device), []).append(g) + from torch._utils import _flatten_dense_tensors + + states: list[tuple[list[Tensor], Tensor, dist.Work]] = [] + for bucket in buckets.values(): + flat = _flatten_dense_tensors(bucket) + w = dist.all_reduce(flat, op=dist.ReduceOp.AVG, group=block_group, async_op=True) + assert w is not None + states.append((bucket, flat, w)) + return states + + def wait_and_unflatten_all_reduce( + self, + states: list[tuple[list[Tensor], Tensor, "dist.Work"]], + ) -> None: + """Wait on each async all-reduce work and copy the reduced flat tensor + back to the original ``.grad`` buffers.""" + from torch._utils import _unflatten_dense_tensors + + for bucket, flat, w in states: + w.wait() + 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, + ) -> tuple[dict[str, Tensor], list["dist.Work"]]: + """PPGD ← CI: irecv full-model CI from the CI rank that owns my slice.""" + 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 + + out: dict[str, Tensor] = {} + works: list[dist.Work] = [] + for site in self.world.all_sites: + C = site_to_c[site] + buf = torch.empty(b_pp, seq_len, C, device=device, dtype=_WIRE_DTYPE) + w = dist.irecv(buf, src=src) + assert w is not None + out[site] = buf + works.append(w) + return out, works + + 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) + dist.send(packed, dst=dst) + + 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] = [] + 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) + assert w is not None + works.append(w) + buffers.append(packed) + 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. + + Coalesced bucketing like the other in-pool reductions. + """ + assert self.my_pool == "ppgd" + if dist.get_world_size(self.world.ppgd_pool_group) <= 1: + return + grads_list = list(grads) + if not grads_list: + return + buckets: dict[tuple[torch.dtype, torch.device], list[Tensor]] = {} + for g in grads_list: + buckets.setdefault((g.dtype, g.device), []).append(g) + from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + + for bucket in buckets.values(): + flat = _flatten_dense_tensors(bucket) + dist.all_reduce(flat, op=dist.ReduceOp.SUM, group=self.world.ppgd_pool_group) + for orig, reduced in zip(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. + + Synchronous variant — kicks off all broadcasts then waits + unpacks. + Use ``async_recv_updated_vu_from_layerwise_kickoff`` + + ``wait_and_unpack_updated_vu`` to overlap with PPGD's target_fwd. + """ + states = self.async_recv_updated_vu_from_layerwise_kickoff(v_templates, u_templates) + return self.wait_and_unpack_updated_vu(states, v_templates, u_templates) + + def async_recv_updated_vu_from_layerwise_kickoff( + self, + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + ) -> list[tuple["LayerwiseBlockGroup", Tensor, "dist.Work"]]: + """Kick off async coalesced V/U recv from every LW block leader. Returns + per-block ``(block_group, packed_buf, work)`` tuples. Caller holds + these across iteration boundaries until + ``wait_and_unpack_updated_vu`` runs. + + Overlap target: PPGD's next-iter target_fwd, which runs on the default + CUDA stream and doesn't depend on V/U. The broadcasts run on NCCL + streams (one per block group's bcast group), so they pipeline. + """ + assert self.my_pool == "ppgd" + bufs: list[tuple[LayerwiseBlockGroup, Tensor, dist.Work]] = [] + 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)) + return bufs + + def wait_and_unpack_updated_vu( + self, + states: list[tuple["LayerwiseBlockGroup", Tensor, "dist.Work"]], + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + """Wait on each block's broadcast and unpack the contiguous packed + tensor back into per-site V/U dicts (upcasting to the templates' dtype). + Returns ``(v_new, u_new)`` ready for ``components[s].V.copy_(...)``. + """ + v_new: dict[str, Tensor] = {} + u_new: dict[str, Tensor] = {} + for bg, packed, w in states: + w.wait() + owned = bg.owned_sites + 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 diff --git a/param_decomp/three_pool/optimize.py b/param_decomp/three_pool/optimize.py new file mode 100644 index 000000000..cd1c8a762 --- /dev/null +++ b/param_decomp/three_pool/optimize.py @@ -0,0 +1,1058 @@ +"""``ThreePoolTrainer`` and ``optimize_three_pool`` — 3-pool sibling of +:class:`param_decomp.optimize.Trainer` / :func:`param_decomp.optimize.optimize`. + +Mirrors the single-pool call shape: caller hands in ``target_model``, +dataloader, configs, sink. Internal validation, per-pool wiring, cross-pool +comms, and the layerwise streaming loss strategy are all hidden behind the +class boundary. + + * **CI pool** trains the CI fn (replicated across ranks; DP-sharded across + batch). Holds CI fn + AdamW state. Each step: target_fwd → CI fn fwd → + broadcast CI to LW + PPGD → dead-time prefetch H_{T+1} → fused backward + seeded by imp_min + per-site g_CI from LW + PPGD → in-pool all-reduce → + AdamW. See :mod:`param_decomp.three_pool.step_ci`. + + * **Layerwise pool** trains V/U (block-DDP within group; sharded across + sites). Recv CI → faithfulness + layerwise stoch recon → send g_CI back + → recv g_VU from PPGD → combine → in-block all-reduce → AdamW → async + ship updated V/U → PPGD. See :mod:`param_decomp.three_pool.step_layerwise`. + + * **PPGD pool** is a stateless full V/U replica. Recv CI → PPGD warmup + + final recon → backward seeds V/U + CI grads (no outer source step; + warmup inner loop owns sources) → sum-reduce V/U within PPGD pool → + send g_VU to LW + g_CI to CI → recv updated V/U. See + :mod:`param_decomp.three_pool.step_ppgd`. + +Data-handling contract +---------------------- +Every rank must read the FULL global batch on every step (i.e. each batch +tensor from ``train_loader`` has shape ``[batch_global, ...]``). The runner +asserts this. Callers wiring up the loader should pass ``dist_state=None`` so +the data path replicates the batch across ranks instead of sharding it. + +Each step function then slices to its own per-pool batch shard via the +layout's ``my_batch_slice_*`` helpers. The 3-way batch routing in +``layout.py`` (see "Batch-split routing" in its docstring) assumes this +sliced-from-global pattern. +""" + +import itertools +import os +import time +from contextlib import nullcontext +from typing import Any, Self + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor +from torch.utils.data import DataLoader + +from param_decomp._trace import dump_memory_stats, trace +from param_decomp.batch_and_loss_fns import ReconstructionLoss, RunBatch +from param_decomp.component_model import ComponentModel +from param_decomp.configs import Cadence, PDConfig, RuntimeConfig +from param_decomp.decomposition_targets import ( + DecompositionTarget, + resolve_decomposition_targets, +) +from param_decomp.distributed import seed_all_ranks, seed_per_rank +from param_decomp.masks import AllLayersRouter +from param_decomp.metrics.base import LossMetricConfig +from param_decomp.metrics.importance_minimality import ImportanceMinimalityLossConfig +from param_decomp.metrics.persistent_pgd_recon import ( + PersistentPGDReconLossConfig, + validate_pgd_scope, +) +from param_decomp.metrics.persistent_pgd_state import PersistentPGDState +from param_decomp.optimize import load_optimizer_state_by_name, optimizer_state_by_name +from param_decomp.run_sink import ThreePoolRunSink +from param_decomp.schedule import get_scheduled_value +from param_decomp.sdpa_strict import verify_flash_attention_available +from param_decomp.three_pool.checkpoint import gather_full_state_dict_to_rank0 +from param_decomp.three_pool.config import ThreePoolConfig +from param_decomp.three_pool.layout import ( + LayerwiseBlockGroup, + ThreePoolLayout, + build_world, +) +from param_decomp.three_pool.profiler import PhaseProfiler +from param_decomp.three_pool.reductions import ( + aggregate_losses_to_rank0, + aggregate_max_memory_to_rank0, +) +from param_decomp.three_pool.runtime import _ThreePoolRuntime +from param_decomp.three_pool.step_ci import step_ci +from param_decomp.three_pool.step_layerwise import ( + finalize_layerwise_async_drain, + run_faithfulness_warmup_layerwise, + step_layerwise, +) +from param_decomp.three_pool.step_ppgd import finalize_ppgd_async_drain, step_ppgd +from param_decomp.torch_helpers import loop_dataloader +from param_decomp.training_state import ThreePoolTrainingState +from param_decomp.two_pool.loss_strategy import LayerwiseLossStrategy + +# Loss-metric type discriminators required for the 3-pool training path. +# Same set as two_pool — three-pool reuses the same loss-metric vocabulary. +REQUIRED_LOSS_METRIC_TYPES: frozenset[str] = frozenset( + { + "FaithfulnessLoss", + "ImportanceMinimalityLoss", + "StochasticReconLayerwiseLoss", + "PersistentPGDReconLoss", + } +) + +FORBIDDEN_LOSS_METRIC_TYPES: frozenset[str] = frozenset( + { + "StochasticReconLoss", + "StochasticReconSubsetLoss", + "StochasticReconSubsetCEAndKL", + "PersistentPGDReconSubsetLoss", + "CIMaskedReconLoss", + "CIMaskedReconLayerwiseLoss", + "CIMaskedReconSubsetLoss", + "UnmaskedReconLoss", + "PGDReconLoss", + "PGDReconLayerwiseLoss", + "PGDReconSubsetLoss", + "CIMaskedAttnPatternsReconLoss", + "StochasticAttnPatternsReconLoss", + "StochasticHiddenActsReconLoss", + "CIHiddenActsReconLoss", + } +) + + +class ThreePoolTrainer: + """Stateful 3-pool trainer. + + Construction wires up the runtime bundle, world layout, ComponentModel, + layerwise loss strategy, and the per-pool optimizer (CI / LW have one; + PPGD has none — see module docstring). The PPGD state itself is built on + the first batch of :meth:`run` because its source tensor shapes depend on + the data's sequence dims. + + Resume support is rank-local: :meth:`snapshot` produces a self-contained + :class:`~param_decomp.trainer_snapshot.TrainerSnapshot` whose ``resume`` + half carries this rank's slice, and :meth:`from_snapshot` reconstructs + one from it. The lab's resume loader composes per-rank shards across + ranks. The snapshot's ``consumable`` half is the gathered full state + (rank 0 only; ``None`` elsewhere) — the gather happens inside + :meth:`snapshot`, so all ranks must call it in sync. + """ + + pd_config: PDConfig + runtime_config: RuntimeConfig + three_pool_config: ThreePoolConfig + reconstruction_loss: ReconstructionLoss + component_model: ComponentModel + layout: ThreePoolLayout + strategy: LayerwiseLossStrategy + optimizer: torch.optim.Optimizer | None + ppgd_state: PersistentPGDState | None + step: int + + def __init__( + self, + *, + target_model: nn.Module, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, + pd_config: PDConfig, + runtime_config: RuntimeConfig, + three_pool_config: ThreePoolConfig, + ) -> None: + assert dist.is_initialized(), ( + "init the distributed process group before constructing ThreePoolTrainer" + ) + self.pd_config = pd_config + self.runtime_config = runtime_config + self.three_pool_config = three_pool_config + self.reconstruction_loss = reconstruction_loss + self._run_batch = run_batch + self._target_model = target_model + self.step = 0 + + trace("ThreePoolTrainer.__init__: enter") + # Verify FA can dispatch on the CI fn's largest SDPA shape. If it + # can't (head_dim > 128, missing kernel, etc.), error here rather + # than silently fall back to a 5-10x slower math kernel during + # training. Skipping the global-toggle approach to stay compatible + # with torch.compile's fake-tensor trace. + ci_attn = _ci_attn_shape_or_none(pd_config) + if ci_attn is not None: + d_model, n_heads = ci_attn + verify_flash_attention_available(head_dim=d_model // n_heads) + _validate_pd_config_for_three_pool(pd_config, three_pool_config) + # PPGD runs only on PPGD pool; the relevant per-rank batch is batch // n_ppgd. + validate_pgd_scope( + pd_config.loss_metrics, + batch_size=pd_config.batch_size, + world_size=len(three_pool_config.ppgd_ranks), + ) + + trace("ThreePoolTrainer.__init__: building runtime") + self.runtime = _build_runtime( + target_model=target_model, + pd_config=pd_config, + runtime_config=runtime_config, + three_pool_config=three_pool_config, + run_batch=run_batch, + reconstruction_loss=reconstruction_loss, + ) + + torch.set_float32_matmul_precision("high") + + self._device = torch.device(runtime_config.device) + block_groups = [ + LayerwiseBlockGroup(ranks=tuple(bg.ranks), owned_sites=tuple(bg.owned_sites)) + for bg in three_pool_config.layerwise_block_groups + ] + trace("ThreePoolTrainer.__init__: build_world: enter") + world = build_world( + ci_ranks=list(three_pool_config.ci_ranks), + layerwise_block_groups=block_groups, + ppgd_ranks=list(three_pool_config.ppgd_ranks), + batch_global=self.runtime.batch_global, + device=self._device, + ) + trace("ThreePoolTrainer.__init__: build_world: done") + self.layout = ThreePoolLayout.from_world(world, dist.get_rank()) + decomposition_targets = _decomposition_targets_for_pool( + self.layout, self.runtime.c_per_site + ) + trace( + f"ThreePoolTrainer.__init__: my_pool={self.layout.my_pool} " + f"n_decomp_targets={len(decomposition_targets)}" + ) + + target_model.requires_grad_(False) + # Resync RNG across ranks before V/U + CI fn init — see + # ``two_pool.optimize.TwoPoolTrainer.__init__`` for rationale. + seed_all_ranks(pd_config.seed) + trace("ThreePoolTrainer.__init__: ComponentModel ctor: enter") + self.component_model = ComponentModel( + target_model=target_model, + run_batch=run_batch, + decomposition_targets=decomposition_targets, + ci_config=pd_config.ci_config, + sigmoid_type=pd_config.sigmoid_type, + ) + trace("ThreePoolTrainer.__init__: ComponentModel ctor: done") + # Drop pool-irrelevant params before moving to GPU. RNG draws used to + # init them already happened (in the ctor above), so equivalence with + # single-pool / 2-pool is preserved. + match self.layout.my_pool: + case "layerwise" | "ppgd": + self.component_model.drop_ci_fn() + trace(f"ThreePoolTrainer.__init__: dropped ci_fn ({self.layout.my_pool} pool)") + case "ci": + self.component_model.drop_components() + trace("ThreePoolTrainer.__init__: dropped V/U components (ci pool)") + trace("ThreePoolTrainer.__init__: ComponentModel.to(device): enter") + self.component_model = self.component_model.to(self._device) + trace("ThreePoolTrainer.__init__: ComponentModel.to(device): done") + dump_memory_stats("after ComponentModel.to(device)") + # CI pool: optionally torch.compile the CI fn. Eats compile time on + # step 0 / step 1 (first fwd + first bwd) but should cut ``ci/8_fused_bwd`` + # substantially — that backward through the 2.64B-param CI fn dominates + # the critical path (70% of CI step at batch=48). + if self.layout.my_pool == "ci" and os.environ.get("PD_COMPILE_CI_FN", "").strip() in ( + "1", + "true", + "yes", + ): + assert self.component_model.ci_fn is not None + trace("ThreePoolTrainer.__init__: torch.compile(ci_fn)") + self.component_model.ci_fn = torch.compile(self.component_model.ci_fn) # pyright: ignore[reportAttributeAccessIssue] + # Diverge stochastic RNG per rank for mask sampling. + seed_per_rank(pd_config.seed) + + trace("ThreePoolTrainer.__init__: building LayerwiseLossStrategy") + self.strategy = LayerwiseLossStrategy.from_cfg( + target_model, + use_fused_kl=three_pool_config.use_fused_kl, + unfused_recon=reconstruction_loss, + ) + 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})") + 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) + trace(f"ThreePoolTrainer.__init__: CI fn params={n_params / 1e9:.3f}B") + self.optimizer = torch.optim.AdamW( + [ + { + "params": self._ci_fn_params, + "lr": pd_config.ci_fn_optimizer.lr_schedule.start_val, + } + ], + weight_decay=0.0, + fused=True, + ) + case "layerwise": + 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( + [ + { + "params": self._component_params, + "lr": pd_config.components_optimizer.lr_schedule.start_val, + } + ], + weight_decay=0.0, + fused=True, + ) + case "ppgd": + pass # ppgd_state constructed lazily from first batch in run() + trace("ThreePoolTrainer.__init__: optimizer build: done") + dump_memory_stats("after optimizer build") + trace("ThreePoolTrainer.__init__: exit") + + # ============================ Atomic cfg + state ============================ + + def _named_params_for_my_optimizer(self) -> list[tuple[str, nn.Parameter]]: + """The ``(name, param)`` pairs in the order they were added to this rank's + optimizer. ``CI`` pool returns ``ci_fn.*`` pairs; ``LW`` pool returns + ``components..*`` pairs for this rank's owned sites; ``PPGD`` has + no optimizer and returns ``[]``.""" + match self.layout.my_pool: + case "ci": + assert self.component_model.ci_fn is not None + return [(f"ci_fn.{n}", p) for n, p in self.component_model.ci_fn.named_parameters()] + case "layerwise": + out: list[tuple[str, nn.Parameter]] = [] + for site in self.layout.my_owned_sites: + for n, p in self.component_model.components[site].named_parameters(): + out.append((f"components.{site}.{n}", p)) + return out + case "ppgd": + return [] + + def snapshot(self) -> ThreePoolTrainingState | None: + """Canonical point-in-time view of the 3-pool trainer. + + On rank 0 (the only rank a sink consumes from), assembles a + topology-free :class:`ThreePoolTrainingState`: + + * ``component_model``: the full gathered model state (every site's V/U + plus the CI fn) from :func:`gather_full_state_dict_to_rank0`. + * ``components_optimizer`` / ``ci_fn_optimizer``: per-parameter + optimizer state keyed by param name, merged across all ranks. + ``components_optimizer`` comes from the layerwise pool (each LW + rank contributes its owned sites); ``ci_fn_optimizer`` comes from + the CI pool (all CI ranks have the same state via in-pool + all-reduce, so any suffices). + * ``ppgd_state_by_rank``: PPGD adversarial sources, keyed by PPGD + rank id. Genuinely rank-coupled for ``PerBatchPerPositionScope`` + sources (sized by the rank-local batch slice). + + Returns ``None`` on non-rank-0 — the lab sink is silent on those + ranks anyway, and the trainer never needs their canonical state. + """ + gathered_model = gather_full_state_dict_to_rank0( + layout=self.layout, + component_model=self.component_model, + target_model=self._target_model, + run_batch=self._run_batch, + ci_config=self.pd_config.ci_config, + sigmoid_type=self.pd_config.sigmoid_type, + c_per_site=self.runtime.c_per_site, + device=self._device, + ) + + 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, + } + 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 + + world_size = self.layout.world.world_size + if self.layout.my_rank != 0: + dist.gather_object(my_contribution, None, dst=0) + return None + + gathered: list[dict[str, Any] | None] = [None] * world_size + dist.gather_object(my_contribution, gathered, dst=0) + components_by_name: dict[str, dict[str, Any]] = {} + ci_fn_by_name: dict[str, dict[str, Any]] = {} + ppgd_by_rank: dict[int, dict[str, Any]] = {} + for r, c in enumerate(gathered): + assert c is not None + pool: str = c["pool"] + match pool: + case "layerwise": + components_by_name.update(c["optimizer_by_name"]) + case "ci": + ci_fn_by_name.update(c["optimizer_by_name"]) + case "ppgd": + if "ppgd" in c: + ppgd_by_rank[r] = c["ppgd"] + case _: + raise AssertionError(f"unknown pool {pool!r} in rank-{r} contribution") + assert gathered_model is not None # rank 0 always has the gathered model + return ThreePoolTrainingState( + step=self.step, + pd_config=self.pd_config.model_dump(), + runtime_config=self.runtime_config.model_dump(), + three_pool_config=self.three_pool_config.model_dump(), + layout_fingerprint=_layout_fingerprint(self.layout), + component_model=gathered_model, + components_optimizer=components_by_name, + ci_fn_optimizer=ci_fn_by_name, + ppgd_state_by_rank=ppgd_by_rank, + ) + + @classmethod + def from_snapshot( + cls, + snapshot: ThreePoolTrainingState, + *, + target_model: nn.Module, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, + cfg_overrides: dict[str, Any] | None = None, + ) -> Self: + """Reconstruct a 3-pool trainer from a canonical snapshot. + + Each rank must arrive with the canonical state populated (typically by + reading ``training_.pth`` on every rank from a shared filesystem). + Each rank extracts its own slice of the canonical state by rank id and + loads it into its locally-owned sub-state. + """ + pd_dict = snapshot.pd_config + if cfg_overrides is not None: + pd_dict = {**pd_dict, **cfg_overrides} + pd_config = PDConfig.model_validate(pd_dict) + runtime_config = RuntimeConfig.model_validate(snapshot.runtime_config) + three_pool_config = ThreePoolConfig.model_validate(snapshot.three_pool_config) + + trainer = cls( + target_model=target_model, + run_batch=run_batch, + reconstruction_loss=reconstruction_loss, + pd_config=pd_config, + runtime_config=runtime_config, + three_pool_config=three_pool_config, + ) + saved_fp = snapshot.layout_fingerprint + current_fp = _layout_fingerprint(trainer.layout) + assert saved_fp == current_fp, ( + f"3-pool layout fingerprint mismatch on resume:\n" + f" saved: {saved_fp}\n" + f" current: {current_fp}\n" + ) + trainer._load_canonical_state(snapshot) + return trainer + + def _load_canonical_state(self, state: ThreePoolTrainingState) -> None: + """Each rank extracts the slice of the canonical state it owns.""" + self.step = state.step + # The full gathered model state has every site's V/U + the CI fn. Each + # rank's locally-constructed ComponentModel only has the keys it owns, + # so we filter to the local subset before loading. `strict=False` lets + # us drop the canonical entries this rank doesn't hold. + local_model_keys = set(self.component_model.state_dict().keys()) + 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) + + # ============================ Training loop ============================ + + def run( + self, + train_loader: DataLoader[Any], + sink: ThreePoolRunSink, + cadence: Cadence, + profiler: PhaseProfiler | None = None, + ) -> None: + """Advance training from ``self.step`` to ``self.pd_config.steps``.""" + trace("Trainer.run: enter") + pd_config = self.pd_config + layout = self.layout + runtime = self.runtime + n_steps = pd_config.steps + defer_vu_opt = self.three_pool_config.defer_vu_opt + device = self._device + + train_iterator = loop_dataloader(train_loader) + # Loader skip-replay (resumed mid-trajectory). + for _ in range(self.step): + next(train_iterator) + + # Peek first batch (after any skip) for PPGD source-shape sizing. + trace("Trainer.run: first_batch peek: enter") + first_batch = next(train_iterator) + trace("Trainer.run: first_batch peek: done") + 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: + trace("Trainer.run: PPGDState ctor: enter") + ppgd_cfg = runtime.ppgd_cfg + self.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, + use_delta_component=True, + optimizer_cfg=ppgd_cfg.optimizer, + scope=ppgd_cfg.scope, + use_sigmoid_parameterization=ppgd_cfg.use_sigmoid_parameterization, + n_warmup_steps=ppgd_cfg.n_warmup_steps, + n_samples=ppgd_cfg.n_samples, + 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 + trace("Trainer.run: PPGDState ctor: done") + + if ( + self.step == 0 + and layout.my_pool == "layerwise" + and pd_config.faithfulness_warmup_steps > 0 + ): + trace( + f"Trainer.run: faithfulness warmup: enter ({pd_config.faithfulness_warmup_steps} steps)" + ) + run_faithfulness_warmup_layerwise( + component_model=self.component_model, + component_params=self._component_params, + n_steps=pd_config.faithfulness_warmup_steps, + lr=pd_config.faithfulness_warmup_lr, + weight_decay=pd_config.faithfulness_warmup_weight_decay, + numel_global=self.runtime.numel_global, + ) + trace("Trainer.run: faithfulness warmup: done") + + components_lr_schedule = pd_config.components_optimizer.lr_schedule + ci_fn_lr_schedule = pd_config.ci_fn_optimizer.lr_schedule + + profiler_ctx = profiler if profiler is not None else nullcontext() + h_cache_ci: dict[str, Tensor] | None = None + # Async-pipeline state threaded across iterations on LW + PPGD pools. + pending_all_reduce_lw: list[tuple[list[Tensor], Tensor, dist.Work]] | None = None + pending_recv_vu_ppgd: list[tuple[Any, Tensor, dist.Work]] | None = None + + def _to_device(b: Any) -> Any: + """Move a batch yielded by the train loader to this rank's GPU. + + 3-pool's step functions assume the batch is already on-device + (mirroring 2-pool's `_extract_batch_tensor`). The loader produces + CPU tensors; moving here keeps the step functions thin. + """ + if b is None: + return None + if isinstance(b, Tensor): + return b.to(device) + if isinstance(b, dict) and "input_ids" in b: + return {**b, "input_ids": b["input_ids"].to(device)} + if isinstance(b, list | tuple) and len(b) > 0 and isinstance(b[0], Tensor): + return type(b)([b[0].to(device), *b[1:]]) + raise TypeError(f"Unsupported batch type from DataLoader: {type(b).__name__}") + + with profiler_ctx: + # 2-batch peek window: batch_T is the step's batch, batch_T_plus_1 is the + # next-step batch peeked early for CI pool's dead-time prefetch. + trace("Trainer.run: pre-loop batch peek: enter") + batch_T = _to_device(next(train_iterator)) + batch_T_plus_1 = _to_device(next(train_iterator, None)) + trace("Trainer.run: pre-loop batch peek: done, entering training loop") + + for step in range(self.step, n_steps): + self.step = step + _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( + step, n_steps, ci_fn_lr_schedule + ) + elif layout.my_pool == "layerwise": + lr_step = max(step - 1, 0) if defer_vu_opt else step + self.optimizer.param_groups[0]["lr"] = get_scheduled_value( + lr_step, n_steps, components_lr_schedule + ) + + if profiler is not None: + dist.barrier() + + torch.cuda.synchronize(device) + step_start = time.perf_counter() + + # batch_T should already be on this rank's device (placed by _to_device). + if isinstance(batch_T, Tensor): + 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, ( + 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.component_model, + self.optimizer, + self._ci_fn_params, + batch_T=batch_T, + batch_T_plus_1=next_batch_for_prefetch, + h_cache_T=h_cache_ci, + cfg=runtime, + current_frac_of_training=step / n_steps if n_steps > 0 else 0.0, + profiler=profiler, + ) + case "layerwise": + assert self.optimizer is not None, ( + f"LW rank {layout.my_rank} missing optimizer" + ) + assert layout.my_owned_sites, ( + f"LW rank {layout.my_rank} has no owned_sites — empty block" + ) + metrics, pending_all_reduce_lw = step_layerwise( + layout, + self.component_model, + self.optimizer, + self._all_params, + batch_T, + runtime, + self.strategy, + defer_vu_opt=defer_vu_opt, + prev_pending_all_reduce=pending_all_reduce_lw, + profiler=profiler, + ) + case "ppgd": + assert self.ppgd_state is not None, ( + f"PPGD rank {layout.my_rank} has no ppgd_state — lazy init failed" + ) + metrics, pending_recv_vu_ppgd = step_ppgd( + layout, + self.component_model, + self.ppgd_state, + batch_T, + runtime, + self.strategy, + step=step, + n_steps=n_steps, + defer_vu_opt=defer_vu_opt, + prev_pending_recv_vu=pending_recv_vu_ppgd, + profiler=profiler, + ) + # Catch silent NaN propagation early. Covers both the per-rank + # display scalars (``loss/*``) and the raw aggregation + # ingredients (``_raw/*``) the logger sums across the pool. + for k, v in metrics.items(): + if k.startswith("loss/") or k.startswith("_raw/"): + assert v == v, f"NaN in metrics[{k!r}] at step {step}" # NaN != NaN + + # current_stream().synchronize() — not torch.cuda.synchronize() — so we + # don't wait for *all* CUDA streams. In ``defer_vu_opt=True`` mode, PPGD + # has a pending async dist.broadcast (irecv side) on a NCCL stream that + # only completes when LW step N+1 phase B4 fires its matching broadcast. + # Waiting for that here would deadlock against ``_log_train_metrics`` + # (LW rank 0 blocks on ``dist.recv`` from PPGD leader → PPGD leader + # blocks here on the irecv → LW can't reach phase B4). The default + # stream carries the step's compute, which is all we need for an + # accurate ``step_ms`` measurement. + torch.cuda.current_stream(device).synchronize() + step_ms = (time.perf_counter() - step_start) * 1000.0 + trace(f"Trainer.run: step {step}: done in {step_ms:.1f}ms") + if step % cadence.train_log_every == 0: + dump_memory_stats(f"step {step} done") + + if step % cadence.train_log_every == 0: + _log_train_metrics( + metrics=metrics, + layout=layout, + device=device, + step=step, + step_ms=step_ms, + runtime=runtime, + optimizer=self.optimizer, + sink=sink, + ) + + if cadence.should_save(step): + snap = self.snapshot() + if snap is not None: + sink.checkpoint(snap) + + batch_T = ( + batch_T_plus_1 + if batch_T_plus_1 is not None + else _to_device(next(train_iterator)) + ) + batch_T_plus_1 = _to_device(next(train_iterator, None)) + + if profiler is not None: + profiler.step() + + # Drain the final iter's deferred opt (async mode only). Without this, + # the saved checkpoint would be missing the last iter's update. + if defer_vu_opt: + match layout.my_pool: + case "layerwise": + assert self.optimizer is not None + if pending_all_reduce_lw is not None: + self.optimizer.param_groups[0]["lr"] = get_scheduled_value( + n_steps - 1, n_steps, components_lr_schedule + ) + finalize_layerwise_async_drain( + layout, + self.component_model, + self.optimizer, + self._all_params, + pending_all_reduce_lw, + runtime.grad_clip_norm_components, + ) + pending_all_reduce_lw = None + case "ppgd": + if pending_recv_vu_ppgd is not None: + finalize_ppgd_async_drain( + layout, + self.component_model, + pending_recv_vu_ppgd, # type: ignore[arg-type] + ) + pending_recv_vu_ppgd = None + case "ci": + pass # CI pool doesn't defer + + self.step = n_steps + snap = self.snapshot() + if snap is not None: + sink.checkpoint(snap) + + +def optimize_three_pool( + target_model: nn.Module, + train_loader: DataLoader[Any], + *, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, + pd_config: PDConfig, + runtime_config: RuntimeConfig, + three_pool_config: ThreePoolConfig, + cadence: Cadence, + sink: ThreePoolRunSink, + profiler: PhaseProfiler | None = None, +) -> None: + """Train a ComponentModel under the 3-pool strategy. + + Thin wrapper over :class:`ThreePoolTrainer` for callers that don't need + resumption or interactive control. + """ + trainer = ThreePoolTrainer( + target_model=target_model, + run_batch=run_batch, + reconstruction_loss=reconstruction_loss, + pd_config=pd_config, + runtime_config=runtime_config, + three_pool_config=three_pool_config, + ) + trainer.run(train_loader, sink, cadence, profiler=profiler) + + +def _layout_fingerprint(layout: ThreePoolLayout) -> dict[str, Any]: + """Compact summary of the 3-pool world layout. Compared at resume time.""" + return { + "world_size": layout.world.world_size, + "ci_ranks": list(layout.world.ci_ranks), + "ppgd_ranks": list(layout.world.ppgd_ranks), + "n_layerwise_blocks": len(layout.world.layerwise_block_groups), + "my_rank": layout.my_rank, + "my_pool": layout.my_pool, + "owned_sites": list(layout.my_owned_sites) if layout.my_pool == "layerwise" else [], + } + + +def _validate_pd_config_for_three_pool( + pd_config: PDConfig, + three_pool_config: ThreePoolConfig, +) -> None: + """Fail loudly on any PDConfig the 3-pool path can't honour.""" + by_type: dict[str, LossMetricConfig] = {m.type: m for m in pd_config.loss_metrics} + + missing = sorted(REQUIRED_LOSS_METRIC_TYPES - set(by_type)) + assert not missing, ( + f"3-pool requires these loss metrics: {sorted(REQUIRED_LOSS_METRIC_TYPES)}.\n" + f"Missing: {missing}. Got: {sorted(by_type)}." + ) + illegal = sorted(FORBIDDEN_LOSS_METRIC_TYPES & set(by_type)) + assert not illegal, ( + f"3-pool does not implement these loss metrics (they would be silently ignored): " + f"{illegal}. Remove them or extend the 3-pool path." + ) + + for name in REQUIRED_LOSS_METRIC_TYPES: + assert by_type[name].coeff is not None, ( + f"pd_config.loss_metrics[{name!r}].coeff is required for 3-pool training" + ) + + n_per_block = len(three_pool_config.layerwise_block_groups[0].ranks) + n_ci = len(three_pool_config.ci_ranks) + n_ppgd = len(three_pool_config.ppgd_ranks) + bs = pd_config.batch_size + assert bs % n_per_block == 0, ( + f"pd_config.batch_size ({bs}) must be divisible by N_per_block ({n_per_block}) " + f"= len(layerwise_block_groups[0].ranks)" + ) + assert bs % n_ci == 0, ( + f"pd_config.batch_size ({bs}) must be divisible by N_ci ({n_ci}) = len(ci_ranks)" + ) + assert bs % n_ppgd == 0, ( + f"pd_config.batch_size ({bs}) must be divisible by N_ppgd ({n_ppgd}) = len(ppgd_ranks)" + ) + + assert pd_config.use_delta_component, ( + "3-pool requires pd_config.use_delta_component=True (hardcoded in LW's " + "layerwise stoch recon + PPGD's PPGD warmup)." + ) + + assert pd_config.sampling == "continuous", ( + "3-pool hardcodes `sampling='continuous'` in CI pool's CI computation; " + f"got pd_config.sampling={pd_config.sampling!r}." + ) + assert pd_config.n_mask_samples == 1, ( + "3-pool draws exactly one stochastic mask per site per step in LW; " + f"got pd_config.n_mask_samples={pd_config.n_mask_samples}." + ) + + assert pd_config.identity_decomposition_targets is None, ( + "3-pool path does not call `insert_identity_operations_`; " + "`identity_decomposition_targets` would be silently ignored." + ) + + # Convention: rank 0 must be the Layerwise pool's block 0 leader. + assert three_pool_config.layerwise_block_groups[0].ranks[0] == 0, ( + "Convention: rank 0 must be the LW pool's block 0 leader (so reductions " + "can ship CI/PPGD pool losses to rank 0). Reorder layerwise_block_groups " + "so the first group starts with rank 0." + ) + + ppgd_cfg = by_type["PersistentPGDReconLoss"] + assert isinstance(ppgd_cfg, PersistentPGDReconLossConfig) + assert ppgd_cfg.start_frac == 0.0, ( + "3-pool path does not implement PersistentPGDReconLoss.start_frac > 0; " + "PPGD always runs from step 0." + ) + + +def _build_runtime( + target_model: nn.Module, + pd_config: PDConfig, + runtime_config: RuntimeConfig, + three_pool_config: ThreePoolConfig, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, +) -> _ThreePoolRuntime: + """Assemble the step-context bundle from configs + target.""" + targets = resolve_decomposition_targets(target_model, pd_config.decomposition_targets) + c_per_site = {t.module_path: t.C for t in targets} + numel_global = 0 + for t in targets: + w = target_model.get_submodule(t.module_path).weight + assert isinstance(w, Tensor) + numel_global += w.numel() + + for bg in three_pool_config.layerwise_block_groups: + for site in bg.owned_sites: + assert site in c_per_site, ( + f"site '{site}' in layerwise block group but not in " + f"pd_config.decomposition_targets after pattern expansion. " + f"Available: {sorted(c_per_site)[:5]}..." + ) + + by_type: dict[str, LossMetricConfig] = {m.type: m for m in pd_config.loss_metrics} + ppgd_cfg = by_type["PersistentPGDReconLoss"] + imp_min_cfg = by_type["ImportanceMinimalityLoss"] + assert isinstance(ppgd_cfg, PersistentPGDReconLossConfig) + assert isinstance(imp_min_cfg, ImportanceMinimalityLossConfig) + + def _coeff(name: str) -> float: + c = by_type[name].coeff + assert c is not None + return float(c) + + block_groups = tuple( + LayerwiseBlockGroup(ranks=tuple(bg.ranks), owned_sites=tuple(bg.owned_sites)) + for bg in three_pool_config.layerwise_block_groups + ) + + return _ThreePoolRuntime( + ci_ranks=tuple(three_pool_config.ci_ranks), + layerwise_block_groups=block_groups, + ppgd_ranks=tuple(three_pool_config.ppgd_ranks), + batch_global=pd_config.batch_size, + c_per_site=c_per_site, + ci_config=pd_config.ci_config, + sigmoid_type=pd_config.sigmoid_type, + run_batch=run_batch, + reconstruction_loss=reconstruction_loss, + ppgd_cfg=ppgd_cfg, + coeff_faith=_coeff("FaithfulnessLoss"), + coeff_imp=_coeff("ImportanceMinimalityLoss"), + coeff_stoch=_coeff("StochasticReconLayerwiseLoss"), + coeff_ppgd=_coeff("PersistentPGDReconLoss"), + imp_min_pnorm=imp_min_cfg.pnorm, + imp_min_beta=imp_min_cfg.beta, + imp_min_eps=imp_min_cfg.eps, + imp_min_p_anneal_start_frac=imp_min_cfg.p_anneal_start_frac, + imp_min_p_anneal_final_p=imp_min_cfg.p_anneal_final_p, + imp_min_p_anneal_end_frac=imp_min_cfg.p_anneal_end_frac, + lr_components=pd_config.components_optimizer.lr_schedule.start_val, + lr_ci_fn=pd_config.ci_fn_optimizer.lr_schedule.start_val, + grad_clip_norm_components=pd_config.components_optimizer.grad_clip_norm, + grad_clip_norm_ci_fn=pd_config.ci_fn_optimizer.grad_clip_norm, + numel_global=numel_global, + bf16_autocast=runtime_config.autocast_bf16, + use_fused_kl=three_pool_config.use_fused_kl, + ) + + +def _ci_attn_shape_or_none(pd_config: PDConfig) -> tuple[int, int] | None: + """Pull ``(d_model, n_heads)`` from the CI fn's attn config if the CI fn + has one (transformer variants), else ``None``. Used to derive the SDPA + shape for FA startup verification. + """ + ci_cfg = pd_config.ci_config + # Walk the nested config tree without a hard dependency on the CI types + # (mode/fn_type discriminators bury the attn config differently per variant). + for attr in ("simple_transformer_ci_cfg", "transformer_cfg"): + nested = getattr(ci_cfg, attr, None) + if nested is None: + continue + attn = getattr(nested, "attn_config", None) + d_model = getattr(nested, "d_model", None) + if attn is not None and d_model is not None: + return d_model, attn.n_heads + return None + + +def _decomposition_targets_for_pool( + layout: ThreePoolLayout, c_per_site: dict[str, int] +) -> list[DecompositionTarget]: + """CI/PPGD pools: every site (full CI fn / full V/U replica). + Layerwise pool: only this rank's owned sites.""" + match layout.my_pool: + case "ci" | "ppgd": + sites = layout.world.all_sites + case "layerwise": + sites = layout.my_owned_sites + return [DecompositionTarget(module_path=s, C=c_per_site[s]) for s in sites] + + +def _assert_full_global_batch(batch: Any, batch_global: int) -> None: + """The 3-pool data contract: every rank reads the FULL global batch on + every step so each pool can slice to its own DP shard. + """ + if isinstance(batch, Tensor): + actual = batch.shape[0] + elif isinstance(batch, dict) and "input_ids" in batch: + actual = batch["input_ids"].shape[0] + elif isinstance(batch, list | tuple) and len(batch) > 0 and isinstance(batch[0], Tensor): + actual = batch[0].shape[0] + else: + raise TypeError(f"Unsupported batch type from DataLoader: {type(batch).__name__}") + assert actual == batch_global, ( + f"3-pool requires each rank to read the FULL global batch; got batch with " + f"leading-dim {actual}, expected {batch_global}. Most likely cause: the data " + f"loader was built with a non-None dist_state (which shards the batch). " + f"For 3-pool, pass dist_state=None to build_loader." + ) + + +def _seq_dims_from_batch(batch: Any) -> tuple[int, ...]: + """Sequence dims (everything past the leading batch dim) of a sample batch.""" + if isinstance(batch, Tensor): + return tuple(batch.shape[1:]) + if isinstance(batch, dict) and "input_ids" in batch: + return tuple(batch["input_ids"].shape[1:]) + if isinstance(batch, list | tuple) and len(batch) > 0 and isinstance(batch[0], Tensor): + return tuple(batch[0].shape[1:]) + raise TypeError(f"Cannot infer seq dims from batch of type {type(batch).__name__}") + + +def _log_train_metrics( + *, + metrics: dict[str, float], + layout: ThreePoolLayout, + device: torch.device, + step: int, + step_ms: float, + runtime: _ThreePoolRuntime, + optimizer: torch.optim.Optimizer | None, + sink: ThreePoolRunSink, +) -> None: + """Aggregate per-pool metrics to rank 0 and dispatch to ``sink``.""" + combined = aggregate_losses_to_rank0(metrics, layout, device) + mem_combined = aggregate_max_memory_to_rank0(layout, device) + + # Reduce step_ms with MAX across LW pool (slowest LW rank is the wall-clock floor). + step_ms_t = torch.tensor([step_ms], device=device) + if layout.my_pool == "layerwise": + dist.all_reduce(step_ms_t, op=dist.ReduceOp.MAX, group=layout.world.layerwise_pool_group) + + if layout.my_rank != 0 or combined is None: + return + + if mem_combined is not None: + combined.update(mem_combined) + combined["perf/step_ms"] = step_ms_t.item() + combined["loss/total"] = ( + runtime.coeff_faith * combined["loss/faith"] + + runtime.coeff_imp * combined["loss/imp"] + + runtime.coeff_stoch * combined["loss/stoch"] + + runtime.coeff_ppgd * combined["loss/ppgd"] + ) + assert layout.my_pool == "layerwise", "rank 0 must be in LW pool (per validator)" + assert optimizer is not None + combined["schedules/lr/components"] = optimizer.param_groups[0]["lr"] + sink.console( + f"--- Step {step} ---", + *(f"train/{name}: {value:.6g}" for name, value in combined.items()), + ) + sink.log({f"train/{k}": v for k, v in combined.items()}, step=step) diff --git a/param_decomp/three_pool/profiler.py b/param_decomp/three_pool/profiler.py new file mode 100644 index 000000000..d06dcd5cb --- /dev/null +++ b/param_decomp/three_pool/profiler.py @@ -0,0 +1,120 @@ +"""``PhaseProfiler`` for 3-pool training. + +Same machinery as ``two_pool.profiler.PhaseProfiler`` (CPU + CUDA events into +a Chrome trace JSON readable by Perfetto), parameterized on the 3-pool literal +``"ci" | "layerwise" | "ppgd"`` instead of ``"a" | "b"``. Duplicated rather +than imported to keep each subsystem owning its own surface. +""" + +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import torch +import torch.profiler + +from param_decomp._trace import phase_trace_enabled, trace + + +@dataclass +class PhaseProfiler: + """Thin wrapper around ``torch.profiler.profile`` for 3-pool training. + + Records CPU + CUDA events into a Chrome trace JSON that loads directly into + Perfetto. Captures kernel timings, NCCL ops, and GPU memory allocations. + + Use as a context manager around the training loop, with ``phase(name)`` + annotations marking logical step phases ("ci/1_ci_fn_fwd", "lw/3_layerwise", + "pgd/4_ppgd_warmup", etc.) inside. Call ``step()`` once per training + iteration so the schedule progresses. + + Trace schedule by default skips the first 20 iters (CUDA / JIT warmup) then + records 3 active steps and dumps the trace. + """ + + enabled: bool = False + out_dir: Path | None = None + rank: int = 0 + pool: Literal["ci", "layerwise", "ppgd"] = "ci" + skip_first: int = 20 + active: int = 3 + _prof: torch.profiler.profile | None = None + + def __enter__(self) -> "PhaseProfiler": + if not self.enabled: + return self + assert self.out_dir is not None, "PhaseProfiler enabled but no out_dir given" + self.out_dir.mkdir(parents=True, exist_ok=True) + trace_path = self.out_dir / f"trace_{self.pool}_rank{self.rank}.json" + + def on_trace_ready(prof: "torch.profiler.profile") -> None: + prof.export_chrome_trace(str(trace_path)) + print(f"[profiler rank{self.rank}] wrote {trace_path}", flush=True) + + self._prof = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule( + skip_first=self.skip_first, wait=0, warmup=1, active=self.active, repeat=1 + ), + on_trace_ready=on_trace_ready, + record_shapes=False, + profile_memory=True, + with_stack=False, + ) + self._prof.__enter__() + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + if self._prof is not None: + self._prof.__exit__(exc_type, exc, tb) + self._prof = None + + @contextmanager + def phase(self, name: str) -> "Iterator[None]": + """Annotate a logical step phase. + + When ``PD_PHASE_TRACE=1``, emits entry + exit traces tagged with + memory deltas. Entry: ``cur=X.XXgb``. Exit: ``peak=X.XXgb end=X.XXgb + delta=±X.XXgb``. ``peak`` is the within-phase peak (we + ``reset_peak_memory_stats`` at entry); ``delta`` is end-current minus + entry-current, which captures whether the phase persistently allocates. + + Note: resetting peak per phase makes any concurrent reader of + ``max_memory_allocated`` see only this phase's peak, not the + whole-step peak. That's fine in debug mode (where this is enabled); + in production ``PD_PHASE_TRACE`` is off and this is a no-op. + + The ``record_function`` wrapper is only active when the torch + profiler is enabled. + """ + do_trace = phase_trace_enabled() + device = None + before_gb = 0.0 + if do_trace: + device = torch.cuda.current_device() + torch.cuda.reset_peak_memory_stats(device) + before_gb = torch.cuda.memory_allocated(device) / 1e9 + trace(f"phase: {name} cur={before_gb:.2f}gb") + + inner = torch.profiler.record_function(name) if self.enabled else nullcontext() + try: + with inner: + yield + finally: + if do_trace: + peak_gb = torch.cuda.max_memory_allocated(device) / 1e9 + end_gb = torch.cuda.memory_allocated(device) / 1e9 + trace( + f"phase: {name} end peak={peak_gb:.2f}gb " + f"end={end_gb:.2f}gb delta={end_gb - before_gb:+.2f}gb" + ) + + def step(self) -> None: + """Advance the profiler schedule (call once per training iteration).""" + if self._prof is not None: + self._prof.step() diff --git a/param_decomp/three_pool/reductions.py b/param_decomp/three_pool/reductions.py new file mode 100644 index 000000000..1118ef7af --- /dev/null +++ b/param_decomp/three_pool/reductions.py @@ -0,0 +1,126 @@ +"""Cross-pool reductions for logging. + +Same approach as ``two_pool.reductions``: each step function emits per-rank +display scalars (``loss/*``) plus raw additive ingredients (``_raw/*``) that +the logger SUM-reduces across each pool, then finalizes into global scalars +on rank 0. + +Global scalars: + ``faith_global = SUM(faith_num) / SUM(faith_den)`` (LW pool) + ``stoch_global = SUM(stoch_num) / SUM(stoch_den)`` (LW pool) + ``imp_global = SUM(imp_num)`` (CI pool — see note) + ``ppgd_global = SUM(ppgd_num) / SUM(ppgd_den)`` (PPGD pool) + +Note on imp: the CI pool already all-reduces ``per_component_sums`` + +``n_examples`` SUM-wise across the CI pool inside its loss computation, so +every CI rank's ``loss_imp`` scalar IS already the global value. The step +function divides by ``n_ci`` before exposing as ``_raw/imp_num`` so that the +pool-wide all-reduce SUM gives back the global value exactly once. + +LW pool's all-reduce SUM uses a ``1 / n_per_block`` pre-scale (see the +two-pool reductions module docstring for the math). + +Memory uses MAX (the bottleneck rank is what matters). +""" + +import torch +import torch.distributed as dist + +from param_decomp.three_pool.layout import ThreePoolLayout + +LW_RAW_KEYS: tuple[str, ...] = ( + "_raw/faith_num", + "_raw/faith_den", + "_raw/stoch_num", + "_raw/stoch_den", +) +CI_RAW_KEYS: tuple[str, ...] = ("_raw/imp_num",) +PPGD_RAW_KEYS: tuple[str, ...] = ("_raw/ppgd_num", "_raw/ppgd_den") + + +def aggregate_max_memory_to_rank0( + layout: ThreePoolLayout, + device: torch.device, +) -> dict[str, float] | None: + """MAX-reduce CUDA peak memory within each pool; non-rank-0 pool leaders + send their value to rank 0. + + Returns ``{mem/_peak_gb for pool in (lw, ci, ppgd)}`` on rank 0, + ``None`` everywhere else. + """ + peak_gb = torch.cuda.max_memory_allocated() / 1e9 + val = torch.tensor([peak_gb], device=device) + match layout.my_pool: + case "layerwise": + dist.all_reduce(val, op=dist.ReduceOp.MAX, group=layout.world.layerwise_pool_group) + if layout.my_rank == 0: + lw_peak = val.item() + ci_val = torch.empty(1, device=device) + pgd_val = torch.empty(1, device=device) + dist.recv(ci_val, src=layout.world.ci_ranks[0]) + dist.recv(pgd_val, src=layout.world.ppgd_ranks[0]) + return { + "mem/lw_peak_gb": lw_peak, + "mem/ci_peak_gb": ci_val.item(), + "mem/ppgd_peak_gb": pgd_val.item(), + } + return None + case "ci": + dist.all_reduce(val, op=dist.ReduceOp.MAX, group=layout.world.ci_pool_group) + if layout.my_is_pool_leader: + dist.send(val, dst=0) + return None + case "ppgd": + dist.all_reduce(val, op=dist.ReduceOp.MAX, group=layout.world.ppgd_pool_group) + if layout.my_is_pool_leader: + dist.send(val, dst=0) + return None + + +def aggregate_losses_to_rank0( + loss_dict: dict[str, float], + layout: ThreePoolLayout, + device: torch.device, +) -> dict[str, float] | None: + """SUM raw (num, den) ingredients within each pool, ship CI's and PPGD's + sums to rank 0, and finalize the global ``loss/*`` scalars there. + """ + match layout.my_pool: + case "layerwise": + keys = list(LW_RAW_KEYS) + n_per_block = layout.world.n_per_block + vals = torch.tensor( + [loss_dict[k] / n_per_block for k in keys], device=device, dtype=torch.float64 + ) + dist.all_reduce(vals, op=dist.ReduceOp.SUM, group=layout.world.layerwise_pool_group) + if layout.my_rank == 0: + lw = {k: vals[i].item() for i, k in enumerate(keys)} + ci_keys = list(CI_RAW_KEYS) + ci_vals = torch.empty(len(ci_keys), device=device, dtype=torch.float64) + dist.recv(ci_vals, src=layout.world.ci_ranks[0]) + ci = {k: ci_vals[i].item() for i, k in enumerate(ci_keys)} + pgd_keys = list(PPGD_RAW_KEYS) + pgd_vals = torch.empty(len(pgd_keys), device=device, dtype=torch.float64) + dist.recv(pgd_vals, src=layout.world.ppgd_ranks[0]) + pgd = {k: pgd_vals[i].item() for i, k in enumerate(pgd_keys)} + return { + "loss/faith": lw["_raw/faith_num"] / lw["_raw/faith_den"], + "loss/stoch": lw["_raw/stoch_num"] / lw["_raw/stoch_den"], + "loss/imp": ci["_raw/imp_num"], + "loss/ppgd": pgd["_raw/ppgd_num"] / pgd["_raw/ppgd_den"], + } + return None + case "ci": + keys = list(CI_RAW_KEYS) + vals = torch.tensor([loss_dict[k] for k in keys], device=device, dtype=torch.float64) + dist.all_reduce(vals, op=dist.ReduceOp.SUM, group=layout.world.ci_pool_group) + if layout.my_is_pool_leader: + dist.send(vals, dst=0) + return None + case "ppgd": + keys = list(PPGD_RAW_KEYS) + vals = torch.tensor([loss_dict[k] for k in keys], device=device, dtype=torch.float64) + dist.all_reduce(vals, op=dist.ReduceOp.SUM, group=layout.world.ppgd_pool_group) + if layout.my_is_pool_leader: + dist.send(vals, dst=0) + return None diff --git a/param_decomp/three_pool/runtime.py b/param_decomp/three_pool/runtime.py new file mode 100644 index 000000000..ff1f93f87 --- /dev/null +++ b/param_decomp/three_pool/runtime.py @@ -0,0 +1,80 @@ +"""Internal step-context bundle for 3-pool training. + +``_ThreePoolRuntime`` glues serializable config (``ThreePoolConfig``) with the +caller-supplied runtime callables and the derived per-site C from the actual +target. It's the parameter shape that ``step_ci`` / ``step_layerwise`` / +``step_ppgd`` consume; ``optimize_three_pool`` builds one internally per call. + +Not part of the public API. +""" + +from dataclasses import dataclass + +from param_decomp.batch_and_loss_fns import ReconstructionLoss, RunBatch +from param_decomp.ci_fns import CiConfig +from param_decomp.ci_sigmoids import SigmoidType +from param_decomp.metrics.persistent_pgd_recon import PersistentPGDReconLossConfig +from param_decomp.three_pool.layout import LayerwiseBlockGroup + +__all__ = ["_ThreePoolRuntime"] + + +@dataclass(frozen=True) +class _ThreePoolRuntime: + """Internal bundle passed to the three step functions. + + Mirrors ``two_pool.runtime._TwoPoolRuntime`` but typed against the + 3-pool topology + the multi-rank CI pool. Coefficients and per-metric + config are flattened in here so step functions don't have to re-parse + ``pd_config.loss_metrics``. + """ + + # Topology + ci_ranks: tuple[int, ...] + layerwise_block_groups: tuple[LayerwiseBlockGroup, ...] + ppgd_ranks: tuple[int, ...] + batch_global: int + + # Per-site C (resolved from decomposition_targets) + c_per_site: dict[str, int] + + # CI fn shape + ci_config: CiConfig + sigmoid_type: SigmoidType + + # Caller-supplied callables + run_batch: RunBatch + reconstruction_loss: ReconstructionLoss + + # PPGD config (the full LossMetricConfig, since PersistentPGDState + # consumes several of its fields) + ppgd_cfg: PersistentPGDReconLossConfig + + # Loss coefficients + coeff_faith: float + coeff_imp: float + coeff_stoch: float + coeff_ppgd: float + + # Importance minimality knobs + imp_min_pnorm: float + imp_min_beta: float + imp_min_eps: float + imp_min_p_anneal_start_frac: float + imp_min_p_anneal_final_p: float | None + imp_min_p_anneal_end_frac: float + + # Optimizer LRs (start values; schedules consumed by optimize_three_pool + # which mutates optimizer.param_groups[i]['lr'] each step) + lr_components: float + lr_ci_fn: float + grad_clip_norm_components: float | None + grad_clip_norm_ci_fn: float | None + # Total numel across all decomposition sites' weight tensors. Used by the + # LW pool's ``_faithfulness_loss`` so multi-pool's per-element grad matches + # single-pool's (which divides faith by global numel, not rank-local). + numel_global: int + + # Substrate + bf16_autocast: bool + use_fused_kl: bool diff --git a/param_decomp/three_pool/step_ci.py b/param_decomp/three_pool/step_ci.py new file mode 100644 index 000000000..ec966faff --- /dev/null +++ b/param_decomp/three_pool/step_ci.py @@ -0,0 +1,317 @@ +"""CI pool training step. + +CI pool is new under 3-pool. Each CI rank holds the full CI fn (replicated) +and processes one DP shard of the batch. + +Phases (numbered to match ``DESIGN.md`` ``ci/N``): + + 0. Step 0 only: target_fwd to build H_T (subsequent steps reuse the prev + iter's prefetch). + 1. CI fn fwd on H_T → CI_T per site. Graph retained for the fused bwd at 8. + 2. Async send CI_T → LW (per-site, sub-sliced) + PPGD (full-model, + sub-sliced). Kicks the NIC so 4 + 5 + 6 can pipeline. + 3. imp_min loss on ``CI.upper_leaky`` (still on graph — gradient enters + the fused bwd at 8). + 4. Dead-time prefetch: target_fwd(batch T+1) → H_{T+1}. GPU work overlaps + with the NIC sends in 2 and the downstream pool work it triggered. + 5. Recv g_CI from LW (per-site, stitched across K_lw_per_ci slices). + 6. Recv g_CI from PPGD (per-site, stitched across K_ppgd_per_ci slices). + 7. Assemble g_CI_total per site on this CI rank's batch slice. + 8. Fused autograd backward: imp_min via ``upper_leaky``, downstream grads + injected on ``lower_leaky``. One pass through the CI fn graph. + 9. In-pool AVG-reduce on CI fn grads (standard DDP). + 9b. Cross-pool grad clip on CI fn (n_replicas=n_ci dedup). + 10. AdamW step. + 11. Wait the step-2 async sends to flush before storage gets reused next iter. + +Returns ``(metrics, h_cache_T_plus_1)``. The runner threads the prefetched +cache into the next call as ``h_cache_T``. Phases 2 + 4 give the headline +overlap (NIC sends concurrent with the prefetch target fwd). +""" + +# pyright: reportArgumentType=false + +from typing import Any + +import torch +import torch.distributed as dist +import torch.distributed.nn.functional as dist_fn +import torch.nn as nn +from torch import Tensor + +from param_decomp.component_model import CIOutputs, ComponentModel +from param_decomp.grad_clip import cross_pool_clip_grad_norm +from param_decomp.metrics.importance_minimality import ( + _finalize as _finalize_imp_min, +) +from param_decomp.metrics.importance_minimality import ( + _get_linear_annealed_p, + _per_component_sums, +) +from param_decomp.three_pool.layout import ThreePoolLayout +from param_decomp.three_pool.profiler import PhaseProfiler +from param_decomp.three_pool.runtime import _ThreePoolRuntime +from param_decomp.two_pool.runtime import autocast_bf16 + + +def step_ci( + layout: ThreePoolLayout, + component_model: ComponentModel, + optimizer: torch.optim.Optimizer, + ci_fn_params: list[nn.Parameter], + batch_T: Any, + batch_T_plus_1: Any | None, + h_cache_T: dict[str, Tensor] | None, + cfg: _ThreePoolRuntime, + current_frac_of_training: float, + profiler: PhaseProfiler | None = None, +) -> tuple[dict[str, float], dict[str, Tensor] | None]: + """One CI-pool training step. + + Returns ``(metrics, h_cache_T_plus_1)``. On step 0, ``h_cache_T`` is + ``None`` and we compute it inline (phase 0). On the last step, + ``batch_T_plus_1`` is ``None`` and the prefetch is skipped. + """ + p = profiler if profiler is not None else PhaseProfiler(enabled=False) + assert layout.my_pool == "ci" + 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) + + if h_cache_T is None: + with p.phase("ci/0_target_fwd_T_sync"): + h_cache_T = _target_fwd_and_cache(component_model, batch_T_local, cfg.bf16_autocast) + + with p.phase("ci/1_ci_fn_fwd"), autocast_bf16(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) + + with p.phase("ci/2_async_send_ci"): + 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) + + with p.phase("ci/3_imp_min"): + 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: + with p.phase("ci/4_prefetch_target_fwd"): + h_cache_T_plus_1 = _target_fwd_and_cache( + component_model, batch_T_plus_1_local, cfg.bf16_autocast + ) + + with p.phase("ci/5_recv_g_ci_from_lw"): + g_ci_lw = layout.recv_g_ci_from_layerwise(cfg.c_per_site, seq_len, device) + with p.phase("ci/6_recv_g_ci_from_ppgd"): + 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) + + optimizer.zero_grad(set_to_none=True) + _fused_backward_through_ci_fn(loss_imp, ci, g_ci_total, layout, cfg, p) + + with p.phase("ci/9_in_pool_allreduce"): + layout.all_reduce_ci_fn_grads(ci_fn_params) + if cfg.grad_clip_norm_ci_fn is not None: + with p.phase("ci/9b_grad_clip"): + cross_pool_clip_grad_norm( + ci_fn_params, + cfg.grad_clip_norm_ci_fn, + group=layout.world.ci_pool_group, + n_replicas=layout.world.n_ci, + ) + with p.phase("ci/10_opt_step"): + optimizer.step() + + with p.phase("ci/11_wait_sends"): + for w in [*send_works_lw, *send_works_pgd]: + w.wait() + del send_bufs_lw, send_bufs_pgd + + # 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. + metrics = { + "loss/imp": loss_imp.item(), + "_raw/imp_num": loss_imp.item() / layout.world.n_ci, + } + return metrics, h_cache_T_plus_1 + + +def _slice_batches_for_ci( + batch_T: Any, batch_T_plus_1: Any | None, layout: ThreePoolLayout +) -> tuple[Any, Any | None]: + """Pull this CI rank's DP shard of batch T and T+1. + + CI pool is multi-rank DP across the global batch; every CI rank gets a + disjoint slice. + """ + sl = layout.my_batch_slice_ci() + batch_T_local = batch_T[sl] if isinstance(batch_T, Tensor) else batch_T + batch_T_plus_1_local = ( + batch_T_plus_1[sl] + if (batch_T_plus_1 is not None and isinstance(batch_T_plus_1, Tensor)) + else batch_T_plus_1 + ) + return batch_T_local, batch_T_plus_1_local + + +def _seq_len_from_ci(ci_lower: dict[str, Tensor]) -> int: + """Extract seq_len from CI value shape. Asserts [B, S, C] layout.""" + sample_ci = next(iter(ci_lower.values())) + assert sample_ci.ndim == 3, f"expected CI shape [B, S, C]; got {sample_ci.shape}" + return sample_ci.shape[1] + + +def _assert_ci_shapes( + ci_lower: dict[str, Tensor], + layout: ThreePoolLayout, + seq_len: int, + cfg: _ThreePoolRuntime, +) -> 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. + """ + batch_local_ci = layout.world.batch_local_ci + for s, c in cfg.c_per_site.items(): + t = ci_lower[s] + assert t.shape == (batch_local_ci, seq_len, c), ( + f"ci.lower_leaky[{s!r}] shape {tuple(t.shape)} != " + f"expected ({batch_local_ci}, {seq_len}, {c})" + ) + + +def _assemble_g_ci_total( + g_ci_lw: dict[str, Tensor], + g_ci_pgd: dict[str, Tensor], + layout: ThreePoolLayout, + cfg: _ThreePoolRuntime, + seq_len: int, +) -> dict[str, Tensor]: + """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] = {} + for s in layout.world.all_sites: + c = cfg.c_per_site[s] + lw, pgd = g_ci_lw[s], g_ci_pgd[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})" + ) + assert pgd.shape == (batch_local_ci, seq_len, c), ( + 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 + + +def _fused_backward_through_ci_fn( + loss_imp: Tensor, + ci: CIOutputs, + g_ci_total: dict[str, Tensor], + layout: ThreePoolLayout, + cfg: _ThreePoolRuntime, + p: PhaseProfiler, +) -> None: + """Phase ci/8. Backward through the CI fn graph. + + Two gradient seeds enter the graph: + * ``coeff_imp * loss_imp`` — flows via ``ci.upper_leaky``. Its backward + traverses the autograd-aware ``dist_fn.all_reduce`` (96 NCCL + broadcasts back to every CI rank) before reaching the CI fn output. + * ``g_CI_total[s]`` per site — injected directly on ``ci.lower_leaky[s]``. + 96 separate gradient seeds rejoining at the shared CI fn output. + + Diagnostic split: each seed runs its own ``torch.autograd.backward`` call + with ``retain_graph=True`` on the first so the second still sees the + graph. Gradient accumulation onto the CI fn params is the same as one + fused call. This is purely so the per-phase profiler can attribute time + between the two backward paths — to find out which one dominates and + where to optimize next. + """ + 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] + with p.phase("ci/8a_bwd_lower_leaky_only"): + torch.autograd.backward( + tensors=lower_leaky_tensors, + grad_tensors=g_ci_total_seeds, + retain_graph=True, + ) + with p.phase("ci/8b_bwd_imp_min_only"): + torch.autograd.backward(tensors=[scaled_imp], grad_tensors=[None]) + + +def _target_fwd_and_cache( + component_model: ComponentModel, + batch: Any, + bf16_autocast: bool, +) -> dict[str, Tensor]: + """target_fwd (no grad) returning the per-site pre-weight act cache. + + Used by phase 0 (on-demand H_T) and phase 4 (dead-time H_{T+1} prefetch). + Cache is upcast to fp32 so the downstream CI fn fwd gets fp32 inputs. + """ + with torch.no_grad(), autocast_bf16(bf16_autocast): + out = component_model(batch, cache_type="input") + return {k: v.to(torch.float32) for k, v in out.cache.items()} + + +def _importance_minimality_loss( + ci_upper: dict[str, Tensor], + current_frac_of_training: float, + cfg: _ThreePoolRuntime, + ci_pool_group: dist.ProcessGroup, + n_ci_pool: int, +) -> Tensor: + """Exact (across CI pool) importance-minimality loss. + + Each CI rank computes ``per_component_sums`` on its slice; we SUM-reduce + them across the CI pool with the autograd-aware all_reduce. ``n_examples`` + is uniform across CI ranks (same batch_local_ci) so we multiply rather + than reduce. + + Autograd note: ``torch.distributed.nn.functional.all_reduce`` is the + autograd-aware variant — forward sums, backward broadcasts the upstream + gradient unchanged to every rank's input (correct for SUM since + ``∂global/∂local_i = 1`` for all i). + """ + annealed_p = _get_linear_annealed_p( + current_frac_of_training=current_frac_of_training, + initial_p=cfg.imp_min_pnorm, + p_anneal_start_frac=cfg.imp_min_p_anneal_start_frac, + p_anneal_final_p=cfg.imp_min_p_anneal_final_p, + p_anneal_end_frac=cfg.imp_min_p_anneal_end_frac, + ) + per_component_sums, n_examples = _per_component_sums( + ci_upper_leaky=ci_upper, pnorm=annealed_p, eps=cfg.imp_min_eps + ) + if n_ci_pool > 1: + per_component_sums = { + k: dist_fn.all_reduce(v, op=dist.ReduceOp.SUM, group=ci_pool_group) + for k, v in per_component_sums.items() + } + n_examples = n_examples * n_ci_pool + return _finalize_imp_min( + per_component_sums=per_component_sums, + n_examples=n_examples, + beta=cfg.imp_min_beta, + # world_size=1 because per_component_sums + n_examples are now global, + # so the log term inside _finalize computes log2(1 + global_sum) directly. + world_size=1, + ) diff --git a/param_decomp/three_pool/step_layerwise.py b/param_decomp/three_pool/step_layerwise.py new file mode 100644 index 000000000..876249cf7 --- /dev/null +++ b/param_decomp/three_pool/step_layerwise.py @@ -0,0 +1,427 @@ +"""Layerwise pool training step. + +Recast of ``two_pool.pool_a.step_pool_a`` with the CI-fn-side bits moved out +to the CI pool, plus an optional async pipeline (``defer_vu_opt=True``) that +hides the V/U opt step + V/U ship-back behind T+1's CI fn forward window on +the CI pool. + +Phases (numbered to match ``DESIGN.md`` ``lw/N``): + + A1. Post async irecv for CI_T from the owning CI rank (overlaps with A2). + A2. target_fwd(batch_T) → L_T on this rank's batch slice. + B. Async mode only: finalize prev iter's deferred all_reduce → grad clip → + AdamW step → async ship V/U. Blocking waits here overlap with A2's + kernels (default CUDA stream) on the GPU. + C. Zero ``param.grad`` for fresh accumulation. + D1. Faithfulness loss + backward (V/U-only, doesn't need CI). + D2. Wait CI recv; re-leaf as fp32 ``requires_grad=True`` so the layerwise + backward populates ``leaf.grad`` for D4. + D3. Layerwise stoch recon, streaming per owned site — the + semantically meaningful for-loop lives at this step level (one + forward+backward per site bounds peak activation memory). + D4. Send g_CI back to CI pool (per-rank, on owned sites). + D5. Recv g_VU from PPGD pool (block leader recvs, then in-block bcast). + D6. Combine V/U grads: faith + layerwise already in ``.grad``; add PPGD's. + E. Tail. Sync mode: in-block all_reduce → grad clip → AdamW → async send + V/U. Async mode: kickoff async all_reduce, return its pending state. + +Phases 1 + 3 give the headline overlap (CI recv on the NIC while +faithfulness runs on the GPU). The phase-B + phase-E async kickoff give the +deferred-mode overlap (they hide behind T+1's CI fn forward window on CI pool). +""" + +# pyright: reportArgumentType=false, reportOperatorIssue=false, reportAttributeAccessIssue=false + +from typing import Any + +import torch +import torch.distributed as dist # noqa: F401 (used in type hints) +import torch.nn as nn +from torch import Tensor + +from param_decomp._trace import phase_trace_enabled, trace +from param_decomp.component_model import ComponentModel +from param_decomp.grad_clip import cross_pool_clip_grad_norm +from param_decomp.masks import make_mask_infos +from param_decomp.three_pool.layout import ThreePoolLayout +from param_decomp.three_pool.profiler import PhaseProfiler +from param_decomp.three_pool.runtime import _ThreePoolRuntime +from param_decomp.two_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp.two_pool.runtime import autocast_bf16 + +PendingAllReduce = list[tuple[list[Tensor], Tensor, "dist.Work"]] + + +def step_layerwise( + layout: ThreePoolLayout, + component_model: ComponentModel, + optimizer: torch.optim.Optimizer, + all_params: list[nn.Parameter], + batch: Any, + cfg: _ThreePoolRuntime, + strategy: LayerwiseLossStrategy, + *, + defer_vu_opt: bool, + prev_pending_all_reduce: PendingAllReduce | None, + profiler: PhaseProfiler | None = None, +) -> tuple[dict[str, float], PendingAllReduce | None]: + """One LW step. Branches on ``defer_vu_opt`` for sync vs async pipeline. + + Sync (``defer_vu_opt=False``): A → D → sync tail (all_reduce, clip, opt, + async send V/U). Returns ``(metrics, None)``. + + Async (``defer_vu_opt=True``): A → B-finalize-prev (concurrent w/ A2) → + D → kickoff async all_reduce, return its pending state. Caller must + thread it across iters and drain via ``finalize_layerwise_async_drain``. + """ + p = profiler if profiler is not None else PhaseProfiler(enabled=False) + 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): + with p.phase("lw/A1_post_async_recv_ci"): + ci_recv, ci_recv_works = 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 p.phase("lw/A2_target_fwd"), torch.no_grad(), autocast_bf16(cfg.bf16_autocast): + target_local = component_model(batch_local).detach() + + if defer_vu_opt and prev_pending_all_reduce is not None: + _finalize_prev_iter_async( + layout, component_model, optimizer, all_params, prev_pending_all_reduce, cfg, p + ) + + for param in all_params: + param.grad = None + + with strategy.context(component_model.target_model): + with p.phase("lw/D1_faith"): + loss_faith, faith_sum_sq_t, faith_numel = _faithfulness_loss( + component_model, device, cfg.numel_global + ) + (cfg.coeff_faith * loss_faith).backward() + + with p.phase("lw/D2_wait_ci_recv"): + for w in ci_recv_works: + w.wait() + 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 p.phase("lw/D3_layerwise"), autocast_bf16(cfg.bf16_autocast): + stoch_total_value = 0.0 + 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_value += (loss_s / n_positions).item() + stoch_n_owned = len(layout.my_owned_sites) + loss_stoch_value = stoch_total_value / stoch_n_owned + + with p.phase("lw/D4_send_g_ci"): + 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" + ) + layout.send_g_ci_to_ci_pool(g_ci_owned) + + v_grads_pgd, u_grads_pgd = _recv_g_vu_from_ppgd(layout, component_model, p) + _combine_vu_grads_in_place(component_model, layout, v_grads_pgd, u_grads_pgd, p) + + metrics = { + "loss/faith": loss_faith.item(), + "loss/stoch": loss_stoch_value, + "_raw/faith_num": faith_sum_sq_t.item(), + "_raw/faith_den": float(faith_numel), + "_raw/stoch_num": stoch_total_value, + "_raw/stoch_den": float(stoch_n_owned), + } + + if defer_vu_opt: + with p.phase("lw/E_kickoff_async_allreduce"): + new_pending = layout.async_all_reduce_grads_in_block_kickoff(all_params) + return metrics, new_pending + + _sync_tail(layout, component_model, optimizer, all_params, cfg, p) + return metrics, None + + +def finalize_layerwise_async_drain( + layout: ThreePoolLayout, + component_model: ComponentModel, + optimizer: torch.optim.Optimizer, + all_params: list[nn.Parameter], + pending_all_reduce: PendingAllReduce, + grad_clip_norm: float | None, +) -> None: + """End-of-training drain in async mode: finish the final iter's deferred + opt step. Skips the V/U async send since training is over. + """ + assert layout.my_pool == "layerwise" + _wait_pending_weight_send(component_model) + layout.wait_and_unflatten_all_reduce(pending_all_reduce) + if grad_clip_norm is not None: + cross_pool_clip_grad_norm( + all_params, + grad_clip_norm, + group=layout.world.layerwise_pool_group, + n_replicas=layout.world.n_per_block, + ) + optimizer.step() + + +def run_faithfulness_warmup_layerwise( + *, + component_model: ComponentModel, + component_params: list[nn.Parameter], + n_steps: int, + lr: float, + weight_decay: float, + numel_global: int, +) -> None: + """Single-pool-equivalent faithfulness warmup on the LW pool only. + + CI pool has no V/U; PPGD pool's V/U is a transient replica that gets + overwritten each step. So warmup only makes sense on LW. Mirrors + ``two_pool.pool_a.run_faithfulness_warmup_pool_a``. + """ + warmup_opt = torch.optim.AdamW(component_params, lr=lr, weight_decay=weight_decay) + for _ in range(n_steps): + warmup_opt.zero_grad() + device = component_params[0].device + loss, _, _ = _faithfulness_loss(component_model, device, numel_global) + loss.backward() + warmup_opt.step() + del warmup_opt + torch.cuda.empty_cache() + + +def _slice_batch_for_layerwise(batch: Any, layout: ThreePoolLayout) -> tuple[Any, int]: + """Pull this LW rank's batch slice + extract its seq_len.""" + sl = layout.my_batch_slice_lw() + batch_local = batch[sl] if isinstance(batch, Tensor) else batch + if isinstance(batch_local, Tensor): + seq_len = batch_local.shape[1] if batch_local.ndim >= 2 else 1 + else: + assert isinstance(batch_local, dict) and "input_ids" in batch_local + seq_len = batch_local["input_ids"].shape[1] + return batch_local, seq_len + + +def _finalize_prev_iter_async( + layout: ThreePoolLayout, + component_model: ComponentModel, + optimizer: torch.optim.Optimizer, + all_params: list[nn.Parameter], + prev_pending_all_reduce: PendingAllReduce, + cfg: _ThreePoolRuntime, + p: PhaseProfiler, +) -> None: + """Phase lw/B (async mode only). Finish the previous iter's deferred opt. + + Sequence: wait pending V/U send → wait + unflatten the prev all_reduce → + cross-pool grad clip → AdamW step → async ship updated V/U to PPGD. + Blocking waits in here overlap with phase A2's target_fwd kernels (which + run on the default CUDA stream while NCCL waits hit their own stream). + """ + with p.phase("lw/B1_wait_prev_weight_send"): + _wait_pending_weight_send(component_model) + with p.phase("lw/B2_wait_and_unflatten_allreduce"): + layout.wait_and_unflatten_all_reduce(prev_pending_all_reduce) + if cfg.grad_clip_norm_components is not None: + with p.phase("lw/B2b_grad_clip"): + cross_pool_clip_grad_norm( + all_params, + cfg.grad_clip_norm_components, + group=layout.world.layerwise_pool_group, + n_replicas=layout.world.n_per_block, + ) + with p.phase("lw/B3_opt_step"): + optimizer.step() + with p.phase("lw/B4_async_send_vu"): + _async_send_owned_vu_to_ppgd(component_model, layout) + + +def _releaf_ci_fp32_for_grads( + ci_recv: dict[str, Tensor], owned_sites: tuple[str, ...] +) -> dict[str, Tensor]: + """Upcast CI (bf16 on the wire) to fp32 and re-leaf with ``requires_grad=True`` + so the layerwise backward populates ``leaf.grad`` that the CI pool merges + into its CI-fn fp32 grads. + """ + return { + s: ci_recv[s].detach().to(torch.float32).clone().requires_grad_(True) for s in owned_sites + } + + +def _assert_ci_recv_shapes( + ci_recv_leaves: dict[str, Tensor], + layout: ThreePoolLayout, + seq_len: int, + cfg: _ThreePoolRuntime, +) -> None: + """Sanity-check the CI leaves match what the CI pool said it'd send. + + Catches a wrong ``c_per_site`` config or a per-rank batch mismatch fast. + """ + batch_local_lw = layout.world.batch_local_lw + for s in layout.my_owned_sites: + c = cfg.c_per_site[s] + t = ci_recv_leaves[s] + assert t.shape == (batch_local_lw, seq_len, c), ( + f"ci_recv_leaves[{s!r}] shape {tuple(t.shape)} != " + f"expected ({batch_local_lw}, {seq_len}, {c})" + ) + + +def _layerwise_one_site( + component_model: ComponentModel, + batch_local: Any, + target_local: Tensor, + ci_recv_leaves: dict[str, Tensor], + site: str, + strategy: LayerwiseLossStrategy, +) -> tuple[Tensor, int]: + """Phase lw/D3 (per-site body). One stochastic masked forward + recon. + + Returns ``(sum_loss, n_positions)`` raw — caller scales by + ``coeff_stoch / (n_positions * n_sites_total)`` and calls ``backward()`` + so the per-site graph is freed between iterations (bounds peak memory). + """ + ci_s = ci_recv_leaves[site] + u = torch.rand_like(ci_s) + mask = ci_s + (1 - ci_s) * u + delta = component_model.target_weight(site) - component_model.components[site].weight + delta_mask = torch.rand(ci_s.shape[:-1], device=ci_s.device, dtype=ci_s.dtype) + mask_infos = make_mask_infos( + {site: mask}, + weight_deltas_and_masks={site: (delta, delta_mask)}, + routing_masks="all", + ) + pred = component_model(batch_local, mask_infos=mask_infos) + loss, n_positions = strategy.recon_loss(pred=pred, target=target_local) + return loss, n_positions + + +def _recv_g_vu_from_ppgd( + layout: ThreePoolLayout, + component_model: ComponentModel, + p: PhaseProfiler, +) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + """Phase lw/D5. Recv V/U grads from PPGD pool (leader recvs, in-block bcast).""" + with p.phase("lw/D5_recv_g_vu_from_ppgd"): + 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) + 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" + ) + assert u_grads_pgd[s].shape == component_model.components[s].U.shape, ( + f"u_grads_pgd[{s!r}] shape mismatch from PPGD send" + ) + return v_grads_pgd, u_grads_pgd + + +def _combine_vu_grads_in_place( + component_model: ComponentModel, + layout: ThreePoolLayout, + v_grads_pgd: dict[str, Tensor], + u_grads_pgd: dict[str, Tensor], + p: PhaseProfiler, +) -> None: + """Phase lw/D6. Add PPGD's V/U grads to .grad (which already has faith+lw).""" + with p.phase("lw/D6_combine_vu_grads"): + for s in layout.my_owned_sites: + comp = component_model.components[s] + assert comp.V.grad is not None and comp.U.grad is not None, ( + "faith + layerwise should have populated V/U .grad" + ) + comp.V.grad.add_(v_grads_pgd[s]) + comp.U.grad.add_(u_grads_pgd[s]) + + +def _sync_tail( + layout: ThreePoolLayout, + component_model: ComponentModel, + optimizer: torch.optim.Optimizer, + all_params: list[nn.Parameter], + cfg: _ThreePoolRuntime, + p: PhaseProfiler, +) -> None: + """Phase lw/E (sync mode). Blocking all_reduce → clip → AdamW → async send V/U. + + Functionally equivalent to the pre-deferral 2-pool tail; safe to coexist + with PPGD's sync recv at end of step T. + """ + with p.phase("lw/E1_wait_prev_weight_send"): + _wait_pending_weight_send(component_model) + with p.phase("lw/E2_in_block_allreduce"): + layout.all_reduce_grads_in_block(all_params) + if cfg.grad_clip_norm_components is not None: + with p.phase("lw/E2b_grad_clip"): + cross_pool_clip_grad_norm( + all_params, + cfg.grad_clip_norm_components, + group=layout.world.layerwise_pool_group, + n_replicas=layout.world.n_per_block, + ) + with p.phase("lw/E3_opt_step"): + optimizer.step() + with p.phase("lw/E4_async_send_vu"): + _async_send_owned_vu_to_ppgd(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.""" + 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, + ) + + +def _wait_pending_weight_send(component_model: ComponentModel) -> None: + """Wait + clear any pending async V/U send from a previous iter. + + 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] + + +def _faithfulness_loss( + component_model: ComponentModel, device: torch.device, numel_global: int +) -> tuple[Tensor, Tensor, int]: + """‖W_target − VU.T‖²_F / numel_global, summed across this rank's owned sites. + + See ``two_pool.pool_a._faithfulness_loss`` for why we divide by + ``numel_global`` not ``numel_owned`` — keeps per-element grad scale aligned + with single-pool's, so the unclipped faithfulness warmup converges to the + same V/U as single-pool. + + Returns ``(scalar_loss, sum_sq, numel_owned)`` — the ``numel_owned`` is the + denominator the logger uses for ``SUM(num) / SUM(den)`` global-ratio + reconstruction across blocks. + """ + weight_deltas = component_model.calc_weight_deltas() + sum_sq = torch.zeros((), device=device) + numel_owned = 0 + for d in weight_deltas.values(): + sum_sq = sum_sq + (d**2).sum() + numel_owned += d.numel() + return sum_sq / numel_global, sum_sq, numel_owned diff --git a/param_decomp/three_pool/step_ppgd.py b/param_decomp/three_pool/step_ppgd.py new file mode 100644 index 000000000..4fee11a38 --- /dev/null +++ b/param_decomp/three_pool/step_ppgd.py @@ -0,0 +1,303 @@ +"""PPGD pool training step — split into ``step_ppgd`` and ``finalize_ppgd_async_drain``. + +Recast of ``two_pool.pool_b.step_pool_b`` for 3-pool. Same PPGD math; the +differences are (a) CI comes from the CI pool not LW leaders, (b) g_CI goes +back to CI pool not LW, and (c) no final outer source-step (the warmup inner +loop owns source updates). + +Phases (numbered to match ``DESIGN.md`` ``ppgd/N``): + + A1. Post async irecv for CI_T from the owning CI rank (concurrent with A2). + A2. target_fwd(batch_T) on the rank's PPGD slice → L_T. + B. Async-mode only: wait + copy prev iter's V/U recv (overlaps with A2's + kernels on the default CUDA stream while NCCL waits run on theirs). + D1. Compute weight_deltas (V/U-dependent — fresh now). + D2. Wait CI recv; re-leaf as fp32 for downstream CI grad extraction. + D3. PPGD warmup refines persistent adversarial sources in place. + D4. Recon loss with refined sources. + D5. Backward: ``torch.autograd.grad`` for g_VU + g_CI (no source backward). + D6. Sum-reduce g_VU within PPGD pool → each rank holds the full-batch grad. + D7. Send g_VU to LW block leaders (PPGD-leader-only). + D8. Send g_CI to CI pool (every rank, on its batch slice). + E. Recv updated V/U from LW: + * sync mode → blocking recv at end of step (returns + ``(metrics, None)``); + * async mode → kickoff async recv, return handles for next step's + phase B. + +Architectural note on the absence of a per-site for-loop at the step level: +PPGD really does *one* fused forward + *one* fused backward across all sites +in D3-D5 — the per-site iteration is just gathering tensors into flat +lists for ``torch.autograd.grad`` (an implementation detail of that API +producing multiple grads in one call). The semantically meaningful unit at +the step level is the whole pool, not a single site. +""" + +# pyright: reportArgumentType=false + +from typing import Any + +import torch +import torch.distributed as dist # noqa: F401 (used in type hints) +from torch import Tensor + +from param_decomp.component_model import ComponentModel +from param_decomp.metrics.persistent_pgd_state import PersistentPGDState +from param_decomp.three_pool.layout import LayerwiseBlockGroup, ThreePoolLayout +from param_decomp.three_pool.profiler import PhaseProfiler +from param_decomp.three_pool.runtime import _ThreePoolRuntime +from param_decomp.two_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp.two_pool.runtime import autocast_bf16 + +PendingRecvVU = list[tuple["LayerwiseBlockGroup", Tensor, "dist.Work"]] + + +def step_ppgd( + layout: ThreePoolLayout, + component_model: ComponentModel, + ppgd_state: PersistentPGDState, + batch: Any, + cfg: _ThreePoolRuntime, + strategy: LayerwiseLossStrategy, + step: int, + n_steps: int, + *, + defer_vu_opt: bool, + prev_pending_recv_vu: PendingRecvVU | None, + profiler: PhaseProfiler | None = None, +) -> tuple[dict[str, float], PendingRecvVU | None]: + """One PPGD step. Branches on ``defer_vu_opt`` for sync vs async pipeline. + + Async mode is required when LW defers — otherwise PPGD's blocking sync recv + at end of step T would deadlock against LW's deferred async send (which + doesn't fire until top of T+1). + """ + p = profiler if profiler is not None else PhaseProfiler(enabled=False) + assert layout.my_pool == "ppgd" + device = next(component_model.parameters()).device + all_sites = list(layout.world.all_sites) + + ppgd_state.update_lr(step=step, total_steps=n_steps) + batch_local, seq_len = _slice_batch_for_ppgd(batch, layout) + v_templates, u_templates = _vu_templates(component_model, all_sites) + + with strategy.context(component_model.target_model): + with p.phase("pgd/A1_post_async_recv_ci"): + ci_recv, ci_recv_works = layout.async_recv_ci_from_ci_pool_ppgd( + cfg.c_per_site, seq_len=seq_len, device=device + ) + with p.phase("pgd/A2_target_fwd"), torch.no_grad(), autocast_bf16(cfg.bf16_autocast): + target_out = component_model(batch_local).detach() + + # Async-mode bookkeeping: this iter's V/U arrived in the prev iter's + # async kickoff. Wait + copy here, overlapping the NCCL wait with the + # A2 target_fwd kernels enqueued above. + if defer_vu_opt and prev_pending_recv_vu is not None: + _wait_and_copy_prev_vu_into_model( + layout, component_model, prev_pending_recv_vu, v_templates, u_templates, p + ) + + with p.phase("pgd/D1_calc_weight_deltas"): + weight_deltas = component_model.calc_weight_deltas() + with p.phase("pgd/D2_wait_ci_recv"): + for w in ci_recv_works: + w.wait() + ci_scratch = _releaf_ci_fp32_for_grads(ci_recv) + _assert_ci_scratch_shapes(ci_scratch, layout, seq_len, cfg) + + with p.phase("pgd/D3_warmup"), autocast_bf16(cfg.bf16_autocast): + ppgd_state.warmup( + model=component_model, + batch=batch_local, + target_out=target_out, + ci=ci_scratch, + weight_deltas=weight_deltas, + ) + with p.phase("pgd/D4_recon"), autocast_bf16(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}" + + # Scale by 1/n_ppgd so D6's SUM-reduce produces the full-batch gradient. + total_ppgd = cfg.coeff_ppgd * (sum_loss / n_examples) / layout.world.n_ppgd + + v_grads, u_grads, ci_grads = _autograd_grad_for_vu_and_ci( + total_ppgd, component_model, ci_scratch, all_sites, p + ) + with p.phase("pgd/D6_in_pool_sum_reduce"): + layout.sum_reduce_ppgd_grads([*v_grads.values(), *u_grads.values()]) + with p.phase("pgd/D7_send_g_vu_to_lw"): + layout.send_g_vu_to_layerwise(v_grads, u_grads) + with p.phase("pgd/D8_send_g_ci_to_ci_pool"): + layout.send_g_ci_to_ci_pool_ppgd(ci_grads) + + # Phase-tag the metrics dict construction so the per-phase profiler + # surfaces the cost of the ``.item()`` device-host syncs. With async + # NCCL ops in D6/D7/D8 still in flight on side streams, these + # ``cudaMemcpyAsync + cudaStreamSync`` calls are where PPGD's wait + # actually lives. + with p.phase("pgd/Dx_metrics_sync"): + metrics = { + "loss/ppgd": (sum_loss / n_examples).item(), + "_raw/ppgd_num": sum_loss.item(), + "_raw/ppgd_den": float(n_examples), + } + + if defer_vu_opt: + with p.phase("pgd/E_kickoff_async_recv_vu"): + new_pending = layout.async_recv_updated_vu_from_layerwise_kickoff( + v_templates, u_templates + ) + return metrics, new_pending + + with p.phase("pgd/E_sync_recv_vu"): + v_new, u_new = layout.recv_updated_vu_from_layerwise(v_templates, u_templates) + _copy_vu_into_model_in_place(component_model, v_new, u_new, all_sites) + return metrics, None + + +def finalize_ppgd_async_drain( + layout: ThreePoolLayout, + component_model: ComponentModel, + pending_recv_vu: PendingRecvVU, +) -> None: + """End-of-training drain in async mode: complete the final iter's V/U recv + so the saved checkpoint (gathered from LW pool's V/U separately) is + consistent with what PPGD pool used last. + """ + assert layout.my_pool == "ppgd" + all_sites = list(layout.world.all_sites) + v_templates, u_templates = _vu_templates(component_model, all_sites) + v_new, u_new = layout.wait_and_unpack_updated_vu(pending_recv_vu, v_templates, u_templates) + _copy_vu_into_model_in_place(component_model, v_new, u_new, all_sites) + + +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() + batch_local = batch[sl] if isinstance(batch, Tensor) else batch + if isinstance(batch_local, Tensor): + seq_len = batch_local.shape[1] if batch_local.ndim >= 2 else 1 + else: + assert isinstance(batch_local, dict) and "input_ids" in batch_local + seq_len = batch_local["input_ids"].shape[1] + return batch_local, seq_len + + +def _vu_templates( + component_model: ComponentModel, all_sites: list[str] +) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + """V/U tensors per site — used as recv buffers in async kickoff/drain.""" + v_templates: dict[str, Tensor] = {s: component_model.components[s].V for s in all_sites} + u_templates: dict[str, Tensor] = {s: component_model.components[s].U for s in all_sites} + return v_templates, u_templates + + +def _wait_and_copy_prev_vu_into_model( + layout: ThreePoolLayout, + component_model: ComponentModel, + prev_pending_recv_vu: PendingRecvVU, + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + p: PhaseProfiler, +) -> None: + """Phase ppgd/B (async mode only). Wait + copy prev iter's V/U recv into model. + + Blocking wait overlaps with the A2 target_fwd kernels enqueued just above — + that's the headline win of async mode. + """ + with p.phase("pgd/B_wait_and_copy_prev_vu"): + v_new, u_new = layout.wait_and_unpack_updated_vu( + prev_pending_recv_vu, v_templates, u_templates + ) + _copy_vu_into_model_in_place(component_model, v_new, u_new, list(layout.world.all_sites)) + + +def _releaf_ci_fp32_for_grads(ci_recv: dict[str, Tensor]) -> dict[str, Tensor]: + """Upcast CI (bf16 on the wire) to fp32 and re-leaf with ``requires_grad=True`` + so the autograd.grad call below populates fp32 ``.grad`` that the CI pool + can merge into its fp32 CI-fn grads. + """ + return { + s: v.detach().to(torch.float32).clone().requires_grad_(True) for s, v in ci_recv.items() + } + + +def _assert_ci_scratch_shapes( + ci_scratch: dict[str, Tensor], + layout: ThreePoolLayout, + seq_len: int, + cfg: _ThreePoolRuntime, +) -> None: + """Sanity-check the CI scratch tensors match what the CI pool said it'd send. + + Catches a wrong ``c_per_site`` config or a per-rank batch mismatch fast. + """ + batch_local_ppgd = layout.world.batch_local_ppgd + for s, c in cfg.c_per_site.items(): + t = ci_scratch[s] + assert t.shape == (batch_local_ppgd, seq_len, c), ( + f"ci_scratch[{s!r}] shape {tuple(t.shape)} != " + f"expected ({batch_local_ppgd}, {seq_len}, {c})" + ) + + +def _autograd_grad_for_vu_and_ci( + total_ppgd: Tensor, + component_model: ComponentModel, + ci_scratch: dict[str, Tensor], + all_sites: list[str], + p: PhaseProfiler, +) -> tuple[dict[str, Tensor], dict[str, Tensor], dict[str, Tensor]]: + """Phase ppgd/D5. ``torch.autograd.grad`` for V/U + CI in one fused backward. + + The flat ``params + ci_list`` arrangement is an artifact of + ``autograd.grad`` returning grads in input order — we then split back into + three dicts keyed by site. No source backward here: the warmup inner loop + already updated the sources. ``retain_graph=False`` since this is the last + graph use. + """ + with p.phase("pgd/D5_backward"): + params: list[Tensor] = [] + for s in all_sites: + params.append(component_model.components[s].V) + params.append(component_model.components[s].U) + ci_list = [ci_scratch[s] for s in all_sites] + grads = torch.autograd.grad(total_ppgd, params + ci_list, retain_graph=False) + + n_sites = len(all_sites) + v_grads = {s: grads[2 * i] for i, s in enumerate(all_sites)} + u_grads = {s: grads[2 * i + 1] for i, s in enumerate(all_sites)} + ci_grads = {s: grads[2 * n_sites + i] for i, s in enumerate(all_sites)} + for s in all_sites: + assert v_grads[s].shape == component_model.components[s].V.shape, ( + f"v_grad[{s!r}] shape mismatch" + ) + assert u_grads[s].shape == component_model.components[s].U.shape, ( + f"u_grad[{s!r}] shape mismatch" + ) + assert ci_grads[s].shape == ci_scratch[s].shape, f"ci_grad[{s!r}] shape mismatch" + return v_grads, u_grads, ci_grads + + +def _copy_vu_into_model_in_place( + component_model: ComponentModel, + v_new: dict[str, Tensor], + u_new: dict[str, Tensor], + all_sites: list[str], +) -> None: + """In-place copy of received V/U into the model's component params. + + In-place rather than reassign — defense in depth against anyone keeping a + Python reference to the original tensors (e.g. an optimizer state). + """ + with torch.no_grad(): + for s in all_sites: + component_model.components[s].V.copy_(v_new[s]) + component_model.components[s].U.copy_(u_new[s]) diff --git a/param_decomp/training_state.py b/param_decomp/training_state.py index 35d4b726a..b2121ca22 100644 --- a/param_decomp/training_state.py +++ b/param_decomp/training_state.py @@ -1,8 +1,14 @@ -"""`TrainingState`: the canonical persisted state of a 1-pool training run. +"""Canonical training-state dataclasses persisted to `training_.pth`. -Lives in its own module so both `param_decomp.optimize` (where `Trainer` -produces it) and `param_decomp.run_sink` (where the `RunSink` Protocol -consumes it) can import without a cycle. +`TrainingState` is the 1-pool shape; `ThreePoolTrainingState` is the 3-pool shape. +The two are deliberately separate (no shared base class) — common fields are +duplicated. Pool-specific fields (e.g. `three_pool_config`, `ppgd_state_by_rank`) +live on the concrete dataclass that needs them. + +Lives in its own module so both `param_decomp.optimize` / +`param_decomp.three_pool.optimize` (where trainers produce these) and +`param_decomp.run_sink` (where they're consumed by the sink protocol) can +import without a cycle. """ from dataclasses import dataclass @@ -32,3 +38,33 @@ class TrainingState: components_optimizer: dict[str, dict[str, Any]] ci_fn_optimizer: dict[str, dict[str, Any]] loss_metrics: dict[str, dict[str, Any]] + + +@dataclass(frozen=True) +class ThreePoolTrainingState: + """Canonical 3-pool training state, persisted to `training_.pth`. + + Produced by `ThreePoolTrainer.snapshot()` on rank 0 only — non-rank-0 ranks + return `None` instead (the lab sink is silent on those ranks anyway). + Rank 0 assembles the canonical state from cross-rank gathers: + + * `component_model`: full gathered model state (every site's V/U + the CI fn). + * `components_optimizer` / `ci_fn_optimizer`: per-parameter optimizer state + keyed by parameter name, merged across the layerwise + CI pools. Same + shape as 1-pool's `TrainingState` — topology-independent. + * `ppgd_state_by_rank`: PPGD adversarial sources, keyed by PPGD rank id. + Genuinely rank-coupled for `PerBatchPerPositionScope` (sized by the + rank-local batch slice). + * `layout_fingerprint`: 3-pool world layout summary. Checked on resume to + flag incompatible topologies. + """ + + step: int + pd_config: dict[str, Any] + runtime_config: dict[str, Any] + three_pool_config: dict[str, Any] + layout_fingerprint: dict[str, Any] + component_model: dict[str, Tensor] + components_optimizer: dict[str, dict[str, Any]] + ci_fn_optimizer: dict[str, dict[str, Any]] + ppgd_state_by_rank: dict[int, dict[str, Any]] diff --git a/param_decomp/two_pool/__init__.py b/param_decomp/two_pool/__init__.py new file mode 100644 index 000000000..62aeb7100 --- /dev/null +++ b/param_decomp/two_pool/__init__.py @@ -0,0 +1,26 @@ +"""2-pool training strategy for SPD on large frozen target models. + +Splits GPUs into two heterogeneous pools that run in wall-clock parallel: + + Pool A (home) owns canonical state — components V/U + per-module CI fn + + optimizer state, sharded by site. Runs target+CI forward, + per-site layerwise loss, faithfulness, importance-minimality. + Pool B (scratchpad) stateless — holds a transient replica of all V/U for the + full-model PPGD forward. Returns per-site V/U grads and + gradients w.r.t. the CI values it received. + +``optimize_two_pool`` mirrors :func:`param_decomp.optimize.optimize`'s call +shape. ``TwoPoolConfig`` declares the topology (block groups + pool-B ranks). +""" + +from param_decomp.two_pool.config import BlockGroupSpec, TwoPoolConfig +from param_decomp.two_pool.optimize import TwoPoolTrainer, optimize_two_pool +from param_decomp.two_pool.profiler import PhaseProfiler + +__all__ = [ + "BlockGroupSpec", + "PhaseProfiler", + "TwoPoolConfig", + "TwoPoolTrainer", + "optimize_two_pool", +] diff --git a/param_decomp/two_pool/config.py b/param_decomp/two_pool/config.py new file mode 100644 index 000000000..9ad94ff34 --- /dev/null +++ b/param_decomp/two_pool/config.py @@ -0,0 +1,93 @@ +"""Serializable config for 2-pool training topology. + +``TwoPoolConfig`` carries only what's genuinely 2-pool-specific: which ranks +form pool A's block groups (replicated V/U + CI fn + shared owned sites) and +which ranks form pool B (stateless PPGD replica). + +Everything else (batch size, loss coefficients, optimizer LRs, PPGD config, +faithfulness warmup, LR schedules, autocast, ci_config, sigmoid_type, +module_info) lives on the regular ``RunConfig`` / ``PDConfig`` / +``RuntimeConfig`` that the 2-pool path consumes — so a 2-pool training run is +configured like a normal SPD run, plus this topology block. + +``validate_run_cfg_for_two_pool`` in ``driver_entry.py`` checks that the +RunConfig's ``pd.loss_metrics`` contains the four loss types the 2-pool path +implements (FaithfulnessLoss, ImportanceMinimalityLoss, +StochasticReconLayerwiseLoss, PersistentPGDReconLoss) and that none of the +unsupported loss types are silently set. +""" + +from typing import Self + +from pydantic import Field, model_validator + +from param_decomp.base_config import BaseConfig + + +class BlockGroupSpec(BaseConfig): + """One block-DDP group: the ranks that replicate V/U + CI fn for a shared + set of sites. The first rank is the block leader (canonical actor for + cross-pool sends). Within a group, in-block all-reduce keeps the replicas + in sync after each optimizer step. + + Serializable mirror of `param_decomp.two_pool.layout.BlockGroup`. The + layout module's dataclass is constructed from this at runtime. + """ + + ranks: list[int] = Field(..., description="Ranks that replicate V/U for `owned_sites`.") + owned_sites: list[str] = Field( + ..., description="Module paths in the target that this group owns." + ) + + +class TwoPoolConfig(BaseConfig): + """Topology for 2-pool training. Pairs with a regular ``RunConfig``. + + Every pool-A rank lives in exactly one ``BlockGroupSpec``. Pool-B ranks + are listed explicitly. Topology integrity checks (rank disjointness, + uniform N_per_block) run at validation time; checks that depend on + ``pd.batch_size`` (divisibility) run in ``validate_run_cfg_for_two_pool`` + since they need the paired ``RunConfig`` to evaluate. + """ + + block_groups: list[BlockGroupSpec] = Field( + ..., + description="Pool-A block groups. Each group's ranks replicate V/U + CI fn for " + "the group's owned sites.", + ) + pool_b_ranks: list[int] = Field( + ..., + description="Ranks assigned to pool B (stateless PPGD replica with DP across the batch).", + ) + use_fused_kl: bool = Field( + default=True, + description="If True (default), pool A + pool B bypass the LM head and " + "compute KL via the chunked fused linear+KL kernel — never materializes " + "[b_local, seq, vocab] tensors. If False, the unfused path is used " + "(materialize logits + standard recon_loss). Toggle for ablation of " + "the kernel vs pure topology levers; defaults to fused because the " + "kernel reliably cuts pool A peak ~50%% at b=64 s=2048 with negligible " + "step-time cost.", + ) + + @model_validator(mode="after") + def validate_topology(self) -> Self: + assert self.block_groups, "block_groups must be non-empty" + n_per_block = len(self.block_groups[0].ranks) + assert n_per_block > 0, "block_groups[0].ranks must be non-empty" + for bg in self.block_groups: + assert len(bg.ranks) == n_per_block, ( + f"all block groups must have the same N_per_block; got " + f"{n_per_block} vs {len(bg.ranks)} for sites {bg.owned_sites[:2]}..." + ) + assert bg.owned_sites, "block_group must own at least one site" + pool_a_flat = [r for bg in self.block_groups for r in bg.ranks] + assert len(set(pool_a_flat)) == len(pool_a_flat), ( + "the same rank appears in multiple block groups" + ) + pool_b = set(self.pool_b_ranks) + pool_a = set(pool_a_flat) + assert pool_a.isdisjoint(pool_b), ( + f"pool A and pool B ranks must be disjoint; overlap: {sorted(pool_a & pool_b)}" + ) + return self diff --git a/param_decomp/two_pool/layout.py b/param_decomp/two_pool/layout.py new file mode 100644 index 000000000..410ea2f35 --- /dev/null +++ b/param_decomp/two_pool/layout.py @@ -0,0 +1,774 @@ +"""World / TwoPoolLayout / BlockDDPLayout — the topology data model. + +`World` (and `BlockDDPWorld`) is purely declarative — identical content on every rank, +no per-rank fields. Built once at startup after `dist.init_process_group`. + +`TwoPoolLayout` (and `BlockDDPLayout`) wraps a World, adds this rank's perspective +(my_pool, my_owned_sites, my_slice_idx), and hangs the cross-pool comm orchestration +methods off itself. + +Six comm methods cover the four cross-pool exchanges: + send_owned_ci_to_pool_b ↔ recv_ci_from_owners (CI values, A → B) + send_pool_b_grads_to_owners ↔ recv_grads_from_pool_b (grads, B → A) + send_updated_weights_to_pool_b ↔ recv_updated_weights_from_owners (weights, A → B) + +Block-DDP variant additionally exposes `all_reduce_grads_in_block` for the in-block +DDP sync over replicated V/U + CI fn params within a block group. + +Origin: `nano_param_decomp.two_pool.layout` / `.block_ddp`. See those files for the +implementation arc; this module is the production version. +""" + +# pyright: reportIndexIssue=false, reportArgumentType=false, reportOperatorIssue=false + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Literal + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor + +# All cross-pool tensors are cast to this dtype on the wire. Halves the bytes +# transferred per step (~22 GB → 11 GB for LlamaSimpleMLP-12L at 24×2+16 topology). +# Pool B's forward + PPGD backward run inside bf16 autocast already, so bf16 V/U +# values are precision-neutral on use. V/U grads accumulating into pool A's fp32 +# .grad upcast back to fp32 on receive — same pattern as standard bf16 mixed +# precision training. +_WIRE_DTYPE: torch.dtype = torch.bfloat16 + +# ───────────────────────── plain (1 rank per block group) ───────────────────────── + + +@dataclass(frozen=True) +class World: + """Declarative 2-pool topology — identical content on every rank. + + Each site has exactly one owning pool-A rank. No in-block DDP here. + """ + + world_size: int + pool_a_ranks: tuple[int, ...] + pool_b_ranks: tuple[int, ...] + all_sites: tuple[str, ...] + site_owner: dict[str, int] + batch_global: int + pool_a_group: dist.ProcessGroup + pool_b_group: dist.ProcessGroup + + @property + def n_pool_a(self) -> int: + return len(self.pool_a_ranks) + + @property + def n_pool_b(self) -> int: + return len(self.pool_b_ranks) + + @property + def batch_local_b(self) -> int: + assert self.batch_global % self.n_pool_b == 0 + return self.batch_global // self.n_pool_b + + +def build_world( + pool_a_ranks: list[int], + pool_b_ranks: list[int], + all_sites: list[str], + site_owner: dict[str, int], + batch_global: int, +) -> World: + """Construct a World after dist.init_process_group on every rank. + + Every rank must call `dist.new_group` collectively for both pools. + """ + world_size = dist.get_world_size() + assert len(pool_a_ranks) + len(pool_b_ranks) == world_size, ( + f"pool ranks ({len(pool_a_ranks)} + {len(pool_b_ranks)}) != world_size ({world_size})" + ) + assert set(pool_a_ranks).isdisjoint(set(pool_b_ranks)) + for site, owner in site_owner.items(): + assert owner in pool_a_ranks + assert site in all_sites + pool_a_group = dist.new_group(ranks=pool_a_ranks) + pool_b_group = dist.new_group(ranks=pool_b_ranks) + return World( + world_size=world_size, + pool_a_ranks=tuple(pool_a_ranks), + pool_b_ranks=tuple(pool_b_ranks), + all_sites=tuple(all_sites), + site_owner=dict(site_owner), + batch_global=batch_global, + pool_a_group=pool_a_group, + pool_b_group=pool_b_group, + ) + + +@dataclass(frozen=True) +class TwoPoolLayout: + """This rank's view of the world + cross-pool comm orchestration.""" + + world: World + my_rank: int + my_pool: Literal["a", "b"] + my_is_pool_leader: bool + my_owned_sites: tuple[str, ...] + my_slice_idx: int | None + + @classmethod + def from_world(cls, world: World, my_rank: int) -> "TwoPoolLayout": + if my_rank in world.pool_a_ranks: + return cls( + world=world, + my_rank=my_rank, + my_pool="a", + my_is_pool_leader=(my_rank == world.pool_a_ranks[0]), + my_owned_sites=tuple(s for s in world.all_sites if world.site_owner[s] == my_rank), + my_slice_idx=None, + ) + elif my_rank in world.pool_b_ranks: + return cls( + world=world, + my_rank=my_rank, + my_pool="b", + my_is_pool_leader=(my_rank == world.pool_b_ranks[0]), + my_owned_sites=(), + my_slice_idx=world.pool_b_ranks.index(my_rank), + ) + raise ValueError(f"rank {my_rank} not in any pool") + + def is_my_site(self, site: str) -> bool: + return self.my_pool == "a" and self.world.site_owner[site] == self.my_rank + + def owner_of(self, site: str) -> int: + return self.world.site_owner[site] + + def slice_for_b_idx(self, slice_idx: int) -> slice: + b = self.world.batch_local_b + return slice(slice_idx * b, (slice_idx + 1) * b) + + def my_batch_slice(self) -> slice: + assert self.my_pool == "b" and self.my_slice_idx is not None + b = self.world.batch_local_b + return slice(self.my_slice_idx * b, (self.my_slice_idx + 1) * b) + + # ── Cross-pool comm ── + + def send_owned_ci_to_pool_b(self, ci_owned: dict[str, Tensor]) -> None: + assert self.my_pool == "a" + for site in self.world.all_sites: + if not self.is_my_site(site): + continue + for slice_idx, b_rank in enumerate(self.world.pool_b_ranks): + sl = self.slice_for_b_idx(slice_idx) + dist.send(ci_owned[site][sl].detach().contiguous(), dst=b_rank) + + def recv_ci_from_owners( + self, + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> dict[str, Tensor]: + assert self.my_pool == "b" + out: dict[str, Tensor] = {} + b_local = self.world.batch_local_b + for site in self.world.all_sites: + owner = self.world.site_owner[site] + buf = torch.empty(b_local, seq_len, site_to_c[site], device=device, dtype=dtype) + dist.recv(buf, src=owner) + out[site] = buf + return out + + def send_pool_b_grads_to_owners( + self, + v_grads: dict[str, Tensor], + u_grads: dict[str, Tensor], + ci_grads: dict[str, Tensor], + ) -> None: + assert self.my_pool == "b" + if self.my_is_pool_leader: + for site in self.world.all_sites: + owner = self.world.site_owner[site] + dist.send(v_grads[site].contiguous(), dst=owner) + dist.send(u_grads[site].contiguous(), dst=owner) + for site in self.world.all_sites: + owner = self.world.site_owner[site] + dist.send(ci_grads[site].contiguous(), dst=owner) + + def recv_grads_from_pool_b( + self, + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + ci_lower_owned: dict[str, Tensor], + ) -> tuple[dict[str, Tensor], dict[str, Tensor], dict[str, Tensor]]: + assert self.my_pool == "a" + v_grads: dict[str, Tensor] = {} + u_grads: dict[str, Tensor] = {} + ci_grads: dict[str, Tensor] = {} + + b_leader = self.world.pool_b_ranks[0] + for site in self.world.all_sites: + if not self.is_my_site(site): + continue + v_buf = torch.empty_like(v_templates[site]) + u_buf = torch.empty_like(u_templates[site]) + dist.recv(v_buf, src=b_leader) + dist.recv(u_buf, src=b_leader) + v_grads[site] = v_buf + u_grads[site] = u_buf + + b_local = self.world.batch_local_b + for site in self.world.all_sites: + if not self.is_my_site(site): + continue + full = ci_lower_owned[site] + _, S, C = full.shape + slices: list[Tensor] = [] + for b_rank in self.world.pool_b_ranks: + buf = torch.empty((b_local, S, C), dtype=full.dtype, device=full.device) + dist.recv(buf, src=b_rank) + slices.append(buf) + ci_grads[site] = torch.cat(slices, dim=0) + + return v_grads, u_grads, ci_grads + + def send_updated_weights_to_pool_b( + self, + v_owned: dict[str, Tensor], + u_owned: dict[str, Tensor], + ) -> None: + assert self.my_pool == "a" + for site in self.world.all_sites: + if not self.is_my_site(site): + continue + for b_rank in self.world.pool_b_ranks: + dist.send(v_owned[site].detach().contiguous(), dst=b_rank) + dist.send(u_owned[site].detach().contiguous(), dst=b_rank) + + def recv_updated_weights_from_owners( + self, + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + assert self.my_pool == "b" + v_new: dict[str, Tensor] = {} + u_new: dict[str, Tensor] = {} + for site in self.world.all_sites: + owner = self.world.site_owner[site] + v_buf = torch.empty_like(v_templates[site]) + u_buf = torch.empty_like(u_templates[site]) + dist.recv(v_buf, src=owner) + dist.recv(u_buf, src=owner) + v_new[site] = v_buf + u_new[site] = u_buf + return v_new, u_new + + +# ───────────────────────── Block-DDP (N ranks per block group) ───────────────────────── + + +@dataclass(frozen=True) +class BlockGroup: + """One block-DDP group: the ranks that replicate V/U + CI fn for a shared set of sites. + + The first rank is the block leader — the canonical actor for cross-pool sends. + Within a group, in-block all-reduce keeps the replicas in sync after each + optimizer step. + """ + + ranks: tuple[int, ...] + owned_sites: tuple[str, ...] + + @property + def leader(self) -> int: + return self.ranks[0] + + def __post_init__(self) -> None: + assert len(self.ranks) > 0, "block group must have at least one rank" + assert len(self.ranks) == len(set(self.ranks)), ( + f"duplicate ranks in block group: {self.ranks}" + ) + + +@dataclass(frozen=True) +class BlockDDPWorld: + """Pool A organized into block groups; each block group's ranks replicate V/U + CI fn.""" + + world_size: int + block_groups: tuple[BlockGroup, ...] + pool_b_ranks: tuple[int, ...] + all_sites: tuple[str, ...] + batch_global: int + pool_a_group: dist.ProcessGroup + pool_b_group: dist.ProcessGroup + block_group_groups: tuple[dist.ProcessGroup, ...] + # One process group per block: {block_leader} ∪ {pool_b_ranks}. Used for + # collective A-leader → pool-B broadcasts (and the symmetric receive). + # NCCL uses a tree topology so leader egress drops from N_pool_b× to 1×; + # multiple broadcasts in different groups run concurrently via async_op. + cross_pool_bcast_groups: tuple[dist.ProcessGroup, ...] + + @property + def n_blocks(self) -> int: + return len(self.block_groups) + + @property + def n_per_block(self) -> int: + size = len(self.block_groups[0].ranks) + assert all(len(bg.ranks) == size for bg in self.block_groups) + return size + + @property + def n_pool_a(self) -> int: + return sum(len(bg.ranks) for bg in self.block_groups) + + @property + def n_pool_b(self) -> int: + return len(self.pool_b_ranks) + + @property + def pool_a_ranks(self) -> tuple[int, ...]: + return tuple(r for bg in self.block_groups for r in bg.ranks) + + @property + def batch_local_a(self) -> int: + assert self.batch_global % self.n_per_block == 0 + return self.batch_global // self.n_per_block + + @property + def batch_local_b(self) -> int: + assert self.batch_global % self.n_pool_b == 0 + return self.batch_global // self.n_pool_b + + def block_idx_of_site(self, site: str) -> int: + for i, bg in enumerate(self.block_groups): + if site in bg.owned_sites: + return i + raise KeyError(site) + + def block_leader_of_site(self, site: str) -> int: + return self.block_groups[self.block_idx_of_site(site)].leader + + +def build_block_ddp_world( + block_groups: list[BlockGroup], + pool_b_ranks: list[int], + batch_global: int, +) -> BlockDDPWorld: + world_size = dist.get_world_size() + pool_a_ranks = [r for bg in block_groups for r in bg.ranks] + assert len(pool_a_ranks) + len(pool_b_ranks) == world_size + assert set(pool_a_ranks).isdisjoint(set(pool_b_ranks)) + assert len(set(pool_a_ranks)) == len(pool_a_ranks) + + all_sites = tuple(s for bg in block_groups for s in bg.owned_sites) + assert len(set(all_sites)) == len(all_sites), "a site is owned by more than one block group" + + pool_a_group = dist.new_group(ranks=pool_a_ranks) + pool_b_group = dist.new_group(ranks=pool_b_ranks) + block_group_groups = tuple(dist.new_group(ranks=list(bg.ranks)) for bg in block_groups) + cross_pool_bcast_groups = tuple( + dist.new_group(ranks=[bg.leader, *pool_b_ranks]) for bg in block_groups + ) + + return BlockDDPWorld( + world_size=world_size, + block_groups=tuple(block_groups), + pool_b_ranks=tuple(pool_b_ranks), + all_sites=all_sites, + batch_global=batch_global, + pool_a_group=pool_a_group, + pool_b_group=pool_b_group, + block_group_groups=block_group_groups, + cross_pool_bcast_groups=cross_pool_bcast_groups, + ) + + +@dataclass(frozen=True) +class BlockDDPLayout: + """This rank's view of the block-DDP world + comm + in-block all-reduce.""" + + world: BlockDDPWorld + my_rank: int + my_pool: Literal["a", "b"] + my_block_idx: int | None + my_within_block_idx: int | None + my_is_block_leader: bool + my_owned_sites: tuple[str, ...] + my_slice_idx: int | None + my_is_pool_leader: bool + + @classmethod + def from_world(cls, world: BlockDDPWorld, my_rank: int) -> "BlockDDPLayout": + for bg_idx, bg in enumerate(world.block_groups): + if my_rank in bg.ranks: + within = bg.ranks.index(my_rank) + return cls( + world=world, + my_rank=my_rank, + my_pool="a", + my_block_idx=bg_idx, + my_within_block_idx=within, + my_is_block_leader=(within == 0), + my_owned_sites=bg.owned_sites, + my_slice_idx=None, + my_is_pool_leader=False, + ) + if my_rank in world.pool_b_ranks: + return cls( + world=world, + my_rank=my_rank, + my_pool="b", + my_block_idx=None, + my_within_block_idx=None, + my_is_block_leader=False, + my_owned_sites=(), + my_slice_idx=world.pool_b_ranks.index(my_rank), + my_is_pool_leader=(my_rank == world.pool_b_ranks[0]), + ) + raise ValueError(f"rank {my_rank} not in any pool") + + def is_my_site(self, site: str) -> bool: + return self.my_pool == "a" and site in self.my_owned_sites + + def i_lead_site(self, site: str) -> bool: + return self.is_my_site(site) and self.my_is_block_leader + + def my_batch_slice_a(self) -> slice: + assert self.my_pool == "a" and self.my_within_block_idx is not None + b = self.world.batch_local_a + return slice(self.my_within_block_idx * b, (self.my_within_block_idx + 1) * b) + + def my_batch_slice_b(self) -> slice: + assert self.my_pool == "b" and self.my_slice_idx is not None + b = self.world.batch_local_b + return slice(self.my_slice_idx * b, (self.my_slice_idx + 1) * b) + + def slice_for_b_idx(self, slice_idx: int) -> slice: + b = self.world.batch_local_b + return slice(slice_idx * b, (slice_idx + 1) * b) + + # ── Cross-pool comm (block leader is the canonical actor) ── + + def send_owned_ci_to_pool_b(self, ci_owned: dict[str, Tensor]) -> None: + assert self.my_pool == "a" + for site in self.world.all_sites: + if not self.i_lead_site(site): + continue + for slice_idx, b_rank in enumerate(self.world.pool_b_ranks): + sl = self.slice_for_b_idx(slice_idx) + dist.send(ci_owned[site][sl].detach().contiguous(), dst=b_rank) + + def async_send_owned_ci_to_pool_b( + self, + ci_owned: dict[str, Tensor], + ) -> tuple[list["dist.Work"], list[Tensor]]: + """Async variant — issue isends and return (work_handles, kept-alive buffers). + Caller must wait on the handles AND keep the buffers alive until then. + Cast to bf16 on the wire (halves bytes; pool B uses bf16 autocast). + """ + assert self.my_pool == "a" + works: list[dist.Work] = [] + buffers: list[Tensor] = [] + for site in self.world.all_sites: + if not self.i_lead_site(site): + continue + for slice_idx, b_rank in enumerate(self.world.pool_b_ranks): + sl = self.slice_for_b_idx(slice_idx) + buf = ci_owned[site][sl].detach().to(_WIRE_DTYPE).contiguous() + works.append(dist.isend(buf, dst=b_rank)) + buffers.append(buf) + return works, buffers + + def recv_ci_from_owners( + self, + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> dict[str, Tensor]: + """Receives bf16-on-wire CI, upcasts to caller's requested ``dtype``.""" + assert self.my_pool == "b" + out: dict[str, Tensor] = {} + b_local = self.world.batch_local_b + for site in self.world.all_sites: + leader = self.world.block_leader_of_site(site) + buf = torch.empty(b_local, seq_len, site_to_c[site], device=device, dtype=_WIRE_DTYPE) + dist.recv(buf, src=leader) + out[site] = buf.to(dtype) + return out + + def async_recv_ci_from_owners( + self, + site_to_c: dict[str, int], + seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[dict[str, Tensor], list["dist.Work"]]: + """Async bf16-on-wire CI recv. Caller waits on the works, then casts to ``dtype``. + + We return the raw bf16 buffers; caller-side cast (after wait) lets the + pool-B step keep the work-handle / consume-buffer separation simple. + """ + assert self.my_pool == "b" + out: dict[str, Tensor] = {} + works: list[dist.Work] = [] + del dtype # unused; pool-B caller casts after wait + b_local = self.world.batch_local_b + for site in self.world.all_sites: + leader = self.world.block_leader_of_site(site) + buf = torch.empty(b_local, seq_len, site_to_c[site], device=device, dtype=_WIRE_DTYPE) + works.append(dist.irecv(buf, src=leader)) + out[site] = buf + return out, works + + def send_pool_b_grads_to_owners( + self, + v_grads: dict[str, Tensor], + u_grads: dict[str, Tensor], + ci_grads: dict[str, Tensor], + ) -> None: + """Coalesced + pipelined grad send pool B → pool A leaders. + + Per A-leader, pack all owned sites' grads into one flat tensor. + Use isend so the 24 destinations overlap on the NIC instead of + serializing — each individual send still costs the same wire time + but the per-message setup latency pipelines. + """ + assert self.my_pool == "b" + works: list[dist.Work] = [] + buffers: list[Tensor] = [] # keep alive until waits complete + if self.my_is_pool_leader: + for bg in self.world.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) + assert w is not None + works.append(w) + buffers.append(packed) + for bg in self.world.block_groups: + parts = [ + ci_grads[site].to(_WIRE_DTYPE).contiguous().flatten() for site in bg.owned_sites + ] + packed = torch.cat(parts) + w = dist.isend(packed, dst=bg.leader) + assert w is not None + works.append(w) + buffers.append(packed) + for w in works: + w.wait() + del buffers + + def recv_grads_from_pool_b( + self, + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + ci_lower_owned: dict[str, Tensor], + ) -> tuple[dict[str, Tensor], dict[str, Tensor], dict[str, Tensor]]: + """Coalesced grad recv pool A leader ← pool B. Mirrors send_pool_b_grads_to_owners.""" + assert self.my_pool == "a" + v_grads: dict[str, Tensor] = {} + u_grads: dict[str, Tensor] = {} + ci_grads: dict[str, Tensor] = {} + + b_leader = self.world.pool_b_ranks[0] + + if self.my_is_block_leader: + # ── Post all irecvs upfront so they pipeline. Receiver-side + # setup latency runs concurrently with sender-side egress. + # Recv into bf16 buffers, then upcast on unpack so grads + # accumulate into pool A's fp32 .grad. + my_sites = self.my_owned_sites + + # V/U grads: one coalesced recv from pool B leader. + vu_numel = sum(v_templates[s].numel() + u_templates[s].numel() for s in my_sites) + sample_v = v_templates[my_sites[0]] + vu_buf = torch.empty(vu_numel, dtype=_WIRE_DTYPE, device=sample_v.device) + vu_work = dist.irecv(vu_buf, src=b_leader) + assert vu_work is not None + + # ci_grads: one coalesced recv per pool B rank. + b_local = self.world.batch_local_b + per_b_numel = sum(ci_lower_owned[s][:b_local].numel() for s in my_sites) + sample_ci = ci_lower_owned[my_sites[0]] + ci_bufs: list[Tensor] = [] + ci_works: list[dist.Work] = [] + for b_rank in self.world.pool_b_ranks: + ci_buf = torch.empty(per_b_numel, dtype=_WIRE_DTYPE, device=sample_ci.device) + w = dist.irecv(ci_buf, src=b_rank) + assert w is not None + ci_bufs.append(ci_buf) + ci_works.append(w) + + # Now wait + unpack each as it arrives. Upcast on view. + vu_work.wait() + offset = 0 + for s in my_sites: + v_n = v_templates[s].numel() + u_n = u_templates[s].numel() + v_grads[s] = ( + vu_buf[offset : offset + v_n].view_as(v_templates[s]).to(v_templates[s].dtype) + ) + offset += v_n + u_grads[s] = ( + vu_buf[offset : offset + u_n].view_as(u_templates[s]).to(u_templates[s].dtype) + ) + offset += u_n + + slices_per_site: dict[str, list[Tensor]] = {s: [] for s in my_sites} + for ci_buf, w in zip(ci_bufs, ci_works, strict=True): + w.wait() + offset = 0 + for s in my_sites: + full = ci_lower_owned[s] + _, S, C = full.shape + slice_n = b_local * S * C + slices_per_site[s].append( + ci_buf[offset : offset + slice_n].view(b_local, S, C).to(full.dtype) + ) + offset += slice_n + for s in my_sites: + ci_grads[s] = torch.cat(slices_per_site[s], dim=0) + else: + # Non-leader pool A ranks pre-allocate; populated via broadcast below. + 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]) + ci_grads[s] = torch.empty_like(ci_lower_owned[s]) + + # In-block broadcast leader → other ranks in the group (unchanged). + assert self.my_block_idx is not None + block_group = self.world.block_group_groups[self.my_block_idx] + block_leader_rank = self.world.block_groups[self.my_block_idx].leader + for s in self.my_owned_sites: + v_grads[s] = v_grads[s].contiguous() + u_grads[s] = u_grads[s].contiguous() + ci_grads[s] = ci_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) + dist.broadcast(ci_grads[s], src=block_leader_rank, group=block_group) + + return v_grads, u_grads, ci_grads + + def all_reduce_grads_in_block(self, params: Iterable[nn.Parameter]) -> None: + """Coalesced in-block DDP all-reduce. + + Flatten all param grads of the same dtype + device into one big buffer, + all-reduce once, scatter back. Cuts NCCL message count from `len(params)` + (e.g. ~76 for the CI fn at 64-GPU 24×2+16) to 1 per (dtype, device) + bucket. HTA showed the per-param pattern was costing ~250ms/step. + """ + assert self.my_pool == "a" and self.my_block_idx is not None + block_group = self.world.block_group_groups[self.my_block_idx] + if dist.get_world_size(block_group) <= 1: + return # 1-rank block: in-block DDP is a no-op + + grads: list[Tensor] = [p.grad for p in params if p.grad is not None] + if not grads: + return + # Bucket by (dtype, device) — typically just one bucket in practice. + buckets: dict[tuple[torch.dtype, torch.device], list[Tensor]] = {} + for g in grads: + buckets.setdefault((g.dtype, g.device), []).append(g) + from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + + for (dtype, device), bucket in buckets.items(): + del dtype, device # used as key + flat = _flatten_dense_tensors(bucket) + dist.all_reduce(flat, op=dist.ReduceOp.AVG, group=block_group) + for orig, reduced in zip(bucket, _unflatten_dense_tensors(flat, bucket), strict=True): + orig.copy_(reduced) + + def send_updated_weights_to_pool_b( + self, + v_owned: dict[str, Tensor], + u_owned: dict[str, Tensor], + ) -> None: + assert self.my_pool == "a" + for site in self.world.all_sites: + if not self.i_lead_site(site): + continue + for b_rank in self.world.pool_b_ranks: + dist.send(v_owned[site].detach().contiguous(), dst=b_rank) + dist.send(u_owned[site].detach().contiguous(), dst=b_rank) + + def async_send_updated_weights_to_pool_b( + self, + v_owned: dict[str, Tensor], + u_owned: dict[str, Tensor], + ) -> tuple[list["dist.Work"], list[Tensor]]: + """Coalesced async weight send pool A leader → all pool B ranks. + + Pack all owned sites' V then U into one contiguous tensor; reuse the + same buffer for the isend to every pool B rank (NCCL P2P sends are + read-only, can share a source tensor). Cuts the per-leader send count + from `N_owned × 2 × N_pool_b` (e.g. 4×2×16=128) to `N_pool_b` (16). + + Caller must keep the buffer alive (returned in `buffers`) until all + work handles are done. + """ + assert self.my_pool == "a" + if not self.my_is_block_leader: + return [], [] + my_sites = self.my_owned_sites + assert self.my_block_idx is not None + 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) # owns its own memory; safe across optimizer.step + # One async broadcast to all pool-B ranks via the pre-built + # {leader} ∪ {pool_b_ranks} group. NCCL uses a tree topology so + # leader egress drops from N_pool_b× isends to 1×. + bcast_group = self.world.cross_pool_bcast_groups[self.my_block_idx] + w = dist.broadcast(packed, src=self.my_rank, group=bcast_group, async_op=True) + assert w is not None + return [w], [packed] + + def recv_updated_weights_from_owners( + self, + v_templates: dict[str, Tensor], + u_templates: dict[str, Tensor], + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + """Coalesced + pipelined weight recv pool B ← all pool A leaders. + + Post all 24 irecvs upfront so they pipeline, then wait+unpack in + completion order (skip the wait order — wait() on each in order is + fine because incomplete waits don't block already-arrived ones). + """ + assert self.my_pool == "b" + v_new: dict[str, Tensor] = {} + u_new: dict[str, Tensor] = {} + bufs: list[tuple[BlockGroup, Tensor, dist.Work]] = [] + for bg_idx, bg in enumerate(self.world.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) + # Mirror the leader's broadcast: post into the leader-rooted group + # for this block. async_op=True lets all N_blocks broadcasts run + # concurrently on independent NCCL communicators. + 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)) + for bg, packed, w in bufs: + w.wait() + owned = bg.owned_sites + 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 diff --git a/param_decomp/two_pool/loss_strategy.py b/param_decomp/two_pool/loss_strategy.py new file mode 100644 index 000000000..8133862c4 --- /dev/null +++ b/param_decomp/two_pool/loss_strategy.py @@ -0,0 +1,112 @@ +"""Layerwise loss strategy: pairs the LM-head-bypass context with recon_loss. + +Both pool A's layerwise loop and pool B's PPGD inner loop go through one of +two regimes: + + - **Fused**: bypass LM head; recon_loss is :func:`fused_linear_kl_div` against + the saved lm_head weight. Forwards return pre-LM-head hidden state; the + kernel applies LM head + KL in chunks (no vocab-scale tensor materialized). + + - **Unfused**: no bypass; the configured ``reconstruction_loss`` from PD + config is used directly. Forwards return logits. + +:class:`LayerwiseLossStrategy` encapsulates that pair so the runner never +branches on ``use_fused_kl`` — it just calls ``strategy.context`` and +``strategy.recon_loss``. +""" + +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass +from typing import Any + +import torch.nn as nn +from torch import Tensor + +from param_decomp.batch_and_loss_fns import ReconstructionLoss +from param_decomp.fused_linear_kl import fused_linear_kl_div + + +@contextmanager +def bypass_lm_head(target_model: nn.Module) -> Iterator[Any]: + """Temporarily replace target_model.lm_head with nn.Identity for one step. + + With the bypass active, every forward through ``target_model`` returns the + pre-LM-head hidden state ([b, seq, d_model]) instead of full-vocab logits + ([b, seq, vocab]). The caller must apply the LM head another way — e.g. + via :func:`fused_linear_kl_div` against the saved lm_head's weight. + + At Qwen vocab=152K, b=32+, s=1024+, the unfused path's vocab-scale tensors + cost 5-40 GB per layerwise iter. Combined with the fused kernel, this + drops peak memory to ``O(chunk_size · vocab)``. + """ + saved = target_model.lm_head + target_model.lm_head = nn.Identity() + try: + yield saved + finally: + target_model.lm_head = saved + + +def _make_fused_kl_recon_loss(lm_head_weight: Tensor) -> ReconstructionLoss: + """Build a ``ReconstructionLoss`` that applies the fused linear+KL kernel. + + Both ``pred`` and ``target`` are expected to be pre-LM-head hidden states + of shape ``[..., d_model]`` — the contract that holds when the caller is + running under :func:`bypass_lm_head`. Returns ``(sum_kl, n_positions)`` + matching :func:`recon_loss_kl`'s contract. + """ + + def fn(pred: Tensor, target: Tensor) -> tuple[Tensor, int]: + pred_flat = pred.reshape(-1, pred.shape[-1]) + target_flat = target.reshape(-1, target.shape[-1]) + return fused_linear_kl_div(pred_flat, target_flat, lm_head_weight) + + return fn + + +@dataclass(frozen=True) +class LayerwiseLossStrategy: + """Pairs the LM-head-bypass context with the matching recon_loss callable. + + The shared contract: + + 1. A context manager controlling model-forward output type for the step + (swap lm_head for Identity → forwards return hidden state; or no-op → + forwards return logits). + 2. A ``recon_loss(pred, target) -> (loss, n)`` callable whose pred/target + shapes match what (1) made the forwards produce. + """ + + context: Callable[[nn.Module], AbstractContextManager[Any]] + recon_loss: ReconstructionLoss + + @classmethod + def fused(cls, lm_head_weight: Tensor) -> "LayerwiseLossStrategy": + return cls( + context=bypass_lm_head, + recon_loss=_make_fused_kl_recon_loss(lm_head_weight), + ) + + @classmethod + def unfused(cls, recon_loss: ReconstructionLoss) -> "LayerwiseLossStrategy": + return cls( + context=lambda _model: nullcontext(), + recon_loss=recon_loss, + ) + + @classmethod + def from_cfg( + cls, + target_model: nn.Module, + use_fused_kl: bool, + unfused_recon: ReconstructionLoss, + ) -> "LayerwiseLossStrategy": + """Resolve the strategy from a (target_model, use_fused_kl) pair.""" + if use_fused_kl: + lm_head = target_model.lm_head + assert isinstance(lm_head, nn.Linear), ( + f"expected target_model.lm_head to be nn.Linear; got {type(lm_head)}" + ) + return cls.fused(lm_head.weight) + return cls.unfused(unfused_recon) diff --git a/param_decomp/two_pool/optimize.py b/param_decomp/two_pool/optimize.py new file mode 100644 index 000000000..2f647f0c3 --- /dev/null +++ b/param_decomp/two_pool/optimize.py @@ -0,0 +1,681 @@ +"""``TwoPoolTrainer`` and ``optimize_two_pool`` — 2-pool sibling of +:class:`param_decomp.optimize.Trainer` / :func:`param_decomp.optimize.optimize`. + +Mirrors the single-pool call shape: caller hands in ``target_model``, +dataloader, configs, sink. Internal validation, per-pool wiring, cross-pool +comms, and the layerwise streaming loss strategy are all hidden behind the +class boundary. + + - **Pool A** trains V/U + CI fn. Each pool-A rank holds the components for its + owned sites and runs target+CI forward, per-site streaming layerwise loss, + home losses (faithfulness, importance-minimality), combined backward seeded + by pool B's ci grads, in-block all-reduce, AdamW step. See + :mod:`param_decomp.two_pool.pool_a`. + + - **Pool B** is a stateless PPGD replica that holds full-target V/U replicas + (received from pool A each step). Each pool-B rank does target forward, + PPGD warmup + recon loss, sends V/U + CI grads back, receives updated V/U. + See :mod:`param_decomp.two_pool.pool_b`. +""" + +import itertools +import time +from contextlib import nullcontext +from typing import Any, Self + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor +from torch.utils.data import DataLoader + +from param_decomp.batch_and_loss_fns import ReconstructionLoss, RunBatch +from param_decomp.ci_fns import LayerwiseCiConfig +from param_decomp.component_model import ComponentModel +from param_decomp.configs import Cadence, PDConfig, RuntimeConfig +from param_decomp.decomposition_targets import ( + DecompositionTarget, + resolve_decomposition_targets, +) +from param_decomp.distributed import seed_all_ranks, seed_per_rank +from param_decomp.masks import AllLayersRouter +from param_decomp.metrics.base import LossMetricConfig +from param_decomp.metrics.importance_minimality import ImportanceMinimalityLossConfig +from param_decomp.metrics.persistent_pgd_recon import ( + PersistentPGDReconLossConfig, + validate_pgd_scope, +) +from param_decomp.metrics.persistent_pgd_state import PersistentPGDState +from param_decomp.run_sink import OnePoolRunSink +from param_decomp.schedule import get_scheduled_value +from param_decomp.torch_helpers import loop_dataloader +from param_decomp.training_state import TrainingState +from param_decomp.two_pool.config import TwoPoolConfig +from param_decomp.two_pool.layout import BlockDDPLayout, BlockGroup, build_block_ddp_world +from param_decomp.two_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp.two_pool.pool_a import run_faithfulness_warmup_pool_a, step_pool_a +from param_decomp.two_pool.pool_b import step_pool_b +from param_decomp.two_pool.profiler import PhaseProfiler +from param_decomp.two_pool.reductions import ( + aggregate_losses_to_rank0, + aggregate_max_memory_to_rank0, +) +from param_decomp.two_pool.runtime import _TwoPoolRuntime + +# Loss-metric type discriminators required for the 2-pool training path. +# Each MUST appear in ``pd_config.loss_metrics`` with a non-None ``coeff``. +REQUIRED_LOSS_METRIC_TYPES: frozenset[str] = frozenset( + { + "FaithfulnessLoss", + "ImportanceMinimalityLoss", + "StochasticReconLayerwiseLoss", + "PersistentPGDReconLoss", + } +) + +# Loss metrics 2-pool does not implement; configuring them would be silent. +# Listed explicitly so misconfiguration is loud. +FORBIDDEN_LOSS_METRIC_TYPES: frozenset[str] = frozenset( + { + "StochasticReconLoss", + "StochasticReconSubsetLoss", + "StochasticReconSubsetCEAndKL", + "PersistentPGDReconSubsetLoss", + "CIMaskedReconLoss", + "CIMaskedReconLayerwiseLoss", + "CIMaskedReconSubsetLoss", + "UnmaskedReconLoss", + "PGDReconLoss", + "PGDReconLayerwiseLoss", + "PGDReconSubsetLoss", + "CIMaskedAttnPatternsReconLoss", + "StochasticAttnPatternsReconLoss", + "StochasticHiddenActsReconLoss", + "CIHiddenActsReconLoss", + } +) + + +class TwoPoolTrainer: + """Stateful 2-pool trainer. + + Construction wires up the runtime bundle, world layout, ComponentModel, + layerwise loss strategy, and the per-pool training state (pool-A optimizer + or pool-B PPGD config stash). The PPGD state itself is built on the first + batch of :meth:`run` because its source tensor shapes depend on the data's + sequence dims. + + Resume support is rank-local: :meth:`state_blob` produces a self-contained + cfg + state dict for the current rank, and :meth:`from_blob` reconstructs + one from that dict. The lab's resume loader composes these across ranks. + + Consumable-checkpoint gathering is still TODO (the runtime currently + asserts ``cadence.save_every is None``); the trainer's + :meth:`consumable_model_state_dict` returns this rank's partial view for + now, matching the pre-Trainer behaviour. + """ + + pd_config: PDConfig + runtime_config: RuntimeConfig + two_pool_config: TwoPoolConfig + reconstruction_loss: ReconstructionLoss + component_model: ComponentModel + layout: BlockDDPLayout + strategy: LayerwiseLossStrategy + optimizer: torch.optim.Optimizer | None + ppgd_state: PersistentPGDState | None + step: int + + def __init__( + self, + *, + target_model: nn.Module, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, + pd_config: PDConfig, + runtime_config: RuntimeConfig, + two_pool_config: TwoPoolConfig, + cadence: Cadence, + ) -> None: + assert dist.is_initialized(), ( + "init the distributed process group before constructing TwoPoolTrainer" + ) + self.pd_config = pd_config + self.runtime_config = runtime_config + self.two_pool_config = two_pool_config + self.reconstruction_loss = reconstruction_loss + self._run_batch = run_batch + self.step = 0 + + _validate_pd_config_for_two_pool(pd_config, two_pool_config, cadence) + # PPGD runs only on pool B; the relevant per-rank batch is batch // n_pool_b. + validate_pgd_scope( + pd_config.loss_metrics, + batch_size=pd_config.batch_size, + world_size=len(two_pool_config.pool_b_ranks), + ) + + self.runtime = _build_runtime( + target_model=target_model, + pd_config=pd_config, + runtime_config=runtime_config, + two_pool_config=two_pool_config, + run_batch=run_batch, + reconstruction_loss=reconstruction_loss, + ) + + # TF32 matmuls are ~2-3x faster on H200 with sub-ULP precision loss — fine + # for SPD training where we already use fp32 throughout. + torch.set_float32_matmul_precision("high") + + self._device = torch.device(runtime_config.device) + world = build_block_ddp_world( + block_groups=list(self.runtime.block_groups), + pool_b_ranks=list(self.runtime.pool_b_ranks), + batch_global=self.runtime.batch_global, + ) + self.layout = BlockDDPLayout.from_world(world, dist.get_rank()) + decomposition_targets = _decomposition_targets_for_pool( + self.layout, self.runtime.c_per_site + ) + + target_model.requires_grad_(False) + # Resync RNG across ranks before V/U + CI fn init. DDP partners within + # a block must start with identical params, but anything between + # ``set_seed(pd.seed)`` in _fresh_main and here (loader build, distributed + # init, etc.) can advance the RNG by rank-dependent amounts. Without + # this, partners initialize different V/U and the in-block grad + # all-reduce can't bring them back into sync. + seed_all_ranks(pd_config.seed) + self.component_model = ComponentModel( + target_model=target_model, + run_batch=run_batch, + decomposition_targets=decomposition_targets, + ci_config=pd_config.ci_config, + sigmoid_type=pd_config.sigmoid_type, + ).to(self._device) + # Diverge stochastic RNG per rank again now that init is done (matches + # single-pool's order — see param_decomp.optimize.Trainer.__init__). + seed_per_rank(pd_config.seed) + + self.strategy = LayerwiseLossStrategy.from_cfg( + target_model, + use_fused_kl=two_pool_config.use_fused_kl, + unfused_recon=reconstruction_loss, + ) + + self.optimizer = None + self._all_params: list[nn.Parameter] = [] + self._component_params: list[nn.Parameter] = [] + self._ci_fn_params: list[nn.Parameter] = [] + self.ppgd_state = None + + if self.layout.my_pool == "a": + for name in self.component_model.target_module_paths: + self._component_params.extend(self.component_model.components[name].parameters()) + assert self.component_model.ci_fn is not None, "2-pool 'a' pool must keep its CI fn" + self._ci_fn_params = list(self.component_model.ci_fn.parameters()) + self._all_params = self._component_params + self._ci_fn_params + self.optimizer = torch.optim.AdamW( + [ + { + "params": self._component_params, + "lr": pd_config.components_optimizer.lr_schedule.start_val, + }, + { + "params": self._ci_fn_params, + "lr": pd_config.ci_fn_optimizer.lr_schedule.start_val, + }, + ], + weight_decay=0.0, + fused=True, + ) + # Pool B's PPGD state is constructed lazily in `run` once batch_dims is known + # from the first batch; pending resume state (if any) is applied at that point. + self._pending_ppgd_resume_state: dict[str, Any] | None = None + + # ============================ Atomic cfg + state ============================ + + def snapshot(self) -> TrainingState: + """2-pool snapshot not ported to the concept-level resumption design yet. + + See ``ThreePoolTrainer.snapshot`` for the canonical pattern; the + equivalent for 2-pool is a straightforward port but not in scope here. + """ + raise NotImplementedError( + "TwoPoolTrainer.snapshot not ported to the concept-level resumption design." + ) + + @classmethod + def from_snapshot( + cls, + snapshot: TrainingState, + *, + target_model: nn.Module, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, + cadence: Cadence, + cfg_overrides: dict[str, Any] | None = None, + ) -> Self: + del snapshot, target_model, run_batch, reconstruction_loss, cadence, cfg_overrides + raise NotImplementedError( + "TwoPoolTrainer.from_snapshot not ported to the concept-level resumption design." + ) + + # ============================ Training loop ============================ + + def run( + self, + train_loader: DataLoader[Any], + sink: OnePoolRunSink, + cadence: Cadence, + profiler: PhaseProfiler | None = None, + ) -> None: + """Advance training from ``self.step`` to ``self.pd_config.steps``.""" + pd_config = self.pd_config + layout = self.layout + runtime = self.runtime + + train_iterator = loop_dataloader(train_loader) + + # Loader skip-replay (resumed mid-trajectory) and first-batch peek for + # pool-B PPGD shape — done in one pass. + for _ in range(self.step): + next(train_iterator) + first_batch = next(train_iterator) + train_iterator = itertools.chain([first_batch], train_iterator) + + if layout.my_pool == "b" and self.ppgd_state is None: + ppgd_cfg = runtime.ppgd_cfg + self.ppgd_state = PersistentPGDState( + module_to_c=runtime.c_per_site, + batch_dims=(layout.world.batch_local_b, *_seq_dims_from_batch(first_batch)), + device=self._device, + use_delta_component=True, + optimizer_cfg=ppgd_cfg.optimizer, + scope=ppgd_cfg.scope, + use_sigmoid_parameterization=ppgd_cfg.use_sigmoid_parameterization, + n_warmup_steps=ppgd_cfg.n_warmup_steps, + n_samples=ppgd_cfg.n_samples, + 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.step == 0 and layout.my_pool == "a" and pd_config.faithfulness_warmup_steps > 0: + run_faithfulness_warmup_pool_a( + component_model=self.component_model, + component_params=self._component_params, + n_steps=pd_config.faithfulness_warmup_steps, + lr=pd_config.faithfulness_warmup_lr, + weight_decay=pd_config.faithfulness_warmup_weight_decay, + numel_global=self.runtime.numel_global, + ) + + n_steps = pd_config.steps + profiler_ctx = profiler if profiler is not None else nullcontext() + with profiler_ctx: + for step in range(self.step, n_steps): + self.step = step + if layout.my_pool == "a" and self.optimizer is not None: + self.optimizer.param_groups[0]["lr"] = get_scheduled_value( + step, n_steps, pd_config.components_optimizer.lr_schedule + ) + self.optimizer.param_groups[1]["lr"] = get_scheduled_value( + step, n_steps, pd_config.ci_fn_optimizer.lr_schedule + ) + + # When profiling, barrier ranks at step boundary so both pools share + # a common time origin in the trace. + if profiler is not None: + dist.barrier() + + batch = _extract_batch_tensor(next(train_iterator), self._device) + assert batch.device == self._device, ( + f"2-pool batch device mismatch at step {step}: {batch.device} vs {self._device}" + ) + assert batch.shape[0] == runtime.batch_global, ( + f"2-pool batch_global mismatch at step {step}: " + f"got {batch.shape[0]}, expected {runtime.batch_global}" + ) + + torch.cuda.synchronize(self._device) + step_start = time.perf_counter() + match layout.my_pool: + case "a": + assert self.optimizer is not None + assert layout.my_owned_sites, ( + f"pool-A rank {layout.my_rank} has no owned_sites — empty block" + ) + metrics = step_pool_a( + layout, + self.component_model, + self.optimizer, + self._all_params, + self._component_params, + self._ci_fn_params, + batch, + runtime, + self.strategy, + current_frac_of_training=step / n_steps if n_steps > 0 else 0.0, + profiler=profiler, + ) + case "b": + assert self.ppgd_state is not None, ( + f"pool-B rank {layout.my_rank} has no ppgd_state — lazy init failed" + ) + metrics = step_pool_b( + layout, + self.component_model, + self.ppgd_state, + batch, + runtime, + self.strategy, + step=step, + n_steps=n_steps, + profiler=profiler, + ) + # Loss values produced by either pool should be finite — non-finite means a + # silent NaN somewhere upstream (PPGD update blew up, autocast overflow, etc.). + # Covers both the per-rank display scalars (``loss/*``) and the raw + # aggregation ingredients (``_raw/*``) that the logger sums across blocks. + for k, v in metrics.items(): + if k.startswith("loss/") or k.startswith("_raw/"): + assert v == v, f"NaN in metrics[{k!r}] at step {step}" # NaN != NaN + torch.cuda.synchronize(self._device) + step_ms = (time.perf_counter() - step_start) * 1000.0 + + if step % cadence.train_log_every == 0: + _log_train_metrics( + metrics=metrics, + layout=layout, + device=self._device, + step=step, + step_ms=step_ms, + runtime=runtime, + optimizer=self.optimizer, + sink=sink, + ) + + if cadence.save_every is not None and step > 0 and step % cadence.save_every == 0: + # All ranks call; the sink (resume-aware variant) decides per-rank + # behaviour. The default rank-0-only `RunSink` is a no-op elsewhere. + sink.checkpoint(self.snapshot()) + + if profiler is not None: + profiler.step() + + +def optimize_two_pool( + target_model: nn.Module, + train_loader: DataLoader[Any], + *, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, + pd_config: PDConfig, + runtime_config: RuntimeConfig, + two_pool_config: TwoPoolConfig, + cadence: Cadence, + sink: OnePoolRunSink, + profiler: PhaseProfiler | None = None, +) -> None: + """Train a ComponentModel under the 2-pool strategy. + + Thin wrapper over :class:`TwoPoolTrainer` for callers that don't need + resumption or interactive control. Identical call shape to the pre-Trainer + function. + """ + trainer = TwoPoolTrainer( + target_model=target_model, + run_batch=run_batch, + reconstruction_loss=reconstruction_loss, + pd_config=pd_config, + runtime_config=runtime_config, + two_pool_config=two_pool_config, + cadence=cadence, + ) + trainer.run(train_loader, sink, cadence, profiler=profiler) + + +def _validate_pd_config_for_two_pool( + pd_config: PDConfig, two_pool_config: TwoPoolConfig, cadence: Cadence +) -> None: + """Fail loudly on any PDConfig the 2-pool path can't honour.""" + by_type: dict[str, LossMetricConfig] = {m.type: m for m in pd_config.loss_metrics} + + missing = sorted(REQUIRED_LOSS_METRIC_TYPES - set(by_type)) + assert not missing, ( + f"2-pool requires these loss metrics: {sorted(REQUIRED_LOSS_METRIC_TYPES)}.\n" + f"Missing: {missing}. Got: {sorted(by_type)}." + ) + illegal = sorted(FORBIDDEN_LOSS_METRIC_TYPES & set(by_type)) + assert not illegal, ( + f"2-pool does not implement these loss metrics (they would be silently ignored): " + f"{illegal}. Remove them or extend the 2-pool path." + ) + + for name in REQUIRED_LOSS_METRIC_TYPES: + assert by_type[name].coeff is not None, ( + f"pd_config.loss_metrics[{name!r}].coeff is required for 2-pool training" + ) + + n_per_block = len(two_pool_config.block_groups[0].ranks) + n_pool_b = len(two_pool_config.pool_b_ranks) + bs = pd_config.batch_size + assert bs % n_per_block == 0, ( + f"pd_config.batch_size ({bs}) must be divisible by N_per_block ({n_per_block}) " + f"= len(block_groups[0].ranks)" + ) + assert bs % n_pool_b == 0, ( + f"pd_config.batch_size ({bs}) must be divisible by N_pool_b ({n_pool_b}) " + f"= len(pool_b_ranks)" + ) + + assert pd_config.use_delta_component, ( + "2-pool requires pd_config.use_delta_component=True (hardcoded in pool A's " + "layerwise stoch recon + pool B's PPGD)." + ) + + # Global CI fns can't span all sites under 2-pool because pool A shards + # sites across ranks — each rank only sees its owned sites' activations. + # Use `mode: layerwise` (with `fn_type: transformer` for the transformer variant). + assert isinstance(pd_config.ci_config, LayerwiseCiConfig), ( + "2-pool requires `pd.ci_config.mode == 'layerwise'`. Global CI fns can't " + "span all sites under 2-pool (sites are sharded across pool-A ranks). " + f"Got mode={pd_config.ci_config.mode!r}." + ) + + # Pool A's layerwise streaming loop hardcodes continuous sampling and one + # mask per site; let the user know if their YAML disagrees. + assert pd_config.sampling == "continuous", ( + "2-pool hardcodes `sampling='continuous'` in pool A's CI computation; " + f"got pd_config.sampling={pd_config.sampling!r}." + ) + assert pd_config.n_mask_samples == 1, ( + "2-pool draws exactly one stochastic mask per site per step; " + f"got pd_config.n_mask_samples={pd_config.n_mask_samples}." + ) + + # insert_identity_operations_ isn't called from the 2-pool path, so any + # identity-layer decomposition target would be silently ignored. + assert pd_config.identity_decomposition_targets is None, ( + "2-pool path does not call `insert_identity_operations_`; " + "`identity_decomposition_targets` would be silently ignored." + ) + + # No distributed-aware consumable checkpoint gather yet — rank 0 only holds + # its block's V/U + CI fn, so saving would produce a structurally partial + # file. Resume shards (per-rank) work fine; the assertion guards the + # consumable-model path until the gather is wired up. + assert cadence.save_every is None, ( + "2-pool does not yet implement distributed consumable checkpointing. " + "rank 0 only holds its own block's V/U + CI fn, so `cadence.save_every` " + "would produce a partial checkpoint that can't be reloaded. " + "Set `cadence.save_every: null` for now." + ) + + ppgd_cfg = by_type["PersistentPGDReconLoss"] + assert isinstance(ppgd_cfg, PersistentPGDReconLossConfig) + assert ppgd_cfg.start_frac == 0.0, ( + "2-pool path does not implement PersistentPGDReconLoss.start_frac > 0; " + "PPGD always runs from step 0." + ) + + +def _build_runtime( + target_model: nn.Module, + pd_config: PDConfig, + runtime_config: RuntimeConfig, + two_pool_config: TwoPoolConfig, + run_batch: RunBatch, + reconstruction_loss: ReconstructionLoss, +) -> _TwoPoolRuntime: + """Assemble the step-context bundle from configs + target.""" + targets = resolve_decomposition_targets(target_model, pd_config.decomposition_targets) + c_per_site = {t.module_path: t.C for t in targets} + # Total numel across all decomposition sites' weight tensors. Used by + # ``_faithfulness_loss`` so multi-pool's per-element grad matches single- + # pool's (which divides faith by the global numel, not the rank-local one). + numel_global = 0 + for t in targets: + w = target_model.get_submodule(t.module_path).weight + assert isinstance(w, Tensor) + numel_global += w.numel() + + for bg in two_pool_config.block_groups: + for site in bg.owned_sites: + assert site in c_per_site, ( + f"site '{site}' in two-pool topology but not in pd_config.decomposition_targets " + f"after pattern expansion. Available: {sorted(c_per_site)[:5]}..." + ) + + by_type: dict[str, LossMetricConfig] = {m.type: m for m in pd_config.loss_metrics} + ppgd_cfg = by_type["PersistentPGDReconLoss"] + imp_min_cfg = by_type["ImportanceMinimalityLoss"] + assert isinstance(ppgd_cfg, PersistentPGDReconLossConfig) + assert isinstance(imp_min_cfg, ImportanceMinimalityLossConfig) + + def _coeff(name: str) -> float: + c = by_type[name].coeff + assert c is not None # validated above + return float(c) + + block_groups = tuple( + BlockGroup(ranks=tuple(bg.ranks), owned_sites=tuple(bg.owned_sites)) + for bg in two_pool_config.block_groups + ) + + return _TwoPoolRuntime( + block_groups=block_groups, + pool_b_ranks=tuple(two_pool_config.pool_b_ranks), + batch_global=pd_config.batch_size, + c_per_site=c_per_site, + ci_config=pd_config.ci_config, + sigmoid_type=pd_config.sigmoid_type, + run_batch=run_batch, + reconstruction_loss=reconstruction_loss, + ppgd_cfg=ppgd_cfg, + coeff_faith=_coeff("FaithfulnessLoss"), + coeff_imp=_coeff("ImportanceMinimalityLoss"), + coeff_stoch=_coeff("StochasticReconLayerwiseLoss"), + coeff_ppgd=_coeff("PersistentPGDReconLoss"), + imp_min_pnorm=imp_min_cfg.pnorm, + imp_min_beta=imp_min_cfg.beta, + imp_min_eps=imp_min_cfg.eps, + imp_min_p_anneal_start_frac=imp_min_cfg.p_anneal_start_frac, + imp_min_p_anneal_final_p=imp_min_cfg.p_anneal_final_p, + imp_min_p_anneal_end_frac=imp_min_cfg.p_anneal_end_frac, + lr_components=pd_config.components_optimizer.lr_schedule.start_val, + lr_ci_fn=pd_config.ci_fn_optimizer.lr_schedule.start_val, + grad_clip_norm_components=pd_config.components_optimizer.grad_clip_norm, + grad_clip_norm_ci_fn=pd_config.ci_fn_optimizer.grad_clip_norm, + numel_global=numel_global, + bf16_autocast=runtime_config.autocast_bf16, + use_fused_kl=two_pool_config.use_fused_kl, + ) + + +def _decomposition_targets_for_pool( + layout: BlockDDPLayout, c_per_site: dict[str, int] +) -> list[DecompositionTarget]: + """Pool A: only this rank's owned sites. Pool B: every site (replicated V/U).""" + sites = layout.my_owned_sites if layout.my_pool == "a" else layout.world.all_sites + return [DecompositionTarget(module_path=s, C=c_per_site[s]) for s in sites] + + +def _extract_batch_tensor(batch: Any, device: str | torch.device) -> Tensor: + """Pull a single input tensor out of a dataloader yield and move to device. + + Supports the three common shapes: a bare Tensor, a dict with ``input_ids``, + or a tuple/list whose first element is a Tensor. + """ + if isinstance(batch, Tensor): + return batch.to(device) + if isinstance(batch, dict) and "input_ids" in batch: + return batch["input_ids"].to(device) + if isinstance(batch, list | tuple) and len(batch) > 0 and isinstance(batch[0], Tensor): + return batch[0].to(device) + raise TypeError(f"Unsupported batch type from DataLoader: {type(batch).__name__}") + + +def _seq_dims_from_batch(batch: Any) -> tuple[int, ...]: + """Sequence dims (everything past the leading batch dim) of a sample batch.""" + if isinstance(batch, Tensor): + return tuple(batch.shape[1:]) + if isinstance(batch, dict) and "input_ids" in batch: + return tuple(batch["input_ids"].shape[1:]) + if isinstance(batch, list | tuple) and len(batch) > 0 and isinstance(batch[0], Tensor): + return tuple(batch[0].shape[1:]) + raise TypeError(f"Cannot infer seq dims from batch of type {type(batch).__name__}") + + +def _log_train_metrics( + *, + metrics: dict[str, float], + layout: BlockDDPLayout, + device: torch.device, + step: int, + step_ms: float, + runtime: _TwoPoolRuntime, + optimizer: torch.optim.Optimizer | None, + sink: OnePoolRunSink, +) -> None: + """Aggregate per-pool metrics to rank 0 and dispatch to ``sink``.""" + combined = aggregate_losses_to_rank0(metrics, layout, device) + mem_combined = aggregate_max_memory_to_rank0(layout, device) + + # Reduce step_ms with MAX across pool A (slowest pool A rank is the wall-clock + # floor; pool B should track it). + step_ms_t = torch.tensor([step_ms], device=device) + if layout.my_pool == "a": + dist.all_reduce(step_ms_t, op=dist.ReduceOp.MAX, group=layout.world.pool_a_group) + + if layout.my_rank != 0 or combined is None: + return + + if mem_combined is not None: + combined.update(mem_combined) + combined["perf/step_ms"] = step_ms_t.item() + # grad_norm/* on rank 0 holds the pre-clip global L2 norm (computed by + # cross_pool_clip_grad_norm via an all-reduce SUM across pool A, so it's + # identical on every pool-A rank — rank 0's value is the right one to log). + for k in ("grad_norm/components", "grad_norm/ci_fn"): + if k in metrics: + combined[k] = metrics[k] + combined["loss/total"] = ( + runtime.coeff_faith * combined["loss/faith"] + + runtime.coeff_imp * combined["loss/imp"] + + runtime.coeff_stoch * combined["loss/stoch"] + + runtime.coeff_ppgd * combined["loss/ppgd"] + ) + assert layout.my_pool == "a", "rank 0 must be in pool A" + assert optimizer is not None + combined["schedules/lr/components"] = optimizer.param_groups[0]["lr"] + combined["schedules/lr/ci_fn"] = optimizer.param_groups[1]["lr"] + sink.console( + f"--- Step {step} ---", + *(f"train/{name}: {value:.6g}" for name, value in combined.items()), + ) + sink.log({f"train/{k}": v for k, v in combined.items()}, step=step) diff --git a/param_decomp/two_pool/pool_a.py b/param_decomp/two_pool/pool_a.py new file mode 100644 index 000000000..b5e912a6d --- /dev/null +++ b/param_decomp/two_pool/pool_a.py @@ -0,0 +1,655 @@ +"""Pool-A training step: target+CI forward, layerwise streaming loss, home losses, opt step. + +Pool A trains V/U + CI fn. The numbered phases are: + + 0. Wait for any pending async weight ship from previous step. + 1. Target + CI fn forward — CI graph retained for the combined backward. + 2. Async-send CI values to pool B (overlaps with home-loss compute). + 3. Faithfulness loss (forward only; backward in phase 7). + 4. Importance-minimality loss (forward only; backward in phase 7). + 5. Streaming layerwise stoch recon — at the step level so the per-site loop + is visible. One site per iter: build mask, forward through model, recon + loss, immediate ``.backward()``. Peak memory bounded to ~1× iter + activations. + 6. Recv per-site V/U grads + per-slice CI grads from pool B. + 7. Combined backward: home (faith+imp) + CI-fn graph seeded by combined + (pool-B + layerwise) CI grads. Adds pool B's V/U grads to layerwise's. + 8. In-block DDP all-reduce on grads. + 8b. Cross-pool grad clip on the global norm (matches single-pool's + ``clip_grad_norm_``). + 9. AdamW step. + 10. Async-ship updated V/U to pool B for next step. + 11. Wait on the CI send from phase 2. + +``step_pool_a`` reads as ~40 lines of orchestration. The per-site layerwise +loop stays at the top level — that loop IS the structure of the step. The +named helpers below correspond to the other phases. Profiler tags +(``a/_``) are preserved. +""" + +# pyright: reportArgumentType=false, reportOperatorIssue=false, reportAttributeAccessIssue=false + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +from torch import Tensor + +from param_decomp.component_model import ComponentModel +from param_decomp.grad_clip import cross_pool_clip_grad_norm +from param_decomp.masks import make_mask_infos +from param_decomp.metrics.importance_minimality import ( + _finalize as _finalize_imp_min, +) +from param_decomp.metrics.importance_minimality import ( + _get_linear_annealed_p, + _per_component_sums, +) +from param_decomp.two_pool.layout import BlockDDPLayout +from param_decomp.two_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp.two_pool.profiler import PhaseProfiler +from param_decomp.two_pool.runtime import _TwoPoolRuntime, autocast_bf16 + + +def _faithfulness_loss( + component_model: ComponentModel, device: torch.device, numel_global: int +) -> tuple[Tensor, Tensor, int]: + """Standard faithfulness loss: ``‖W_target − VU.T‖²_F / numel_global`` over this + rank's owned sites. + + Single-pool computes ``sum_sq_global / numel_global``; the per-element + gradient on each ``V_i`` / ``U_i`` is then ``∝ 1 / numel_global``. To match + that gradient scale in multi-pool we still divide by ``numel_global`` (not + rank-local ``numel_owned``) — otherwise per-element gradients scale up by + ``numel_global / numel_owned`` and the trajectory diverges from single-pool + (most visibly during the unclipped faithfulness warmup). + + Returns ``(scalar_loss, sum_sq, numel_owned)``. The raw ``numel_owned`` is + what the logger needs as the denominator for the + ``SUM(num) / SUM(den)`` global-ratio reconstruction across blocks. + """ + weight_deltas = component_model.calc_weight_deltas() + sum_sq = torch.zeros((), device=device) + numel_owned = 0 + for d in weight_deltas.values(): + sum_sq = sum_sq + (d**2).sum() + numel_owned += d.numel() + assert sum_sq.dim() == 0, f"sum_sq should be scalar; got shape {sum_sq.shape}" + assert numel_owned > 0, "rank must own at least one site's params" + return sum_sq / numel_global, sum_sq, numel_owned + + +def _importance_minimality_loss( + ci_upper: dict[str, Tensor], + current_frac_of_training: float, + cfg: _TwoPoolRuntime, +) -> Tensor: + """Importance-minimality loss matching single-pool semantics. + + ``world_size=1`` because the pool-A target+CI forward (phase a/1) runs on + the FULL global batch (only the layerwise stoch loss slices the batch — + NOT the CI fn forward). Each rank's local ``sum`` over its owned sites + already equals the global sum for those sites; cross-block aggregation is + SUM-across-disjoint-site-sets and lives entirely in the logger. + """ + annealed_p = _get_linear_annealed_p( + current_frac_of_training=current_frac_of_training, + initial_p=cfg.imp_min_pnorm, + p_anneal_start_frac=cfg.imp_min_p_anneal_start_frac, + p_anneal_final_p=cfg.imp_min_p_anneal_final_p, + p_anneal_end_frac=cfg.imp_min_p_anneal_end_frac, + ) + per_component_sums, n_examples = _per_component_sums( + ci_upper_leaky=ci_upper, pnorm=annealed_p, eps=cfg.imp_min_eps + ) + return _finalize_imp_min( + per_component_sums=per_component_sums, + n_examples=n_examples, + beta=cfg.imp_min_beta, + world_size=1, + ) + + +def step_pool_a( + layout: BlockDDPLayout, + component_model: ComponentModel, + optimizer: torch.optim.Optimizer, + all_params: list[nn.Parameter], + component_params: list[nn.Parameter], + ci_fn_params: list[nn.Parameter], + batch: Any, + cfg: _TwoPoolRuntime, + strategy: LayerwiseLossStrategy, + current_frac_of_training: float, + profiler: PhaseProfiler | None = None, +) -> dict[str, float]: + """One training step on a pool-A rank. + + Reads as orchestration: each phase delegates to a helper, except the + semantically central per-site layerwise loop which stays inline so the + step's structure is immediately visible. + """ + p = profiler if profiler is not None else PhaseProfiler(enabled=False) + n_sites_total = len(cfg.c_per_site) + + _wait_pending_weight_send(component_model, p) + + # ``strategy.context`` chooses logits-vs-hidden output of the model + # forwards for this step; ``recon_loss`` matches the choice. + with strategy.context(component_model.target_model): + fwd = _target_and_ci_forward(component_model, batch, cfg, layout, p) + ci_send = _async_send_owned_ci_to_pool_b(layout, fwd.ci, p) + home = _home_losses_forward_only(component_model, fwd, cfg, current_frac_of_training, p) + optimizer.zero_grad(set_to_none=True) + + # ── Phase a/5: per-site streaming layerwise loop ── + # The CI-fn graph is retained from phase a/1; we re-leaf the sliced + # CI values so each per-site backward stops at the leaf rather than + # traversing the CI-fn graph (the combined backward in phase a/7 + # does that traversal once with the merged seed). + with p.phase("a/5_layerwise"), autocast_bf16(cfg.bf16_autocast): + sl = layout.my_batch_slice_a() + batch_local = batch[sl] if isinstance(batch, Tensor) else batch + target_local = fwd.target_out[sl].detach() + ci_lower_leaves = _make_sliced_ci_leaves(fwd.ci, sl, layout.my_owned_sites) + stoch_total = 0.0 + for site in layout.my_owned_sites: + pred = _masked_forward_one_site(component_model, batch_local, ci_lower_leaves, site) + loss, n_positions = strategy.recon_loss(pred=pred, target=target_local) + assert loss.dim() == 0, f"recon_loss should return scalar; got {loss.shape}" + # Per-site contribution: ``coeff_stoch * sum_kl_s / (n_pos * + # n_sites_total)``. Backward per iter so this site's autograd + # graph is freed before the next forward. + (cfg.coeff_stoch * loss / (n_positions * n_sites_total)).backward() + stoch_total += (loss / n_positions).item() + stoch_n_owned = len(layout.my_owned_sites) + loss_stoch_scalar = stoch_total / stoch_n_owned + + pool_b_grads = _recv_grads_from_pool_b(layout, component_model, fwd.ci, p) + _combined_backward( + component_model, + fwd.ci, + ci_lower_leaves, + sl, + pool_b_grads, + home, + layout, + cfg, + p, + ) + _in_block_all_reduce_grads(layout, all_params, p) + clip_norms = _cross_pool_grad_clip(component_params, ci_fn_params, layout, cfg, p) + _optimizer_step(optimizer, p) + weight_send = _async_send_updated_vu_to_pool_b(component_model, layout, p) + _wait_async_ci_send(ci_send, p) + _stash_pending_weight_send(component_model, weight_send) + + return _step_metrics(home, loss_stoch_scalar, stoch_total, stoch_n_owned, clip_norms) + + +# ============================================================================= +# Local data bundles for threading state through phases without dropping the +# per-step orchestration function into a wall of positional args. +# ============================================================================= + + +@dataclass +class _ForwardOutputs: + """Output of phase a/1's combined target + CI fn forward.""" + + target_out: Tensor # logits or pre-LM-head hidden, per strategy + ci: Any # CIOutputs — contains upper_leaky + lower_leaky dicts + + +@dataclass +class _HomeLossOutputs: + """Outputs of phases a/3 + a/4 (forward only — backward is in a/7).""" + + loss_faith: Tensor + faith_sum_sq: Tensor + faith_numel: int + loss_imp: Tensor + + +@dataclass +class _PoolBGrads: + v_grads: dict[str, Tensor] + u_grads: dict[str, Tensor] + ci_grads: dict[str, Tensor] + + +@dataclass +class _AsyncWorkHandle: + """Generic (works, buffers) bundle for an async cross-pool send.""" + + works: list[Any] + buffers: list[Any] + + +@dataclass +class _ClipNorms: + components: float + ci_fn: float + + +# ============================================================================= +# Phase helpers (other than the inline layerwise loop above). +# ============================================================================= + + +def _wait_pending_weight_send(component_model: ComponentModel, p: PhaseProfiler) -> None: + """Phase a/0. Wait for any pending async V/U send from the previous step. + + The previous step kicks off an async send in phase a/10 and stashes the + handle on the model. We must wait for it before mutating V/U via the + optimizer step — otherwise the wire reads inconsistent data. + """ + with p.phase("a/0_wait_prev_weight_send"): + 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] + + +def _target_and_ci_forward( + component_model: ComponentModel, + batch: Any, + cfg: _TwoPoolRuntime, + layout: BlockDDPLayout, + p: PhaseProfiler, +) -> _ForwardOutputs: + """Phase a/1. Target forward (cached) + CI fn forward. + + Inside ``bf16_autocast``: target_fwd's SDPA only has a fast kernel for + bf16/fp16 on H200; fp32 falls back to O(N²) math backend. Faithfulness + loss is computed OUTSIDE autocast (see ``_faithfulness_loss``). + + CI fn graph is retained — phase a/7 will backward through it. + """ + with p.phase("a/1_target_and_ci_fwd"), autocast_bf16(cfg.bf16_autocast): + out = component_model(batch, cache_type="input") + target_out = out.output + ci = component_model.calc_causal_importances( + pre_weight_acts=out.cache, + sampling="continuous", + detach_inputs=False, + ) + _assert_ci_shapes(ci, layout, cfg) + _maybe_log_dtype_sanity(component_model, layout, target_out, out.cache, ci, cfg) + return _ForwardOutputs(target_out=target_out, ci=ci) + + +def _assert_ci_shapes(ci: Any, layout: BlockDDPLayout, cfg: _TwoPoolRuntime) -> None: + """CI lower_leaky/upper_leaky are dict[site → [batch, seq, C_s]] full-batch.""" + batch_global = cfg.batch_global + for s in layout.my_owned_sites: + assert s in ci.lower_leaky, f"CI missing site {s!r}" + t = ci.lower_leaky[s] + assert t.ndim == 3, f"ci.lower_leaky[{s!r}] expected 3D [B,S,C], got {t.shape}" + assert t.shape[0] == batch_global, ( + f"ci.lower_leaky[{s!r}] batch dim {t.shape[0]} != batch_global {batch_global}" + ) + assert t.shape[-1] == cfg.c_per_site[s], ( + f"ci.lower_leaky[{s!r}] C dim {t.shape[-1]} != c_per_site {cfg.c_per_site[s]}" + ) + + +def _maybe_log_dtype_sanity( + component_model: ComponentModel, + layout: BlockDDPLayout, + target_out: Tensor, + cache: dict[str, Tensor], + ci: Any, + cfg: _TwoPoolRuntime, +) -> None: + """One-time dtype sanity print on step 0 from rank 0 only.""" + if getattr(component_model, "_dtype_logged", False) or layout.my_rank != 0: + return + sample_act = next(iter(cache.values())) + print( + f"[two_pool sanity rank0] bf16_autocast={cfg.bf16_autocast} " + f"use_fused_kl={cfg.use_fused_kl} " + f"target_out.dtype={target_out.dtype} " + f"target_out.shape={tuple(target_out.shape)} " + f"cached_act.dtype={sample_act.dtype} " + f"ci.lower_leaky.dtype={next(iter(ci.lower_leaky.values())).dtype}", + flush=True, + ) + component_model._dtype_logged = True # type: ignore[attr-defined] + + +def _async_send_owned_ci_to_pool_b( + layout: BlockDDPLayout, + ci: Any, + p: PhaseProfiler, +) -> _AsyncWorkHandle: + """Phase a/2. Async-send owned CI values to pool B. + + Don't block here — pool B is doing its own target_fwd in parallel and + will pull these values when it needs them. We wait at the end of the + step (phase a/11) before the next iter mutates the underlying tensors. + """ + with p.phase("a/2_async_send_ci"): + works, buffers = layout.async_send_owned_ci_to_pool_b( + {s: ci.lower_leaky[s] for s in layout.my_owned_sites} + ) + return _AsyncWorkHandle(works=works, buffers=buffers) + + +def _home_losses_forward_only( + component_model: ComponentModel, + fwd: _ForwardOutputs, + cfg: _TwoPoolRuntime, + current_frac_of_training: float, + p: PhaseProfiler, +) -> _HomeLossOutputs: + """Phases a/3 + a/4. Compute faith and imp losses (forward only). + + Backward folds into the combined backward (a/7) — we want a single fused + backward through faith + imp + CI-fn-seeded-from-layerwise-and-pool-B. + """ + device = fwd.target_out.device + with p.phase("a/3_faith"): + loss_faith, faith_sum_sq, faith_numel = _faithfulness_loss( + component_model, device, cfg.numel_global + ) + with p.phase("a/4_imp"): + loss_imp = _importance_minimality_loss(fwd.ci.upper_leaky, current_frac_of_training, cfg) + assert loss_faith.dim() == 0 and loss_imp.dim() == 0, ( + f"home losses should be scalars; got faith={loss_faith.shape}, imp={loss_imp.shape}" + ) + return _HomeLossOutputs( + loss_faith=loss_faith, + faith_sum_sq=faith_sum_sq, + faith_numel=faith_numel, + loss_imp=loss_imp, + ) + + +def _make_sliced_ci_leaves( + ci: Any, + sl: slice, + owned_sites: tuple[str, ...], +) -> dict[str, Tensor]: + """Build detached, requires-grad leaves of CI at this rank's batch slice. + + Layerwise's per-iter backward stops at these leaves rather than traversing + the (retained) CI-fn graph; the combined backward in phase a/7 picks up + the leaf .grads and merges them with pool B's CI grads to seed a single + fused backward through the CI fn. + """ + return {s: ci.lower_leaky[s][sl].detach().requires_grad_(True) for s in owned_sites} + + +def _masked_forward_one_site( + component_model: ComponentModel, + batch_local: Any, + ci_lower_leaves: dict[str, Tensor], + site: str, +) -> Tensor: + """One layerwise iter's forward: build mask from CI leaf, run component + model with that single-site mask active, return pred. + + The mask is ``ci + (1 - ci) * uniform`` — single-pool layerwise's standard + stochastic mask. ``delta_mask`` is a per-site random gate on the + ``target - V@U.T`` delta component. + """ + ci_s = ci_lower_leaves[site] + u = torch.rand_like(ci_s) + mask = ci_s + (1 - ci_s) * u + delta = component_model.target_weight(site) - component_model.components[site].weight + delta_mask = torch.rand(ci_s.shape[:-1], device=ci_s.device, dtype=ci_s.dtype) + mask_infos = make_mask_infos( + {site: mask}, + weight_deltas_and_masks={site: (delta, delta_mask)}, + routing_masks="all", + ) + return component_model(batch_local, mask_infos=mask_infos) + + +def _recv_grads_from_pool_b( + layout: BlockDDPLayout, + component_model: ComponentModel, + ci: Any, + p: PhaseProfiler, +) -> _PoolBGrads: + """Phase a/6. Recv per-site V/U grads + per-slice CI grads from pool B. + + Pool B already SUM-reduced V/U grads internally and scaled by + ``1/n_pool_b``, so the values arriving are the full-batch V/U + contribution. ``ci_grads`` are full-batch CI tensors (pool B sees all + positions for the rank's owned sites). + """ + with p.phase("a/6_recv_grads_from_b"): + 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} + ci_lower_owned_full = {s: ci.lower_leaky[s] for s in layout.my_owned_sites} + v_grads, u_grads, ci_grads = layout.recv_grads_from_pool_b( + v_templates, + u_templates, + ci_lower_owned_full, + ) + for s in layout.my_owned_sites: + assert v_grads[s].shape == component_model.components[s].V.shape, ( + f"pool-B v_grad shape mismatch for {s!r}: " + f"{v_grads[s].shape} vs {component_model.components[s].V.shape}" + ) + assert u_grads[s].shape == component_model.components[s].U.shape, ( + f"pool-B u_grad shape mismatch for {s!r}: " + f"{u_grads[s].shape} vs {component_model.components[s].U.shape}" + ) + assert ci_grads[s].shape == ci.lower_leaky[s].shape, ( + f"pool-B ci_grad shape mismatch for {s!r}: " + f"{ci_grads[s].shape} vs {ci.lower_leaky[s].shape}" + ) + return _PoolBGrads(v_grads=v_grads, u_grads=u_grads, ci_grads=ci_grads) + + +def _combined_backward( + component_model: ComponentModel, + ci: Any, + ci_lower_leaves: dict[str, Tensor], + sl: slice, + pool_b: _PoolBGrads, + home: _HomeLossOutputs, + layout: BlockDDPLayout, + cfg: _TwoPoolRuntime, + p: PhaseProfiler, +) -> None: + """Phase a/7. Fused backward through home losses + CI-fn graph. + + Two contributions need to merge before the CI-fn backward: + + * Pool B's ``ci_grads[s]`` are FULL-batch shape (pool B saw all + positions for the rank's owned sites). + * Layerwise's ``ci_lower_leaves[s].grad`` lives at this rank's batch + slice ``sl`` only. + + We splat layerwise's slice grad into the full-batch pool-B tensor so the + backward through ``ci.lower_leaky`` sees the combined seed. Pool B's V/U + grads are simply added to the V/U .grad accumulator that layerwise + already populated. + """ + with p.phase("a/7_seed_and_backward"): + combined_ci_grads: dict[str, Tensor] = {} + for s in layout.my_owned_sites: + grad = pool_b.ci_grads[s].clone() + layerwise_slice_grad = ci_lower_leaves[s].grad + assert layerwise_slice_grad is not None, ( + f"layerwise should have populated ci_lower_leaves[{s!r}].grad" + ) + assert layerwise_slice_grad.shape == grad[sl].shape, ( + f"layerwise slice grad shape {layerwise_slice_grad.shape} != " + f"target slice shape {grad[sl].shape} for site {s!r}" + ) + grad[sl] += layerwise_slice_grad + combined_ci_grads[s] = grad + for s in layout.my_owned_sites: + comp = component_model.components[s] + assert comp.V.grad is not None and comp.U.grad is not None, ( + "layerwise should have populated V/U .grad" + ) + comp.V.grad.add_(pool_b.v_grads[s]) + comp.U.grad.add_(pool_b.u_grads[s]) + total_home = cfg.coeff_faith * home.loss_faith + cfg.coeff_imp * home.loss_imp + assert total_home.dim() == 0, f"total_home should be scalar; got {total_home.shape}" + torch.autograd.backward( + tensors=[total_home, *(ci.lower_leaky[s] for s in layout.my_owned_sites)], + grad_tensors=[None, *(combined_ci_grads[s] for s in layout.my_owned_sites)], + ) + + +def _in_block_all_reduce_grads( + layout: BlockDDPLayout, + all_params: list[nn.Parameter], + p: PhaseProfiler, +) -> None: + """Phase a/8. In-block DDP all-reduce on all params' grads. + + DDP partners within a block share the same site set; they batch-shard + layerwise (different slices) but compute identical faith/imp + (data-independent). AVG-within-block reconciles to the correct + full-batch grad. + """ + with p.phase("a/8_in_block_allreduce"): + layout.all_reduce_grads_in_block(all_params) + + +def _cross_pool_grad_clip( + component_params: list[nn.Parameter], + ci_fn_params: list[nn.Parameter], + layout: BlockDDPLayout, + cfg: _TwoPoolRuntime, + p: PhaseProfiler, +) -> _ClipNorms: + """Phase a/8b. Clip global gradient norm across all pool-A blocks. + + Single-pool's ``clip_grad_norm_`` works on a single rank's identical-via- + DDP grads; in multi-pool, parameters are sharded across blocks, so we + must compute the cross-block norm. ``cross_pool_clip_grad_norm`` does the + all-reduce-SUM with ``/n_per_block`` dedup (DDP partners within a block + hold identical grads post-phase-a/8). + """ + components_norm = 0.0 + ci_fn_norm = 0.0 + with p.phase("a/8b_grad_clip"): + n_per_block = layout.world.n_per_block + if cfg.grad_clip_norm_components is not None: + norm_t = cross_pool_clip_grad_norm( + component_params, + cfg.grad_clip_norm_components, + group=layout.world.pool_a_group, + n_replicas=n_per_block, + ) + assert norm_t.dim() == 0, f"clip norm should be scalar; got {norm_t.shape}" + components_norm = norm_t.item() + if cfg.grad_clip_norm_ci_fn is not None: + norm_t = cross_pool_clip_grad_norm( + ci_fn_params, + cfg.grad_clip_norm_ci_fn, + group=layout.world.pool_a_group, + n_replicas=n_per_block, + ) + ci_fn_norm = norm_t.item() + return _ClipNorms(components=components_norm, ci_fn=ci_fn_norm) + + +def _optimizer_step(optimizer: torch.optim.Optimizer, p: PhaseProfiler) -> None: + """Phase a/9. AdamW step on V/U + CI fn params.""" + with p.phase("a/9_opt_step"): + optimizer.step() + + +def _async_send_updated_vu_to_pool_b( + component_model: ComponentModel, + layout: BlockDDPLayout, + p: PhaseProfiler, +) -> _AsyncWorkHandle: + """Phase a/10. Async-ship updated V/U back to pool B for next step. + + The wait happens at the START of the next step (phase a/0), so this send + runs concurrently with pool B's next iter target_fwd + PPGD warmup. + """ + with p.phase("a/10_async_send_weights"): + 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} + works, buffers = layout.async_send_updated_weights_to_pool_b(v_owned, u_owned) + return _AsyncWorkHandle(works=works, buffers=buffers) + + +def _wait_async_ci_send(ci_send: _AsyncWorkHandle, p: PhaseProfiler) -> None: + """Phase a/11. Wait on the CI send from phase a/2. + + Must complete before the next step's CI forward mutates the underlying + ci.lower_leaky storage that the send buffers reference. + """ + with p.phase("a/11_wait_async_ci_send"): + for w in ci_send.works: + w.wait() + ci_send.buffers.clear() + + +def _stash_pending_weight_send( + component_model: ComponentModel, + weight_send: _AsyncWorkHandle, +) -> None: + """Stash the weight-send handle on the model for next step's phase a/0 wait.""" + component_model._pending_weight_sends = ( # type: ignore[attr-defined] + weight_send.works, + weight_send.buffers, + ) + + +def _step_metrics( + home: _HomeLossOutputs, + loss_stoch_scalar: float, + stoch_raw_num: float, + stoch_raw_den: int, + clip_norms: _ClipNorms, +) -> dict[str, float]: + """Per-step metrics dict: per-rank display scalars + raw ingredients for + the cross-block logger reduction + pre-clip grad norms. + + Raw ``_raw/_{num,den}`` are designed so the global value is + ``SUM(num)/SUM(den)`` (for ratio losses faith and stoch) or ``SUM(num)`` + (for additive losses imp). See ``two_pool/reductions.py``. + """ + return { + "loss/faith": home.loss_faith.item(), + "loss/imp": home.loss_imp.item(), + "loss/stoch": loss_stoch_scalar, + "_raw/faith_num": home.faith_sum_sq.item(), + "_raw/faith_den": float(home.faith_numel), + "_raw/imp_num": home.loss_imp.item(), + "_raw/stoch_num": stoch_raw_num, + "_raw/stoch_den": float(stoch_raw_den), + "grad_norm/components": clip_norms.components, + "grad_norm/ci_fn": clip_norms.ci_fn, + } + + +def run_faithfulness_warmup_pool_a( + *, + component_model: ComponentModel, + component_params: list[nn.Parameter], + n_steps: int, + lr: float, + weight_decay: float, + numel_global: int, +) -> None: + """Single-pool-equivalent faithfulness warmup on pool A only. + + Pool B has no V/U params to optimize, so the warmup is a no-op there. + Done once at startup so 2-pool runs initialize the same way as single-pool. + """ + warmup_opt = torch.optim.AdamW(component_params, lr=lr, weight_decay=weight_decay) + for _ in range(n_steps): + warmup_opt.zero_grad() + device = component_params[0].device + loss, _, _ = _faithfulness_loss(component_model, device, numel_global) + loss.backward() + warmup_opt.step() + del warmup_opt + torch.cuda.empty_cache() diff --git a/param_decomp/two_pool/pool_b.py b/param_decomp/two_pool/pool_b.py new file mode 100644 index 000000000..33bfcf381 --- /dev/null +++ b/param_decomp/two_pool/pool_b.py @@ -0,0 +1,345 @@ +# pyright: reportArgumentType=false + +"""Pool-B training step: stateless PPGD inner loop + cross-pool send/recv. + +Pool B is a PPGD replica. Each rank holds full-target V/U replicas (received +from pool A each step) and runs PPGD locally on its batch shard: + + 1. Recv CI values from owning pool-A ranks (async, overlapped with + target_fwd that doesn't depend on CI). + 2. Re-leaf CI as fp32 — it's wired in bf16 to save bandwidth, but pool A + wants fp32 grads back. + 3. PPGD warmup refines the persistent adversarial sources in place. + 4. Final PPGD recon loss with refined sources → autograd.grad to extract + V/U + CI gradients without polluting .grad. + 5. Source-tensor optimizer step on the same loss. + 6. SUM-reduce V/U grads within pool B (so pool A's incoming = full-batch). + 7. Ship grads to owning A ranks; receive updated V/U for next step. + +``step_pool_b`` reads as ~25 lines of orchestration; the named helpers below +correspond one-to-one with the phases above. Profiler tags (``b/_``) +are preserved so HTA traces remain comparable. +""" + +from typing import Any + +import torch +import torch.distributed as dist +from torch import Tensor + +from param_decomp.component_model import ComponentModel +from param_decomp.metrics.persistent_pgd_state import PersistentPGDState +from param_decomp.two_pool.layout import BlockDDPLayout +from param_decomp.two_pool.loss_strategy import LayerwiseLossStrategy +from param_decomp.two_pool.profiler import PhaseProfiler +from param_decomp.two_pool.runtime import _TwoPoolRuntime, autocast_bf16 + + +def step_pool_b( + layout: BlockDDPLayout, + component_model: ComponentModel, + ppgd_state: PersistentPGDState, + batch: Any, + cfg: _TwoPoolRuntime, + strategy: LayerwiseLossStrategy, + step: int, + n_steps: int, + profiler: PhaseProfiler | None = None, +) -> dict[str, float]: + """One training step on a pool-B rank. + + Reads as orchestration: each call delegates to a same-file helper that + implements one numbered phase. See the module docstring for the phase + list. + """ + p = profiler if profiler is not None else PhaseProfiler(enabled=False) + device = next(component_model.parameters()).device + batch_local = _slice_batch_for_pool_b(batch, layout) + + # Mirror single-pool's persistent_pgd_recon: advance LR once per step + # before warmup/forward consume it. + ppgd_state.update_lr(step=step, total_steps=n_steps) + weight_deltas = component_model.calc_weight_deltas() + + # ``strategy.context`` controls whether forwards return logits or pre-LM- + # head hidden state; ``ppgd_state``'s recon_loss was built to match. + with strategy.context(component_model.target_model): + ci_recv, ci_recv_works = _async_recv_ci_from_pool_a(layout, cfg, batch_local, device, p) + target_out = _target_forward_no_grad(component_model, batch_local, cfg, p) + _wait_ci_recv(ci_recv_works, p) + ci_scratch = _releaf_ci_fp32_for_grads(ci_recv) + + _ppgd_inner_warmup( + ppgd_state, component_model, batch_local, target_out, ci_scratch, weight_deltas, cfg, p + ) + sum_loss, n_examples = _ppgd_recon_forward( + ppgd_state, component_model, batch_local, target_out, ci_scratch, weight_deltas, cfg, p + ) + + # Scale by 1/N_pool_b so the upcoming SUM-reduce of V/U grads across pool B + # produces the full-batch gradient (each rank contributes its slice mean ÷ N). + total_ppgd = cfg.coeff_ppgd * (sum_loss / n_examples) / layout.world.n_pool_b + + v_grads, u_grads, ci_grads = _autograd_grad_for_vu_and_ci( + total_ppgd, component_model, ci_scratch, layout, p + ) + _ppgd_source_step_from_total_loss(ppgd_state, total_ppgd, p) + _sum_reduce_vu_grads_within_pool_b(v_grads, u_grads, layout, p) + _send_grads_to_pool_a(layout, v_grads, u_grads, ci_grads, p) + _recv_updated_vu_from_pool_a(component_model, layout, p) + + return _step_metrics(sum_loss, n_examples) + + +# ============================================================================= +# Phase helpers — each implements one numbered phase from the module docstring. +# ============================================================================= + + +def _slice_batch_for_pool_b(batch: Any, layout: BlockDDPLayout) -> Any: + """Pool B batch-shards across its ranks; pull this rank's slice.""" + sl = layout.my_batch_slice_b() + return batch[sl] if isinstance(batch, Tensor) else batch + + +def _async_recv_ci_from_pool_a( + layout: BlockDDPLayout, + cfg: _TwoPoolRuntime, + batch_local: Any, + device: torch.device, + p: PhaseProfiler, +) -> tuple[dict[str, Tensor], list[Any]]: + """Phase b/1. Post irecvs for CI values from owning A ranks. + + Returns ``(ci_recv, ci_recv_works)``. The works run on NIC concurrently + with the next phase's target_fwd on the GPU — pool B doesn't need CI for + target_fwd, so we save the recv latency by overlapping. + """ + with p.phase("b/1_post_async_recv_ci"): + seq_len = batch_local.shape[1] if batch_local.ndim >= 2 else 1 + ci_recv, works = layout.async_recv_ci_from_owners( + cfg.c_per_site, + seq_len=seq_len, + device=device, + dtype=torch.float32, + ) + # Pre-allocated buffers — values aren't filled until ``works`` are waited + # on, but shapes are fixed at allocation time and the post-wait shape + # must match what pool A sends. The assertion here costs nothing and + # would catch a wrong c_per_site config or a per-rank batch size + # mismatch between pools. + batch_local_b = layout.world.batch_local_b + for s, c in cfg.c_per_site.items(): + t = ci_recv[s] + assert t.shape == (batch_local_b, seq_len, c), ( + f"ci_recv[{s!r}] shape {tuple(t.shape)} != expected ({batch_local_b}, {seq_len}, {c})" + ) + return ci_recv, works + + +def _target_forward_no_grad( + component_model: ComponentModel, + batch_local: Any, + cfg: _TwoPoolRuntime, + p: PhaseProfiler, +) -> Tensor: + """Phase b/2. Frozen target forward on the rank's batch slice. + + Detached + no_grad — we only need the output as the recon target. Runs + in parallel with the CI recvs posted in phase b/1. + """ + with p.phase("b/2_target_fwd"), torch.no_grad(), autocast_bf16(cfg.bf16_autocast): + return component_model(batch_local).detach() + + +def _wait_ci_recv(ci_recv_works: list[Any], p: PhaseProfiler) -> None: + """Phase b/3. Block on the irecvs from phase b/1 — should already be done.""" + with p.phase("b/3_wait_ci_recv"): + for w in ci_recv_works: + w.wait() + + +def _releaf_ci_fp32_for_grads(ci_recv: dict[str, Tensor]) -> dict[str, Tensor]: + """Re-leaf received CI tensors as fp32 ``requires_grad=True``. + + CI is shipped in bf16 to save bandwidth; we upcast and re-leaf here so the + leaf has fp32 ``.grad`` that pool A can merge into its fp32 layerwise grads. + """ + return { + s: v.detach().to(torch.float32).clone().requires_grad_(True) for s, v in ci_recv.items() + } + + +def _ppgd_inner_warmup( + ppgd_state: PersistentPGDState, + component_model: ComponentModel, + batch_local: Any, + target_out: Tensor, + ci_scratch: dict[str, Tensor], + weight_deltas: dict[str, Tensor], + cfg: _TwoPoolRuntime, + p: PhaseProfiler, +) -> None: + """Phase b/4. Refine the persistent adversarial sources in place.""" + with p.phase("b/4_ppgd_warmup"), autocast_bf16(cfg.bf16_autocast): + ppgd_state.warmup( + model=component_model, + batch=batch_local, + target_out=target_out, + ci=ci_scratch, + weight_deltas=weight_deltas, + ) + + +def _ppgd_recon_forward( + ppgd_state: PersistentPGDState, + component_model: ComponentModel, + batch_local: Any, + target_out: Tensor, + ci_scratch: dict[str, Tensor], + weight_deltas: dict[str, Tensor], + cfg: _TwoPoolRuntime, + p: PhaseProfiler, +) -> tuple[Tensor, int]: + """Phase b/5. Final recon loss with the refined sources. + + Returns ``(sum_loss, n_examples)`` raw so the logger can SUM-reduce across + pool B to recover ``global_mean = SUM(sum_loss) / SUM(n_examples)``. + """ + with p.phase("b/5_ppgd_recon"), autocast_bf16(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 sum_loss, n_examples + + +def _autograd_grad_for_vu_and_ci( + total_ppgd: Tensor, + component_model: ComponentModel, + ci_scratch: dict[str, Tensor], + layout: BlockDDPLayout, + p: PhaseProfiler, +) -> tuple[dict[str, Tensor], dict[str, Tensor], dict[str, Tensor]]: + """Phase b/6. Extract V/U + CI gradients via ``torch.autograd.grad``. + + Uses ``autograd.grad`` rather than ``.backward()`` so we don't pollute the + components' ``.grad`` accumulator — pool A owns the optimizer, and we'll + ship these grads over the wire instead. ``retain_graph=True`` keeps the + graph alive for the upcoming PPGD source-step (which traverses the same + forward with respect to the source tensors). + """ + with p.phase("b/6_backward"): + all_sites = list(layout.world.all_sites) + params: list[Tensor] = [] + for s in all_sites: + params.append(component_model.components[s].V) + params.append(component_model.components[s].U) + ci_list = [ci_scratch[s] for s in all_sites] + grads = torch.autograd.grad(total_ppgd, params + ci_list, retain_graph=True) + + n_sites = len(all_sites) + v_grads = {s: grads[2 * i] for i, s in enumerate(all_sites)} + u_grads = {s: grads[2 * i + 1] for i, s in enumerate(all_sites)} + ci_grads = {s: grads[2 * n_sites + i] for i, s in enumerate(all_sites)} + for s in all_sites: + assert v_grads[s].shape == component_model.components[s].V.shape, ( + f"v_grad[{s!r}] shape {v_grads[s].shape} != " + f"V.shape {component_model.components[s].V.shape}" + ) + assert u_grads[s].shape == component_model.components[s].U.shape, ( + f"u_grad[{s!r}] shape {u_grads[s].shape} != " + f"U.shape {component_model.components[s].U.shape}" + ) + assert ci_grads[s].shape == ci_scratch[s].shape, ( + f"ci_grad[{s!r}] shape {ci_grads[s].shape} != ci_scratch shape {ci_scratch[s].shape}" + ) + return v_grads, u_grads, ci_grads + + +def _ppgd_source_step_from_total_loss( + ppgd_state: PersistentPGDState, + total_ppgd: Tensor, + p: PhaseProfiler, +) -> None: + """Phase b/7. Update the persistent adversarial sources from the same loss. + + ``retain_graph=False`` here — the upstream call already used + ``retain_graph=True`` to keep the graph alive for this final pass. + """ + with p.phase("b/7_ppgd_source_step"): + source_grads = ppgd_state.get_grads(total_ppgd, retain_graph=False) + ppgd_state.step(source_grads) + + +def _sum_reduce_vu_grads_within_pool_b( + v_grads: dict[str, Tensor], + u_grads: dict[str, Tensor], + layout: BlockDDPLayout, + p: PhaseProfiler, +) -> None: + """Phase b/8. SUM-reduce V/U grads within pool B. + + Combined with the ``/n_pool_b`` scaling on the loss, the SUM gives the + full-batch gradient that pool A would have computed if it had owned PPGD. + """ + with p.phase("b/8_pool_b_allreduce"): + for s in layout.world.all_sites: + dist.all_reduce(v_grads[s], op=dist.ReduceOp.SUM, group=layout.world.pool_b_group) + dist.all_reduce(u_grads[s], op=dist.ReduceOp.SUM, group=layout.world.pool_b_group) + + +def _send_grads_to_pool_a( + layout: BlockDDPLayout, + v_grads: dict[str, Tensor], + u_grads: dict[str, Tensor], + ci_grads: dict[str, Tensor], + p: PhaseProfiler, +) -> None: + """Phase b/9. Ship V/U + CI grads to the owning pool-A ranks.""" + with p.phase("b/9_send_grads_to_a"): + layout.send_pool_b_grads_to_owners(v_grads, u_grads, ci_grads) + + +def _recv_updated_vu_from_pool_a( + component_model: ComponentModel, + layout: BlockDDPLayout, + p: PhaseProfiler, +) -> None: + """Phase b/10. Receive updated V/U from pool A and copy in place. + + In-place copy avoids reallocating ``components[s].V`` / ``.U`` (which would + detach them from the optimizer if pool B had one — defense in depth even + though pool B doesn't). + """ + with p.phase("b/10_recv_weights"): + all_sites = layout.world.all_sites + v_templates = {s: component_model.components[s].V for s in all_sites} + u_templates = {s: component_model.components[s].U for s in all_sites} + v_new, u_new = layout.recv_updated_weights_from_owners(v_templates, u_templates) + with torch.no_grad(): + for s in all_sites: + component_model.components[s].V.copy_(v_new[s]) + component_model.components[s].U.copy_(u_new[s]) + + +def _step_metrics(sum_loss: Tensor, n_examples: int) -> dict[str, float]: + """Per-step metrics dict: per-rank display scalar + raw (num, den) for the + cross-pool-B logger reduction. + + Pool B batch-shards across its ranks. Global mean per position is + ``SUM(sum_loss) / SUM(n_examples)`` across the pool, NOT AVG of per-rank + ratios. ``n_examples`` is the same on every PPGD rank so AVG happens to + coincide, but reporting raw keeps the contract uniform with pool A's + cross-block reductions. + """ + return { + "loss/ppgd": (sum_loss / n_examples).item(), + "_raw/ppgd_num": sum_loss.item(), + "_raw/ppgd_den": float(n_examples), + } diff --git a/param_decomp/two_pool/profiler.py b/param_decomp/two_pool/profiler.py new file mode 100644 index 000000000..2f77a85db --- /dev/null +++ b/param_decomp/two_pool/profiler.py @@ -0,0 +1,86 @@ +"""``PhaseProfiler``: ``torch.profiler.profile`` wrapper for 2-pool training. + +Captures CPU + CUDA events into a Chrome trace JSON that loads directly into +Perfetto (https://ui.perfetto.dev). Use as a context manager around the +training loop with ``phase(name)`` annotations marking logical step phases. +""" + +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import torch +import torch.profiler + + +@dataclass +class PhaseProfiler: + """Thin wrapper around ``torch.profiler.profile`` for 2-pool training. + + Records CPU + CUDA events into a Chrome trace JSON that loads directly into + Perfetto (https://ui.perfetto.dev). Captures kernel timings, NCCL ops, and + GPU memory allocations. + + Use as a context manager around the training loop, with ``phase(name)`` + annotations marking logical step phases ("a/5_layerwise" etc.) inside. + Call ``step()`` once per training iteration so the schedule progresses. + + Trace schedule by default skips the first 20 iters (CUDA / JIT warmup) then + records 3 active steps and dumps the trace. + """ + + enabled: bool = False + out_dir: Path | None = None + rank: int = 0 + pool: Literal["a", "b"] = "a" + skip_first: int = 20 + active: int = 3 + _prof: torch.profiler.profile | None = None + + def __enter__(self) -> "PhaseProfiler": + if not self.enabled: + return self + assert self.out_dir is not None, "PhaseProfiler enabled but no out_dir given" + self.out_dir.mkdir(parents=True, exist_ok=True) + trace_path = self.out_dir / f"trace_pool{self.pool}_rank{self.rank}.json" + + def on_trace_ready(prof: "torch.profiler.profile") -> None: + prof.export_chrome_trace(str(trace_path)) + print(f"[profiler rank{self.rank}] wrote {trace_path}", flush=True) + + self._prof = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule( + skip_first=self.skip_first, wait=0, warmup=1, active=self.active, repeat=1 + ), + on_trace_ready=on_trace_ready, + record_shapes=False, + profile_memory=True, + with_stack=False, + ) + self._prof.__enter__() + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + if self._prof is not None: + self._prof.__exit__(exc_type, exc, tb) + self._prof = None + + @contextmanager + def phase(self, name: str) -> "Iterator[None]": + """Annotate a logical step phase. No-op when disabled.""" + if not self.enabled: + yield + return + with torch.profiler.record_function(name): + yield + + def step(self) -> None: + """Advance the profiler schedule (call once per training iteration).""" + if self._prof is not None: + self._prof.step() diff --git a/param_decomp/two_pool/reductions.py b/param_decomp/two_pool/reductions.py new file mode 100644 index 000000000..50cf1de5e --- /dev/null +++ b/param_decomp/two_pool/reductions.py @@ -0,0 +1,124 @@ +"""Cross-pool reductions for logging. + +The per-step metrics dict carries two flavors of values: + + * Per-rank display scalars (``loss/faith``, ``loss/imp``, ``loss/stoch``, + ``loss/ppgd``) — what the rank computed locally, in whatever per-rank + units the step function chose. Useful for sanity but **not directly + comparable across topologies** (different ranks own different sites and + batch slices). + + * Raw additive ingredients prefixed ``_raw/`` (``faith_num``, ``faith_den``, + ``imp_num``, ``stoch_num``, ``stoch_den``, ``ppgd_num``, ``ppgd_den``). + These are designed so the global scalar is recoverable from a cross-rank + SUM: + + - ``faith_global = SUM(faith_num) / SUM(faith_den)`` + - ``imp_global = SUM(imp_num)`` (no denominator) + - ``stoch_global = SUM(stoch_num) / SUM(stoch_den)`` + - ``ppgd_global = SUM(ppgd_num) / SUM(ppgd_den)`` + +Pool A's aggregator scales every raw value by ``1 / n_per_block`` *before* +the pool-A all-reduce SUM. That single division collapses two reductions +into one and is mathematically equivalent to AVG-within-block then +SUM-across-blocks: + + * For values that are identical across DDP partners in a block (faith and + imp — the CI fn forward runs on the FULL batch, so partners produce the + same scalar): ``sum_over_partners(value / n_per_block) = value``, and the + cross-block SUM recovers ``sum_over_blocks(value)``. + * For values that differ across DDP partners (stoch — partners process + disjoint batch slices): ``sum_over_partners(value / n_per_block) = + mean_over_partners(value)``, i.e. the cross-slice mean per site, which is + what we want before summing across blocks. + +Memory uses MAX (the bottleneck rank is what matters). +""" + +import torch +import torch.distributed as dist + +from param_decomp.two_pool.layout import BlockDDPLayout + +POOL_A_RAW_KEYS: tuple[str, ...] = ( + "_raw/faith_num", + "_raw/faith_den", + "_raw/imp_num", + "_raw/stoch_num", + "_raw/stoch_den", +) +POOL_B_RAW_KEYS: tuple[str, ...] = ( + "_raw/ppgd_num", + "_raw/ppgd_den", +) + + +def aggregate_max_memory_to_rank0( + layout: BlockDDPLayout, + device: torch.device, +) -> dict[str, float] | None: + """MAX-reduce CUDA peak memory within each pool; ship pool B's to rank 0. + + Returns ``{mem/pool_a_peak_gb, mem/pool_b_peak_gb}`` on rank 0, + ``None`` everywhere else. + """ + peak_gb = torch.cuda.max_memory_allocated() / 1e9 + val = torch.tensor([peak_gb], device=device) + if layout.my_pool == "a": + dist.all_reduce(val, op=dist.ReduceOp.MAX, group=layout.world.pool_a_group) + if layout.my_rank == 0: + pool_a_peak = val.item() + b_val = torch.empty(1, device=device) + dist.recv(b_val, src=layout.world.pool_b_ranks[0]) + return { + "mem/pool_a_peak_gb": pool_a_peak, + "mem/pool_b_peak_gb": b_val.item(), + } + return None + else: + dist.all_reduce(val, op=dist.ReduceOp.MAX, group=layout.world.pool_b_group) + if layout.my_is_pool_leader: + dist.send(val, dst=0) + return None + + +def aggregate_losses_to_rank0( + loss_dict: dict[str, float], + layout: BlockDDPLayout, + device: torch.device, +) -> dict[str, float] | None: + """SUM raw (num, den) ingredients within each pool, ship pool B's sums to + rank 0, and finalize to ``loss/`` global scalars there. + + See the module docstring for the math behind the ``1 / n_per_block`` + scale-then-SUM trick used on pool A. + """ + if layout.my_pool == "a": + keys = list(POOL_A_RAW_KEYS) + n_per_block = layout.world.n_per_block + # Scale before all-reduce SUM: see module docstring. + vals = torch.tensor( + [loss_dict[k] / n_per_block for k in keys], device=device, dtype=torch.float64 + ) + dist.all_reduce(vals, op=dist.ReduceOp.SUM, group=layout.world.pool_a_group) + if layout.my_rank == 0: + a = {k: vals[i].item() for i, k in enumerate(keys)} + # Receive pool B's raw sums (already pool-B-wide SUMs). + b_keys = list(POOL_B_RAW_KEYS) + b_vals = torch.empty(len(b_keys), device=device, dtype=torch.float64) + dist.recv(b_vals, src=layout.world.pool_b_ranks[0]) + b = {k: b_vals[i].item() for i, k in enumerate(b_keys)} + return { + "loss/faith": a["_raw/faith_num"] / a["_raw/faith_den"], + "loss/imp": a["_raw/imp_num"], + "loss/stoch": a["_raw/stoch_num"] / a["_raw/stoch_den"], + "loss/ppgd": b["_raw/ppgd_num"] / b["_raw/ppgd_den"], + } + return None + else: + keys = list(POOL_B_RAW_KEYS) + vals = torch.tensor([loss_dict[k] for k in keys], device=device, dtype=torch.float64) + dist.all_reduce(vals, op=dist.ReduceOp.SUM, group=layout.world.pool_b_group) + if layout.my_is_pool_leader: + dist.send(vals, dst=0) + return None diff --git a/param_decomp/two_pool/runtime.py b/param_decomp/two_pool/runtime.py new file mode 100644 index 000000000..36c25c779 --- /dev/null +++ b/param_decomp/two_pool/runtime.py @@ -0,0 +1,71 @@ +"""Internal step-context bundle for 2-pool training. + +``_TwoPoolRuntime`` glues serializable config (``TwoPoolConfig``) with the +caller-supplied runtime callables and the derived per-site C from the actual +target. It's the parameter shape that ``step_pool_a``/``step_pool_b`` consume; +``optimize_two_pool`` builds one internally per call. + +Not part of the public API. +""" + +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass +from typing import Any + +import torch + +from param_decomp.batch_and_loss_fns import ReconstructionLoss, RunBatch +from param_decomp.ci_fns import CiConfig +from param_decomp.ci_sigmoids import SigmoidType +from param_decomp.metrics.persistent_pgd_recon import PersistentPGDReconLossConfig +from param_decomp.two_pool.layout import BlockGroup + +__all__ = ["_TwoPoolRuntime", "autocast_bf16"] + + +@dataclass(frozen=True) +class _TwoPoolRuntime: + """Internal bundle passed to ``step_pool_a``/``step_pool_b``.""" + + block_groups: tuple[BlockGroup, ...] + pool_b_ranks: tuple[int, ...] + batch_global: int + c_per_site: dict[str, int] + ci_config: CiConfig + sigmoid_type: SigmoidType + run_batch: RunBatch + reconstruction_loss: ReconstructionLoss + ppgd_cfg: PersistentPGDReconLossConfig + coeff_faith: float + coeff_imp: float + coeff_stoch: float + coeff_ppgd: float + imp_min_pnorm: float + imp_min_beta: float + imp_min_eps: float + imp_min_p_anneal_start_frac: float + imp_min_p_anneal_final_p: float | None + imp_min_p_anneal_end_frac: float + lr_components: float + lr_ci_fn: float + grad_clip_norm_components: float | None + grad_clip_norm_ci_fn: float | None + # Total numel across all pool-A blocks' V/U weight_deltas. Used by + # ``_faithfulness_loss`` to match single-pool's per-element gradient scale + # (which divides by global numel, not rank-local owned numel). + numel_global: int + bf16_autocast: bool + use_fused_kl: bool + + +def autocast_bf16(enabled: bool) -> AbstractContextManager[Any]: + """bf16 autocast on CUDA when enabled, no-op otherwise. + + Wrapping target/CI forward + PPGD forward in bf16 unlocks PyTorch's + flash/cudnn SDPA backends (math is the only fp32 backend on H200 → ~6× + slower attention). Faithfulness loss (small-number-sensitive) is kept + OUTSIDE this block by the step functions. + """ + if enabled: + return torch.autocast(device_type="cuda", dtype=torch.bfloat16) + return nullcontext() diff --git a/param_decomp_lab/experiments/lm/run.py b/param_decomp_lab/experiments/lm/run.py index 43a8e9d9b..f7da5200d 100644 --- a/param_decomp_lab/experiments/lm/run.py +++ b/param_decomp_lab/experiments/lm/run.py @@ -30,6 +30,9 @@ from param_decomp.distributed import DistributedState, is_main_process from param_decomp.log import logger from param_decomp.optimize import EvalLoop, Trainer +from param_decomp.three_pool import ThreePoolConfig, ThreePoolTrainer +from param_decomp.training_state import ThreePoolTrainingState, TrainingState +from param_decomp.two_pool import TwoPoolConfig, TwoPoolTrainer from param_decomp_lab.batch_and_loss_fns import make_run_batch as _make_run_batch from param_decomp_lab.batch_and_loss_fns import recon_loss_kl from param_decomp_lab.component_model_io import load_component_model @@ -65,6 +68,7 @@ resolve_step, write_provenance, ) +from param_decomp_lab.run_sink import OnePoolSink, ThreePoolSink from param_decomp_lab.seed import set_seed @@ -114,7 +118,15 @@ class LMTargetConfig(BaseConfig): class LMExperimentConfig(ExperimentConfig[LMTargetConfig, LMDataConfig]): - pass + two_pool: TwoPoolConfig | None = None + """When set, training runs under the 2-pool strategy + (:class:`param_decomp.two_pool.TwoPoolTrainer`) instead of single-process + :class:`param_decomp.optimize.Trainer`. Mutually exclusive with ``three_pool``.""" + + three_pool: ThreePoolConfig | None = None + """When set, training runs under the 3-pool strategy + (:class:`param_decomp.three_pool.ThreePoolTrainer`). Mutually exclusive with + ``two_pool``.""" def build_target(target_cfg: LMTargetConfig) -> nn.Module: @@ -254,8 +266,11 @@ def _fresh_main( tags: str | None, run_id: str | None, ) -> None: - """Fresh-run path: parse YAML, build everything, train from step 0.""" + """Fresh-run path: parse YAML, dispatch on pool config, train from step 0.""" cfg = LMExperimentConfig.from_file(config_path) + assert not (cfg.two_pool is not None and cfg.three_pool is not None), ( + "two_pool and three_pool are mutually exclusive; set only one." + ) dist_state = init_distributed() if is_main_process(): @@ -274,31 +289,62 @@ def _fresh_main( ) target_model = build_target(cfg.target) - + is_multi_pool = cfg.two_pool is not None or cfg.three_pool is not None train_loader = build_lm_loader( cfg.target, cfg.data, split="train", device=device, batch_size=cfg.pd.batch_size, - dist_state=dist_state, + # Multi-pool requires the full global batch on every rank — each pool + # slices it locally. Single-pool uses standard DistributedSampler. + dist_state=None if is_multi_pool else dist_state, seed=cfg.pd.seed, ) - eval_loop = _build_eval_loop(cfg, device, dist_state) - sink = init_pd_run(cfg, group=group, tags=tags, run_id=run_id) + if cfg.three_pool is not None: + three_sink = init_pd_run( + cfg, sink_class=ThreePoolSink, group=group, tags=tags, run_id=run_id + ) + try: + three_trainer = ThreePoolTrainer( + target_model=target_model, + run_batch=make_run_batch(cfg.target), + reconstruction_loss=recon_loss_kl, + pd_config=cfg.pd, + runtime_config=cfg.runtime, + three_pool_config=cfg.three_pool, + ) + three_trainer.run(train_loader, three_sink, cfg.cadence) + finally: + three_sink.finish() + return + one_sink = init_pd_run(cfg, sink_class=OnePoolSink, group=group, tags=tags, run_id=run_id) + eval_loop = _build_eval_loop(cfg, device, dist_state) if cfg.two_pool is None else None try: - trainer = Trainer( - target_model=target_model, - run_batch=make_run_batch(cfg.target), - reconstruction_loss=recon_loss_kl, - pd_config=cfg.pd, - runtime_config=cfg.runtime, - ) - trainer.run(train_loader, sink, cfg.cadence, eval_loop) + if cfg.two_pool is not None: + two_trainer = TwoPoolTrainer( + target_model=target_model, + run_batch=make_run_batch(cfg.target), + reconstruction_loss=recon_loss_kl, + pd_config=cfg.pd, + runtime_config=cfg.runtime, + two_pool_config=cfg.two_pool, + cadence=cfg.cadence, + ) + two_trainer.run(train_loader, one_sink, cfg.cadence) + else: + trainer = Trainer( + target_model=target_model, + run_batch=make_run_batch(cfg.target), + reconstruction_loss=recon_loss_kl, + pd_config=cfg.pd, + runtime_config=cfg.runtime, + ) + trainer.run(train_loader, one_sink, cfg.cadence, eval_loop) finally: - sink.finish() + one_sink.finish() def _resume_main( @@ -309,7 +355,7 @@ def _resume_main( run_id: str | None, ) -> None: """Resume-run path: read parent `run_meta.yaml` + `training_.pth`, - rebuild trainer via `Trainer.from_snapshot`, continue training.""" + rebuild trainer via `from_snapshot`, continue training.""" resume_cfg = ResumeConfig.from_file(resume_cfg_path) parent_cfg = LMExperimentConfig.from_file(resume_cfg.from_run / RUN_META_FILENAME) @@ -333,39 +379,81 @@ def _resume_main( resolved_step = resolve_step(resume_cfg.from_run, resume_cfg.step) snapshot = read_training_snapshot(resume_cfg.from_run, resolved_step) - # Override the saved device with the current resume environment. Mutating - # the dict (model_dump output) in place is fine even on a frozen dataclass; - # we're changing a value the dataclass references, not rebinding the field. snapshot.runtime_config["device"] = device target_model = build_target(effective_cfg.target) + is_multi_pool = effective_cfg.two_pool is not None or effective_cfg.three_pool is not None train_loader = build_lm_loader( effective_cfg.target, effective_cfg.data, split="train", device=device, batch_size=effective_cfg.pd.batch_size, - dist_state=dist_state, + dist_state=None if is_multi_pool else dist_state, seed=effective_cfg.pd.seed, ) - eval_loop = _build_eval_loop(effective_cfg, device, dist_state) - sink = init_pd_run(effective_cfg, group=group, tags=tags, run_id=run_id) - if sink.out_dir is not None: + run_batch = make_run_batch(effective_cfg.target) + + if effective_cfg.three_pool is not None: + three_sink = init_pd_run( + effective_cfg, sink_class=ThreePoolSink, group=group, tags=tags, run_id=run_id + ) + if three_sink.out_dir is not None: + write_provenance( + three_sink.out_dir, + ResumeProvenance(parent_run_dir=resume_cfg.from_run, parent_step=resolved_step), + ) + try: + assert isinstance(snapshot, ThreePoolTrainingState), ( + f"3-pool resume needs ThreePoolTrainingState; got {type(snapshot).__name__}" + ) + three_trainer = ThreePoolTrainer.from_snapshot( + snapshot, + target_model=target_model, + run_batch=run_batch, + reconstruction_loss=recon_loss_kl, + ) + three_trainer.run(train_loader, three_sink, effective_cfg.cadence) + finally: + three_sink.finish() + return + + one_sink = init_pd_run( + effective_cfg, sink_class=OnePoolSink, group=group, tags=tags, run_id=run_id + ) + if one_sink.out_dir is not None: write_provenance( - sink.out_dir, + one_sink.out_dir, ResumeProvenance(parent_run_dir=resume_cfg.from_run, parent_step=resolved_step), ) - + eval_loop = ( + _build_eval_loop(effective_cfg, device, dist_state) + if effective_cfg.two_pool is None + else None + ) try: - trainer = Trainer.from_snapshot( - snapshot, - target_model=target_model, - run_batch=make_run_batch(effective_cfg.target), - reconstruction_loss=recon_loss_kl, + assert isinstance(snapshot, TrainingState), ( + f"1-pool / 2-pool resume needs TrainingState; got {type(snapshot).__name__}" ) - trainer.run(train_loader, sink, effective_cfg.cadence, eval_loop) + if effective_cfg.two_pool is not None: + two_trainer = TwoPoolTrainer.from_snapshot( + snapshot, + target_model=target_model, + run_batch=run_batch, + reconstruction_loss=recon_loss_kl, + cadence=effective_cfg.cadence, + ) + two_trainer.run(train_loader, one_sink, effective_cfg.cadence) + else: + trainer = Trainer.from_snapshot( + snapshot, + target_model=target_model, + run_batch=run_batch, + reconstruction_loss=recon_loss_kl, + ) + trainer.run(train_loader, one_sink, effective_cfg.cadence, eval_loop) finally: - sink.finish() + one_sink.finish() def _build_eval_loop( diff --git a/param_decomp_lab/experiments/resid_mlp/run.py b/param_decomp_lab/experiments/resid_mlp/run.py index 7d61e64c7..535230029 100644 --- a/param_decomp_lab/experiments/resid_mlp/run.py +++ b/param_decomp_lab/experiments/resid_mlp/run.py @@ -30,6 +30,7 @@ ) from param_decomp_lab.infra.paths import ModelPath from param_decomp_lab.infra.run_files import resolve_run_files +from param_decomp_lab.run_sink import OnePoolSink from param_decomp_lab.seed import set_seed @@ -148,7 +149,7 @@ def main( ) eval_loop = _build_eval_loop(cfg, device) - sink = init_pd_run(cfg, group=group, tags=tags) + sink = init_pd_run(cfg, sink_class=OnePoolSink, group=group, tags=tags) try: trainer = Trainer( diff --git a/param_decomp_lab/experiments/tms/run.py b/param_decomp_lab/experiments/tms/run.py index 4cd014b37..ec54aa286 100644 --- a/param_decomp_lab/experiments/tms/run.py +++ b/param_decomp_lab/experiments/tms/run.py @@ -30,6 +30,7 @@ ) from param_decomp_lab.infra.paths import ModelPath from param_decomp_lab.infra.run_files import resolve_run_files +from param_decomp_lab.run_sink import OnePoolSink from param_decomp_lab.seed import set_seed @@ -150,7 +151,7 @@ def main( ) eval_loop = _build_eval_loop(cfg, device) - sink = init_pd_run(cfg, group=group, tags=tags) + sink = init_pd_run(cfg, sink_class=OnePoolSink, group=group, tags=tags) try: trainer = Trainer( diff --git a/param_decomp_lab/experiments/utils.py b/param_decomp_lab/experiments/utils.py index d29d6685a..82df68f99 100644 --- a/param_decomp_lab/experiments/utils.py +++ b/param_decomp_lab/experiments/utils.py @@ -14,7 +14,7 @@ from param_decomp_lab.infra.run_files import generate_run_id from param_decomp_lab.infra.settings import PARAM_DECOMP_OUT_DIR from param_decomp_lab.infra.wandb import try_wandb -from param_decomp_lab.run_sink import RunSink +from param_decomp_lab.run_sink import OnePoolSink, ThreePoolSink RUN_META_FILENAME = "run_meta.yaml" @@ -58,30 +58,35 @@ class LMExperimentConfig(ExperimentConfig[LMTargetConfig, LMDataConfig]): wandb: WandbConfig | None = None -def init_pd_run[T: BaseConfig, D: BaseConfig]( +def init_pd_run[T: BaseConfig, D: BaseConfig, S: OnePoolSink | ThreePoolSink]( cfg: ExperimentConfig[T, D], *, + sink_class: type[S], group: str | None, tags: str | None, run_id: str | None = None, -) -> RunSink: +) -> S: """Allocate `run_id` + `out_dir`, write `run_meta.yaml`, return a sink. + `sink_class` picks the pool-specific sink (`OnePoolSink` for 1-pool runs, + `ThreePoolSink` for 3-pool). The choice is the caller's; this helper just + forwards through to the class's `local` / `with_wandb` / `silent` constructors. + Local-only when `cfg.wandb is None`, else wandb-backed. Non-main DDP ranks get a silent no-op sink without touching disk or wandb. `group` is a "launched together" id; `tags` is a comma-separated string of orthogonal labels. """ if not is_main_process(): - return RunSink.silent() + return sink_class.silent() run_id = run_id or generate_run_id("param_decomp") out_dir = PARAM_DECOMP_OUT_DIR / "decompositions" / run_id meta_path = out_dir / RUN_META_FILENAME cfg.to_file(meta_path) keep_last_n = cfg.cadence.keep_last_n_checkpoints if cfg.wandb is None: - return RunSink.local(out_dir, keep_last_n_checkpoints=keep_last_n) + return sink_class.local(out_dir, keep_last_n_checkpoints=keep_last_n) parsed_tags = [s.strip() for s in tags.split(",") if s.strip()] if tags else None - sink = RunSink.with_wandb( + sink = sink_class.with_wandb( out_dir, project=cfg.wandb.project, entity=cfg.wandb.entity, diff --git a/param_decomp_lab/run_sink.py b/param_decomp_lab/run_sink.py index e9a8d2521..89fe1ffe4 100644 --- a/param_decomp_lab/run_sink.py +++ b/param_decomp_lab/run_sink.py @@ -1,12 +1,25 @@ -"""Concrete `RunSink` for the in-repo experiments: local files + optional wandb. +"""Concrete `RunSink` classes used by the in-repo experiments and lab tooling. -Non-main ranks transparently get a no-op sink regardless of which constructor is used. +Two pool-specific sinks (`OnePoolSink`, `ThreePoolSink`) share a private base +(`_LabSinkBase`) for the local-files / wandb / console / log plumbing. The +subclasses differ only in `checkpoint`'s typed parameter — 1-pool takes +`TrainingState`, 3-pool takes `ThreePoolTrainingState`. Each delegates to a +shared `_persist` for the actual save. + +Three constructors on each: + + sink = OnePoolSink.local(out_dir) + sink = OnePoolSink.with_wandb(out_dir, project=..., run_id=..., ...) + sink = OnePoolSink.silent() # tests / quick checks + +(same shape for `ThreePoolSink`). Non-main ranks transparently get a no-op +sink regardless of which constructor is called. """ import json from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, Self import wandb from PIL import Image @@ -15,7 +28,7 @@ from param_decomp.base_config import BaseConfig from param_decomp.distributed import is_main_process from param_decomp.log import logger -from param_decomp.training_state import TrainingState +from param_decomp.training_state import ThreePoolTrainingState, TrainingState from param_decomp_lab.infra.run_files import save_file from param_decomp_lab.infra.wandb import init_wandb, try_wandb @@ -23,9 +36,9 @@ def _local_log(data: dict[str, Any], step: int, out_dir: Path) -> None: """Write a step's metrics, figures, and custom charts to disk. - PIL images go to `{out_dir}/figures/_.png`; `wandb.plot.CustomChart` - payloads go to `{out_dir}/figures/_.json`; everything else is appended - as one JSON line to `{out_dir}/metrics.jsonl`. + PIL images go to ``{out_dir}/figures/_.png``; ``wandb.plot.CustomChart`` + payloads go to ``{out_dir}/figures/_.json``; everything else is + appended as one JSON line to ``{out_dir}/metrics.jsonl``. """ metrics_file = out_dir / "metrics.jsonl" metrics_file.touch(exist_ok=True) @@ -53,20 +66,19 @@ def _local_log(data: dict[str, Any], step: int, out_dir: Path) -> None: @dataclass(frozen=True) -class RunSink: - """Construct via `local`, `with_wandb`, or `silent` (not the dataclass directly). +class _LabSinkBase: + """Shared local-files + wandb plumbing for `OnePoolSink` / `ThreePoolSink`. - Non-main ranks always get a no-op handle. `out_dir=None` disables disk output. + Pool-specific subclasses inherit constructors and the log / console / + finish / `_persist` helpers, and add typed `checkpoint` methods. """ out_dir: Path | None _wandb_active: bool keep_last_n_checkpoints: int | None = None - # =========================== Constructors =========================== - @classmethod - def local(cls, out_dir: Path, *, keep_last_n_checkpoints: int | None = None) -> "RunSink": + def local(cls, out_dir: Path, *, keep_last_n_checkpoints: int | None = None) -> Self: """Sink that writes to local files only (no wandb).""" if not is_main_process(): return cls(out_dir=None, _wandb_active=False) @@ -92,12 +104,8 @@ def with_wandb( group: str | None = None, view_meta: dict[str, Any] | None = None, keep_last_n_checkpoints: int | None = None, - ) -> "RunSink": - """Sink that writes to local files and a wandb run. - - Initializes wandb on the main rank via `init_wandb`; non-main ranks return a - silent no-op. - """ + ) -> Self: + """Sink that writes to local files + a wandb run. Non-main ranks are silent.""" if not is_main_process(): return cls(out_dir=None, _wandb_active=False) out_dir.mkdir(parents=True, exist_ok=True) @@ -119,40 +127,37 @@ def with_wandb( ) @classmethod - def silent(cls) -> "RunSink": + def silent(cls) -> Self: """No-op sink for tests and quick interactive runs.""" return cls(out_dir=None, _wandb_active=False) - # =========================== Output API =========================== - def log(self, metrics: dict[str, Any], step: int) -> None: - """Emit a flat metrics dict to disk and/or wandb. - - Values may be scalars, PIL images, or `wandb.plot.CustomChart` payloads. - """ + """Emit a flat metrics dict to disk and/or wandb.""" if self.out_dir is not None: _local_log(metrics, step, self.out_dir) if self._wandb_active: try_wandb(wandb.log, {k: _wandb_value(v) for k, v in metrics.items()}, step=step) def console(self, *lines: str) -> None: - """Print lines to stderr via `tqdm.write`. No-op on non-main ranks.""" + """Print lines via `tqdm.write`. No-op on non-main ranks.""" if not is_main_process(): return for line in lines: tqdm.write(line) - def checkpoint(self, snapshot: TrainingState) -> None: - """Save the snapshot as two files: `model_.pth` + `training_.pth`. + def finish(self) -> None: + """End-of-run cleanup.""" + if self._wandb_active and wandb.run is not None: + wandb.finish() + + def _persist(self, snapshot: TrainingState | ThreePoolTrainingState) -> None: + """Save the snapshot as `model_.pth` + `training_.pth`. `model_.pth` is just the component-model state dict — the artifact downstream tools (`SavedRun.load_model`, postprocessing) consume. - `training_.pth` is the full `TrainingState` (configs, optimizer - state, metric states, step) needed for resumption. - - No-op when `out_dir is None` (silent sink / non-main rank); wandb upload - only when wandb is active. Prunes older (model, training) pairs after the - write when ``keep_last_n_checkpoints`` is set. + `training_.pth` is the full canonical state. No-op on a silent / + non-main-rank sink. Prunes older (model, training) pairs after the write + when ``keep_last_n_checkpoints`` is set. """ if self.out_dir is None: return @@ -167,10 +172,21 @@ def checkpoint(self, snapshot: TrainingState) -> None: if self.keep_last_n_checkpoints is not None: _prune_old_checkpoints(self.out_dir, keep_last_n=self.keep_last_n_checkpoints) - def finish(self) -> None: - """End-of-run cleanup.""" - if self._wandb_active and wandb.run is not None: - wandb.finish() + +@dataclass(frozen=True) +class OnePoolSink(_LabSinkBase): + """Lab sink for 1-pool runs (satisfies `OnePoolRunSink`).""" + + def checkpoint(self, snapshot: TrainingState) -> None: + self._persist(snapshot) + + +@dataclass(frozen=True) +class ThreePoolSink(_LabSinkBase): + """Lab sink for 3-pool runs (satisfies `ThreePoolRunSink`).""" + + def checkpoint(self, snapshot: ThreePoolTrainingState) -> None: + self._persist(snapshot) def _wandb_value(v: Any) -> Any: diff --git a/param_decomp_lab/tests/test_gpt2.py b/param_decomp_lab/tests/test_gpt2.py index 7fd712647..5aba57ed8 100644 --- a/param_decomp_lab/tests/test_gpt2.py +++ b/param_decomp_lab/tests/test_gpt2.py @@ -22,7 +22,7 @@ from param_decomp_lab.batch_and_loss_fns import make_run_batch, recon_loss_kl from param_decomp_lab.eval_metrics.ci_l0 import CI_L0, CI_L0Config from param_decomp_lab.experiments.lm.data import LMDataConfig, create_lm_data_loader -from param_decomp_lab.run_sink import RunSink +from param_decomp_lab.run_sink import OnePoolSink from param_decomp_lab.seed import set_seed @@ -101,7 +101,7 @@ def collate_input_ids(batch: list[dict[str, Tensor]]) -> Tensor: collate_fn=collate_input_ids, ) - sink = RunSink.local(tmp_path) + sink = OnePoolSink.local(tmp_path) cadence = Cadence(train_log_every=50, save_every=None) eval_loop = EvalLoop( loader=eval_loader, diff --git a/param_decomp_lab/tests/test_resid_mlp.py b/param_decomp_lab/tests/test_resid_mlp.py index 5399d27a0..bf9728ebb 100644 --- a/param_decomp_lab/tests/test_resid_mlp.py +++ b/param_decomp_lab/tests/test_resid_mlp.py @@ -16,7 +16,7 @@ from param_decomp_lab.batch_and_loss_fns import recon_loss_mse, run_batch_first_element from param_decomp_lab.experiments.resid_mlp.data import ResidMLPDataset from param_decomp_lab.experiments.resid_mlp.models import ResidMLP, ResidMLPModelConfig -from param_decomp_lab.run_sink import RunSink +from param_decomp_lab.run_sink import OnePoolSink from param_decomp_lab.seed import set_seed @@ -103,7 +103,7 @@ def test_resid_mlp_decomposition_happy_path(tmp_path: Path) -> None: train_loader = DataLoader(train_dataset, batch_size=None) eval_loader = DataLoader(eval_dataset, batch_size=None) - sink = RunSink.local(tmp_path) + sink = OnePoolSink.local(tmp_path) cadence = Cadence(train_log_every=50, save_every=None) eval_loop = EvalLoop( loader=eval_loader, diff --git a/param_decomp_lab/tests/test_resumption.py b/param_decomp_lab/tests/test_resumption.py index d3770f57f..58a544521 100644 --- a/param_decomp_lab/tests/test_resumption.py +++ b/param_decomp_lab/tests/test_resumption.py @@ -29,7 +29,7 @@ read_training_snapshot, resolve_step, ) -from param_decomp_lab.run_sink import RunSink +from param_decomp_lab.run_sink import OnePoolSink class TinyLinear(nn.Module): @@ -87,7 +87,7 @@ def _cadence() -> Cadence: def test_run_sink_writes_model_and_training_files(tmp_path: Path) -> None: """A fresh run writes both `model_.pth` and `training_.pth` per save.""" run_dir = tmp_path / "run" - sink = RunSink.local(run_dir) + sink = OnePoolSink.local(run_dir) trainer = Trainer( target_model=TinyLinear(), @@ -119,7 +119,7 @@ def test_resume_round_trip_matches_uninterrupted_run(tmp_path: Path) -> None: runtime_config=_runtime(), ) full_sink_dir = tmp_path / "full" - trainer_full.run(_loader(), RunSink.local(full_sink_dir), _cadence()) + trainer_full.run(_loader(), OnePoolSink.local(full_sink_dir), _cadence()) final_full = {k: v.clone() for k, v in trainer_full.component_model.state_dict().items()} # Phase 1: train 2 steps, the sink writes training_2.pth. @@ -132,7 +132,7 @@ def test_resume_round_trip_matches_uninterrupted_run(tmp_path: Path) -> None: pd_config=_pd_config(steps=2), runtime_config=_runtime(), ) - trainer_half.run(_loader(), RunSink.local(parent_dir), _cadence()) + trainer_half.run(_loader(), OnePoolSink.local(parent_dir), _cadence()) assert (parent_dir / "training_2.pth").is_file() # Phase 2: resume from parent's training_2.pth, train to step 4. @@ -149,7 +149,7 @@ def test_resume_round_trip_matches_uninterrupted_run(tmp_path: Path) -> None: ) assert trainer_resumed.step == 2 resumed_dir = tmp_path / "resumed" - trainer_resumed.run(_loader(), RunSink.local(resumed_dir), _cadence()) + trainer_resumed.run(_loader(), OnePoolSink.local(resumed_dir), _cadence()) resumed_final = trainer_resumed.component_model.state_dict() assert final_full.keys() == resumed_final.keys() @@ -167,7 +167,7 @@ def test_resolve_step_finds_latest(tmp_path: Path) -> None: pd_config=_pd_config(steps=4), runtime_config=_runtime(), ) - trainer.run(_loader(), RunSink.local(parent_dir), _cadence()) + trainer.run(_loader(), OnePoolSink.local(parent_dir), _cadence()) assert resolve_step(parent_dir, "latest") == 4 assert resolve_step(parent_dir, 2) == 2 diff --git a/param_decomp_lab/tests/test_tms.py b/param_decomp_lab/tests/test_tms.py index 8c7e87cbc..d2adf9ad2 100644 --- a/param_decomp_lab/tests/test_tms.py +++ b/param_decomp_lab/tests/test_tms.py @@ -23,7 +23,7 @@ from param_decomp_lab.experiments.tms.data import SparseFeatureDataset from param_decomp_lab.experiments.tms.models import TMSModel, TMSModelConfig, TMSTrainConfig from param_decomp_lab.experiments.tms.train_tms import get_model_and_dataloader, train -from param_decomp_lab.run_sink import RunSink +from param_decomp_lab.run_sink import OnePoolSink from param_decomp_lab.seed import set_seed @@ -98,7 +98,7 @@ def test_tms_decomposition_happy_path(tmp_path: Path) -> None: train_loader = DataLoader(dataset, batch_size=None) eval_loader = DataLoader(dataset, batch_size=None) - sink = RunSink.local(tmp_path) + sink = OnePoolSink.local(tmp_path) cadence = Cadence(train_log_every=2, save_every=None) eval_loop = EvalLoop( loader=eval_loader, From d695195b0f13923798dcd6ce6bd70231711f7f14 Mon Sep 17 00:00:00 2001 From: Oliver Clive-Griffin Date: Tue, 26 May 2026 16:16:50 +0000 Subject: [PATCH 002/551] ci_fn: per-stage backward timing via CUDA events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GlobalSharedTransformerCiFn exposes enable_bwd_profile() that registers backward hooks on each boundary tensor (output of input projector, between each transformer block, output of output head). On backward, the hooks record torch.cuda.Events whose elapsed_time differences give the GPU time spent in each stage's backward. ThreePoolTrainer enables it on the CI pool under PD_CI_FN_BWD_PROFILE=1, and step_ci emits one trace line per stage after ci/8a so the existing analyzer surfaces them. Outcome at GPT-2 XL Q/K production scale: total GPU bwd through the CI fn is 34 ms while the ci/8a phase wall time is ~620 ms — the bwd is CPU/dispatch-bound, not compute-bound. GPU runs at ~30% of B200 bf16 peak through the bwd. Also: optimize.run() now uses an env-var-gated pre-step barrier (PD_TORCH_PROFILE_RANKS triggers it on every rank, not just profiled ones) so future torch.profiler runs don't deadlock at the per-step sync point. Co-Authored-By: Claude Opus 4.7 (1M context) --- param_decomp/ci_fns.py | 59 ++++++++++++++++++++++++++++- param_decomp/three_pool/optimize.py | 21 +++++++++- param_decomp/three_pool/step_ci.py | 26 +++++++++++++ 3 files changed, 104 insertions(+), 2 deletions(-) diff --git a/param_decomp/ci_fns.py b/param_decomp/ci_fns.py index e72cc633b..d6b10272e 100644 --- a/param_decomp/ci_fns.py +++ b/param_decomp/ci_fns.py @@ -300,6 +300,60 @@ def __init__( for _ in range(n_layers) ] ) + # Lazily-enabled per-stage backward timing. ``enable_bwd_profile`` populates this + # with one ``torch.cuda.Event(enable_timing=True)`` per boundary tensor; the next + # ``forward`` call registers backward hooks that record those events as gradient + # flows through. ``compute_bwd_breakdown`` returns ``elapsed_time`` between + # consecutive events. + self._bwd_events: dict[str, torch.cuda.Event] = {} + + def enable_bwd_profile(self) -> None: + """Create CUDA events for per-stage backward timing. + + Boundaries instrumented (in forward-execution order): + ``x_0`` (input projector out) → ``x_1`` ... ``x_n`` (block outputs) → ``output`` + (output head out). On backward, hooks record events in reverse order, so + ``elapsed_time(output, x_n)`` = output-head bwd, ``elapsed_time(x_{i+1}, x_i)`` = + block i bwd. The input projector's bwd is anchored by ``post_bwd`` (recorded by + the caller after ``torch.autograd.backward``) because ``concatenated`` doesn't + require grad — bwd stops at the input projector, so there's no hook anchor there. + """ + labels = [ + "output", + *(f"x_{i}" for i in range(self.n_transformer_layers, -1, -1)), + "post_bwd", + ] + self._bwd_events = {label: torch.cuda.Event(enable_timing=True) for label in labels} + + def record_post_bwd_event(self) -> None: + """Record the ``post_bwd`` event. Call right after ``torch.autograd.backward``.""" + if self._bwd_events: + self._bwd_events["post_bwd"].record() + + def compute_bwd_breakdown(self) -> dict[str, float]: + """Per-stage bwd times in ms. Empty if ``enable_bwd_profile`` was never called. + + Caller must ``torch.cuda.synchronize()`` first to ensure events are recorded. + """ + if not self._bwd_events: + return {} + n = self.n_transformer_layers + ev = self._bwd_events + out: dict[str, float] = {"output_head_bwd": ev["output"].elapsed_time(ev[f"x_{n}"])} + for i in range(n, 0, -1): + out[f"block_{i - 1}_bwd"] = ev[f"x_{i}"].elapsed_time(ev[f"x_{i - 1}"]) + out["input_projector_bwd"] = ev["x_0"].elapsed_time(ev["post_bwd"]) + return out + + def _maybe_hook(self, tensor: Tensor, label: str) -> None: + """Register a backward hook that records the event for ``label``. + + Skips if events disabled or the tensor doesn't require grad (no bwd will reach it). + """ + if not self._bwd_events or not tensor.requires_grad: + return + event = self._bwd_events[label] + tensor.register_hook(lambda _grad, e=event: e.record()) @override def forward( @@ -320,13 +374,16 @@ def forward( added_seq_dim = True x = projected - for block in self._blocks: + self._maybe_hook(x, "x_0") + for i, block in enumerate(self._blocks): x = block(x) + self._maybe_hook(x, f"x_{i + 1}") output = self._output_head(x) if added_seq_dim: output = output.squeeze(-2) + self._maybe_hook(output, "output") split_outputs = torch.split(output, self.split_sizes, dim=-1) outputs = {name: split_outputs[i] for i, name in enumerate(self.layer_order)} diff --git a/param_decomp/three_pool/optimize.py b/param_decomp/three_pool/optimize.py index cd1c8a762..ad77522e2 100644 --- a/param_decomp/three_pool/optimize.py +++ b/param_decomp/three_pool/optimize.py @@ -267,6 +267,19 @@ def __init__( assert self.component_model.ci_fn is not None trace("ThreePoolTrainer.__init__: torch.compile(ci_fn)") self.component_model.ci_fn = torch.compile(self.component_model.ci_fn) # pyright: ignore[reportAttributeAccessIssue] + if self.layout.my_pool == "ci" and os.environ.get("PD_CI_FN_BWD_PROFILE", "").strip() in ( + "1", + "true", + "yes", + ): + from param_decomp.ci_fns import GlobalSharedTransformerCiFn + + assert self.component_model.ci_fn is not None + for m in self.component_model.ci_fn.modules(): + if isinstance(m, GlobalSharedTransformerCiFn): + m.enable_bwd_profile() + trace("ThreePoolTrainer.__init__: enabled CI fn bwd-stage profile") + break # Diverge stochastic RNG per rank for mask sampling. seed_per_rank(pd_config.seed) @@ -566,6 +579,12 @@ def run( ci_fn_lr_schedule = pd_config.ci_fn_optimizer.lr_schedule profiler_ctx = profiler if profiler is not None else nullcontext() + # All-ranks barrier flag — read identically on every rank from the env + # so even unprofiled ranks join the pre-step ``dist.barrier()`` when + # profiling is enabled somewhere in the world. Without this, a barrier + # gated on the local ``profiler is not None`` deadlocks any rank not + # in ``PD_TORCH_PROFILE_RANKS``. + pre_step_barrier_enabled = bool(os.environ.get("PD_TORCH_PROFILE_RANKS", "").strip()) h_cache_ci: dict[str, Tensor] | None = None # Async-pipeline state threaded across iterations on LW + PPGD pools. pending_all_reduce_lw: list[tuple[list[Tensor], Tensor, dist.Work]] | None = None @@ -614,7 +633,7 @@ def _to_device(b: Any) -> Any: lr_step, n_steps, components_lr_schedule ) - if profiler is not None: + if pre_step_barrier_enabled: dist.barrier() torch.cuda.synchronize(device) diff --git a/param_decomp/three_pool/step_ci.py b/param_decomp/three_pool/step_ci.py index ec966faff..c2e020cae 100644 --- a/param_decomp/three_pool/step_ci.py +++ b/param_decomp/three_pool/step_ci.py @@ -39,6 +39,7 @@ import torch.nn as nn from torch import Tensor +from param_decomp._trace import trace from param_decomp.component_model import CIOutputs, ComponentModel from param_decomp.grad_clip import cross_pool_clip_grad_norm from param_decomp.metrics.importance_minimality import ( @@ -117,6 +118,7 @@ def step_ci( optimizer.zero_grad(set_to_none=True) _fused_backward_through_ci_fn(loss_imp, ci, g_ci_total, layout, cfg, p) + _maybe_emit_ci_fn_bwd_breakdown(component_model) with p.phase("ci/9_in_pool_allreduce"): layout.all_reduce_ci_fn_grads(ci_fn_params) @@ -219,6 +221,30 @@ def _assemble_g_ci_total( return g_ci_total +def _maybe_emit_ci_fn_bwd_breakdown(component_model: ComponentModel) -> None: + """Emit per-stage CI fn bwd times as ``trace()`` lines when the bwd profile is on. + + Records the ``post_bwd`` event immediately (anchoring the input projector's bwd + end time on the stream), synchronizes to flush events, then walks the CI fn + modules to find the ``GlobalSharedTransformerCiFn`` instance and emits one + ``phase: ci/8a_stage_