three_pool: relax cross-pool divisibility to bidirectional

CLAUDE.md

@@ -117,6 +117,7 @@ from param_decomp.batch_and_loss_fns import RunBatch, ReconstructionLoss
 | `param_decomp_lab/investigate/` | `param_decomp_lab/investigate/CLAUDE.md` | Agent investigation of a research question |
 | `param_decomp_lab/app/` | `param_decomp_lab/app/CLAUDE.md` | Web visualization (FastAPI + Svelte) |
 | `param_decomp_lab/experiments/lm/pretrain/` | `param_decomp_lab/experiments/lm/pretrain/CLAUDE.md` | LM target-model pretraining |
+| `param_decomp_lab/three_pool/` | `param_decomp_lab/three_pool/CLAUDE.md` | 3-pool (CI / Layerwise / PPGD) parallel training; cross-pool batch routing |
 
 ## Saved-run layout
 

param_decomp_lab/tests/test_three_pool_batch_edge.py

@@ -0,0 +1,97 @@
+"""Unit tests for ``BatchEdge`` — the cross-pool batch-slice geometry.
+
+CPU-only, no torch.distributed. Exercises both fan directions: CI-coarse
+(``n_ci`` divides ``n_down``) and CI-fine / inverted (``n_down`` divides
+``n_ci``), plus the square case. The invariant under test is that the
+``(ci_slice, down_slice)`` overlaps tile each pool's local batch tensor exactly
+once and agree on the global rows they cover from both sides — which is what
+makes the cross-pool sends/recvs stitch back to the single-pool batch.
+"""
+
+import pytest
+
+from param_decomp_lab.three_pool.layout import BatchEdge
+
+# (n_ci, n_down, batch_global), spanning coarse / square / inverted regimes.
+EDGE_CASES = [
+    (2, 4, 8),  # CI coarse (legacy forward path), fanout 2
+    (2, 2, 8),  # square
+    (4, 2, 8),  # CI fine (inverted), fanout 2
+    (4, 4, 8),  # square, all-equal
+    (8, 4, 8),  # CI fine, fanout 2, b_down > b_ci
+    (4, 8, 8),  # CI coarse, fanout 2
+    (1, 4, 8),  # single CI rank fans to all
+    (6, 6, 12),  # square non-power-of-two
+]
+
+
+@pytest.mark.parametrize(("n_ci", "n_down", "batch_global"), EDGE_CASES)
+def test_overlaps_tile_local_tensors_and_agree_globally(
+    n_ci: int, n_down: int, batch_global: int
+) -> None:
+    edge = BatchEdge(n_ci=n_ci, n_down=n_down, batch_global=batch_global)
+
+    # Every downstream slice is fully tiled by its CI overlaps (no gaps/overlaps).
+    for j in range(n_down):
+        covered = sorted(
+            (s.start, s.stop)
+            for s in (edge.overlap_within_down(i, j) for i in edge.ci_slices_for_down_slice(j))
+        )
+        cursor = 0
+        for start, stop in covered:
+            assert start == cursor, f"gap/overlap in down slice {j}: {covered}"
+            cursor = stop
+        assert cursor == edge.b_down
+
+    # Every CI slice is fully tiled by its downstream overlaps.
+    for i in range(n_ci):
+        covered = sorted(
+            (s.start, s.stop)
+            for s in (edge.overlap_within_ci(i, j) for j in edge.down_slices_for_ci_slice(i))
+        )
+        cursor = 0
+        for start, stop in covered:
+            assert start == cursor, f"gap/overlap in CI slice {i}: {covered}"
+            cursor = stop
+        assert cursor == edge.b_ci
+
+    # Each overlap covers the SAME global rows from both sides, with matching length.
+    for i in range(n_ci):
+        for j in edge.down_slices_for_ci_slice(i):
+            ci_sub = edge.overlap_within_ci(i, j)
+            down_sub = edge.overlap_within_down(i, j)
+            assert (ci_sub.stop - ci_sub.start) == (down_sub.stop - down_sub.start)
+            ci_global = (i * edge.b_ci + ci_sub.start, i * edge.b_ci + ci_sub.stop)
+            down_global = (j * edge.b_down + down_sub.start, j * edge.b_down + down_sub.stop)
+            assert ci_global == down_global
+
+
+@pytest.mark.parametrize(("n_ci", "n_down", "batch_global"), EDGE_CASES)
+def test_pairing_is_symmetric(n_ci: int, n_down: int, batch_global: int) -> None:
+    """``j in down_slices_for_ci_slice(i)`` iff ``i in ci_slices_for_down_slice(j)``."""
+    edge = BatchEdge(n_ci=n_ci, n_down=n_down, batch_global=batch_global)
+    for i in range(n_ci):
+        for j in edge.down_slices_for_ci_slice(i):
+            assert i in edge.ci_slices_for_down_slice(j)
+    for j in range(n_down):
+        for i in edge.ci_slices_for_down_slice(j):
+            assert j in edge.down_slices_for_ci_slice(i)
+
+
+def test_fanout_direction() -> None:
+    coarse = BatchEdge(n_ci=2, n_down=4, batch_global=8)
+    assert coarse.ci_is_coarse and coarse.fanout == 2
+    # One CI rank feeds 2 down ranks; each down rank reads from 1 CI rank.
+    assert coarse.down_slices_for_ci_slice(0) == (0, 1)
+    assert coarse.ci_slices_for_down_slice(3) == (1,)
+
+    fine = BatchEdge(n_ci=4, n_down=2, batch_global=8)
+    assert not fine.ci_is_coarse and fine.fanout == 2
+    # One down rank gathers from 2 CI ranks; each CI rank feeds 1 down rank.
+    assert fine.ci_slices_for_down_slice(0) == (0, 1)
+    assert fine.down_slices_for_ci_slice(3) == (1,)
+
+
+def test_rejects_non_cross_divisible() -> None:
+    with pytest.raises(AssertionError):
+        BatchEdge(n_ci=6, n_down=4, batch_global=12)

param_decomp_lab/three_pool/CLAUDE.md

@@ -0,0 +1,68 @@
+# three_pool
+
+3-pool training strategy for SPD on large frozen targets. Splits GPUs into
+three rank-disjoint, wall-clock-parallel pools: **CI** (replicated CI fn, DP
+across batch), **Layerwise/LW** (V/U sharded by site, block-DDP within group),
+**PPGD** (stateless full V/U replica + persistent sources, DP across batch).
+A dedicated unsharded CI pool is what makes a *global shared transformer* CI fn
+physically realizable. See `DESIGN.md` for the per-step dependency graph and the
+pipelining tricks.
+
+## Files
+
+| File | What it covers |
+|---|---|
+| `config.py` | `ThreePoolConfig` — serializable topology (which ranks form each pool, LW block groups). Pairs with a regular `PDConfig`. Topology integrity asserts live in `validate_topology`. |
+| `layout.py` | `World` (declarative topology + process groups), `ThreePoolLayout` (this rank's view + the six cross-pool comm methods), `BatchEdge` (cross-pool batch-slice geometry). |
+| `optimize.py` | `ThreePoolTrainer` / `optimize_three_pool` — composition root; mirrors `param_decomp.optimize.optimize`. PDConfig compatibility asserts in `_validate_pd_config_for_three_pool`. |
+| `step_ci.py`, `step_layerwise.py`, `step_ppgd.py` | The three per-pool step functions. Each calls the layout's cross-pool exchanges. |
+| `reductions.py` | Per-pool SUM/MAX reductions for logging; the `1/n_per_block` LW collapse trick. |
+| `eval_step.py` | Cross-pool eval pass; builds `MetricContext` on PPGD, ships CIOutputs CI→PPGD. |
+| `checkpoint.py` | Gather full state dict to rank 0. |
+
+## Batch-split routing (the cross-pool wrinkle)
+
+CI/LW/PPGD each shard the global batch on their own axis, so CI values (and the
+grads coming back) must be routed across pools along the batch dim. The
+constraint (`config.validate_topology` + `optimize._validate_pd_config_*`):
+
+- Each arity divides the global batch: `B % N_ci`, `B % N_per_block`, `B % N_ppgd`.
+- Each cross-pool edge (CI↔LW, CI↔PPGD) is **cross-divisible**: one arity
+  divides the other, in *either* direction.
+
+Cross-divisibility keeps every CI↔downstream overlap a whole, aligned sub-slice.
+`BatchEdge` (in `layout.py`) is the single source of truth for the geometry and
+handles **both** fan directions uniformly:
+
+- **CI coarse** (`N_ci ≤ N_down`): one CI rank fans a sub-slice out to
+  `fanout = N_down/N_ci` downstream ranks; grads stitch back fanout-to-one.
+- **CI fine / inverted** (`N_ci > N_down`): one downstream rank gathers CI from
+  `fanout = N_ci/N_down` CI ranks (concat) and scatters grads to those same K.
+
+`BatchEdge.{ci_slices_for_down_slice, down_slices_for_ci_slice}` give the rank
+pairings; `{overlap_within_ci, overlap_within_down}` give the matching sub-slice
+on each side. The exchange methods in `ThreePoolLayout` are written once against
+this API and never branch on regime. Pairs where neither arity divides the other
+(ragged overlaps) are out of scope and rejected by the validator.
+
+`PendingCiRecv` holds one packet (CI-coarse) or `fanout` packets (CI-fine, one
+per source CI rank) and stitches them into the downstream rank's full `[b_down]`
+CI tensor on `wait_and_unpack()`.
+
+## Equivalence harness
+
+`scripts/hetero_topology_equiv/` compares 3-pool loss curves to a single-pool
+reference (same seed/losses/coeffs/steps/batch). `compare.py` reads `train/loss/*`
+from `metrics.jsonl` and reports per-term 3p/1p ratios. Use a relaxation-
+requiring topology (e.g. `threepool_inverted.yaml`: `N_ci=4, N_per_block=2,
+N_ppgd=2`) to exercise the inverted path, and a current-legal one
+(`threepool_legal.yaml`: `N_ci=2, N_per_block=4`) to guard against regression.
+`faith` is ~bit-stable; `stoch`/`imp`/`ppgd` track within the ~1-5%
+non-determinism band.
+
+## Gotchas
+
+- Rank 0 must be the LW block-0 leader (reductions ship CI/PPGD losses to rank 0).
+- No distributed-aware checkpoint resumption yet (`save_every` stays None).
+- All cross-pool p2p goes through `world.cross_pool_p2p_group`, structurally
+  separate from the default communicator (barriers) to avoid NCCL wedging.

param_decomp_lab/three_pool/DESIGN.md

@@ -191,17 +191,33 @@ The defining wrinkle is **3-way batch slicing**:
 - 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.
+CI→PPGD sends could become **many-to-many along the batch dim**. To keep
+routing tractable, each cross-pool edge is constrained to be **cross-divisible**:
+one arity divides the other (`N_ci | N_lw_per_block OR N_lw_per_block | N_ci`,
+and likewise for `N_ci / N_pgd`). The batch-divisibility of each arity itself
+(`B % N_ci`, `B % N_lw_per_block`, `B % N_pgd`) is also required. Cross-
+divisibility makes every CI↔downstream overlap a **whole, aligned sub-slice**
+(uniform integer fan-in/out) — never ragged. The two regimes per edge:
+
+- **CI coarse** (`N_ci ≤ N_down`): one CI rank fans a sub-slice out to
+  `K = N_down/N_ci` downstream ranks; grads stitch back K-to-one. (Original
+  MVP direction.)
+- **CI fine / inverted** (`N_ci > N_down`): one downstream rank gathers from
+  `K = N_ci/N_down` CI ranks (concat) and scatters grads back to those same K.
+
+`layout.BatchEdge` captures one edge symmetrically and answers every routing
+question (which ranks pair, and the overlap sub-slice on each side) for both
+regimes, so the exchange methods never branch on direction. Pairs where
+*neither* arity divides the other (e.g. `N_ci=6, N_per_block=4`) remain out of
+scope — they'd produce ragged overlaps straddling slice boundaries — and the
+validator rejects them.
 
 ## 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.
+1. **Batch-split divisibility.** RESOLVED: each cross-pool edge's two arities
+   must be cross-divisible (one divides the other), in EITHER direction; each
+   arity must also divide the global batch. The validator rejects anything else
+   loudly. `BatchEdge` handles both fan directions; see "Routing complexity".
 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

param_decomp_lab/three_pool/config.py

@@ -16,7 +16,7 @@
 global shared transformer 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
+cross-divisibility with Layerwise/PPGD arities) 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.
@@ -54,21 +54,24 @@ class ThreePoolConfig(BaseConfig):
 
     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.
+      * ``N_ci`` and ``N_per_block_layerwise`` must be cross-divisible — one
+        divides the other. Whichever is smaller owns the coarser batch slice;
+        the finer pool's shards each nest in exactly one coarse shard, so the
+        CI↔LW exchange is a clean one-to-K fan-out (and K-to-one reduction)
+        in either direction.
+      * ``N_ci`` and ``N_ppgd`` must be cross-divisible for the same reason on
+        the CI↔PPGD side.
+
+    These two constraints keep every batch-slice overlap a whole, aligned
+    sub-slice (uniform integer fan-in/out), avoiding ragged many-to-many
+    routing. 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).",
+        "Must be cross-divisible with both N_per_block_layerwise and N_ppgd "
+        "(see class docstring).",
     )
     layerwise_block_groups: list[LayerwiseBlockGroupSpec] = Field(
         ...,
@@ -132,17 +135,21 @@ def validate_topology(self) -> Self:
             f"Layerwise pool and PPGD pool must be rank-disjoint; overlap: {lw_pgd_overlap}"
         )
 
-        # Batch-split divisibility (MVP constraint — see DESIGN.md open Q1).
+        # Batch-split divisibility (see DESIGN.md open Q1). Each cross-pool edge
+        # (CI↔LW, CI↔PPGD) must be a clean one-to-K fan-out in EITHER direction:
+        # one arity divides the other, so every batch-slice overlap is a whole,
+        # aligned sub-slice. (The batch-divisibility of each arity itself is
+        # enforced against pd.batch_size in _validate_pd_config_for_three_pool.)
         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_per_block % n_ci == 0 or n_ci % n_per_block == 0, (
+            f"N_ci ({n_ci}) and N_per_block_layerwise ({n_per_block}) must be "
+            f"cross-divisible (one divides the other) so each batch-slice overlap "
+            f"is a whole sub-slice. Tip: make one a multiple of the other."
         )
-        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."
+        assert n_ppgd % n_ci == 0 or n_ci % n_ppgd == 0, (
+            f"N_ci ({n_ci}) and N_ppgd ({n_ppgd}) must be cross-divisible (one "
+            f"divides the other) so each batch-slice overlap is a whole sub-slice. "
+            f"Tip: make one a multiple of the other."
         )
         return self