diff --git a/CLAUDE.md b/CLAUDE.md index 266093cf4..1b54072cc 100644 --- a/CLAUDE.md +++ b/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 diff --git a/param_decomp_lab/tests/test_three_pool_batch_edge.py b/param_decomp_lab/tests/test_three_pool_batch_edge.py new file mode 100644 index 000000000..18b95a6ef --- /dev/null +++ b/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) diff --git a/param_decomp_lab/three_pool/CLAUDE.md b/param_decomp_lab/three_pool/CLAUDE.md new file mode 100644 index 000000000..f9dd0a0d9 --- /dev/null +++ b/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. diff --git a/param_decomp_lab/three_pool/DESIGN.md b/param_decomp_lab/three_pool/DESIGN.md index f07df7716..4a6813f8a 100644 --- a/param_decomp_lab/three_pool/DESIGN.md +++ b/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 diff --git a/param_decomp_lab/three_pool/config.py b/param_decomp_lab/three_pool/config.py index c1616263f..6858464fc 100644 --- a/param_decomp_lab/three_pool/config.py +++ b/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 diff --git a/param_decomp_lab/three_pool/layout.py b/param_decomp_lab/three_pool/layout.py index 80e286b52..64507db52 100644 --- a/param_decomp_lab/three_pool/layout.py +++ b/param_decomp_lab/three_pool/layout.py @@ -25,14 +25,19 @@ The defining wrinkle is **3-way batch slicing**: CI/LW/PPGD each shard the global batch on their own axis. The constraint (enforced in -``ThreePoolConfig.validate_topology``) is: - - 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. +``ThreePoolConfig.validate_topology``) is that each cross-pool edge's two +arities are cross-divisible — one divides the other: + + N_ci | N_per_block_layerwise OR N_per_block_layerwise | N_ci + N_ci | N_ppgd OR N_ppgd | N_ci + +so each CI↔LW / CI↔PPGD edge is a clean one-to-K fan-out in WHICHEVER +direction the smaller arity owns the coarser slice. When CI is coarser +(``N_ci`` smaller) one CI rank fans a sub-slice out to K downstream ranks and +stitches K grads back. When CI is finer (``N_ci`` larger) one downstream rank +gathers from K CI ranks (concat) and scatters grads back to those K. The +``BatchEdge`` geometry (see "Batch-split routing" below) answers both +directions uniformly so the exchange methods don't branch on regime. """ # pyright: reportIndexIssue=false, reportArgumentType=false, reportOperatorIssue=false @@ -122,34 +127,58 @@ def flush_nccl_event_timings() -> None: @dataclass(frozen=True) -class PendingCiRecv: - """One coalesced CI-values irecv, held until ``wait_and_unpack()``. +class _CiRecvPacket: + """One incoming CI-values irecv covering ``row_slice`` of the destination. - The packed buffer carries ``sites`` worth of CI values (in order) as - ``b * seq_len * c_s`` ``_WIRE_DTYPE`` elements each. ``wait_and_unpack`` - blocks on the underlying ``dist.Work`` then materializes per-site - ``[b, seq_len, c_s]`` views into the packed buffer (no copy). + The packed buffer holds, for each site in ``sites`` order, ``overlap_len * + seq_len * c_s`` ``_WIRE_DTYPE`` elements, where ``overlap_len`` is the length + of ``row_slice`` along the batch dim. """ packed: torch.Tensor work: "dist.Work" + row_slice: slice + + +@dataclass(frozen=True) +class PendingCiRecv: + """One or more coalesced CI-values irecvs, held until ``wait_and_unpack()``. + + Covers the full ``[b_down]`` CI tensor for this downstream rank's owned + sites. In the CI-coarse regime that's a single packet spanning the whole + tensor; in the CI-fine regime it's ``fanout`` packets, each filling a + disjoint ``row_slice`` from a different CI rank. ``wait_and_unpack`` blocks + on all works, stitches each packet's overlap rows into per-site + ``[b_down, seq_len, c_s]`` destinations, and returns them. + """ + + packets: tuple[_CiRecvPacket, ...] sites: tuple[str, ...] site_to_c: dict[str, int] - b: int + b_down: int seq_len: int + device: torch.device def wait_and_unpack(self) -> dict[str, torch.Tensor]: - self.work.wait() - out: dict[str, torch.Tensor] = {} - offset = 0 - for s in self.sites: - c_s = self.site_to_c[s] - numel = self.b * self.seq_len * c_s - out[s] = self.packed[offset : offset + numel].view(self.b, self.seq_len, c_s) - offset += numel - assert offset == self.packed.numel(), ( - f"unpack size mismatch: consumed {offset} of {self.packed.numel()}" - ) + out = { + s: torch.empty( + self.b_down, self.seq_len, self.site_to_c[s], device=self.device, dtype=_WIRE_DTYPE + ) + for s in self.sites + } + for packet in self.packets: + packet.work.wait() + overlap_len = packet.row_slice.stop - packet.row_slice.start + offset = 0 + for s in self.sites: + c_s = self.site_to_c[s] + numel = overlap_len * self.seq_len * c_s + view = packet.packed[offset : offset + numel].view(overlap_len, self.seq_len, c_s) + out[s][packet.row_slice].copy_(view) + offset += numel + assert offset == packet.packed.numel(), ( + f"unpack size mismatch: consumed {offset} of {packet.packed.numel()}" + ) return out @@ -265,67 +294,119 @@ def block_idx_of_site(self, site: str) -> int: 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) ── + # ── Batch-split routing (the cross-pool wrinkle) ── + # + # CI/LW/PPGD each shard the global batch on their own axis. The validator + # constrains each cross-pool edge (CI↔LW, CI↔PPGD) so the two arities are + # cross-divisible: one slice nests an integer number of the other's slices. + # ``BatchEdge`` captures one such edge symmetrically — it answers every + # routing question for both the "CI coarser" and "CI finer" regimes — so the + # exchange methods don't branch on direction. @property - def k_lw_per_ci(self) -> int: - """How many LW batch shards (per block) fit inside one CI batch shard. + def ci_lw_edge(self) -> "BatchEdge": + return BatchEdge(n_ci=self.n_ci, n_down=self.n_per_block, batch_global=self.batch_global) - 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 ci_ppgd_edge(self) -> "BatchEdge": + return BatchEdge(n_ci=self.n_ci, n_down=self.n_ppgd, batch_global=self.batch_global) + + +@dataclass(frozen=True) +class BatchEdge: + """Symmetric batch-slice geometry for one cross-pool edge (CI ↔ downstream). + + ``n_down`` is the downstream pool's batch arity (``n_per_block`` for LW, + ``n_ppgd`` for PPGD). The validator guarantees ``n_ci`` and ``n_down`` are + cross-divisible, so exactly one of two regimes holds: + + * ``ci_is_coarse`` (``n_ci <= n_down``, ``n_ci | n_down``): each CI slice + contains ``fanout = n_down // n_ci`` whole downstream slices. CI fans a + sub-slice out to each; grads come back to be stitched. One downstream + rank pairs with exactly one CI rank. + * not ``ci_is_coarse`` (``n_ci > n_down``, ``n_down | n_ci``): each + downstream slice contains ``fanout = n_ci // n_down`` whole CI slices. + One downstream rank gathers from ``fanout`` CI ranks (concat) and + scatters grads back to those same ``fanout`` CI ranks. One CI rank pairs + with exactly one downstream rank. + + All slice arithmetic is in units of the FINER pool's shard, so every overlap + is a whole, aligned sub-slice. + """ + + n_ci: int + n_down: int + batch_global: int + + def __post_init__(self) -> None: + assert self.n_down % self.n_ci == 0 or self.n_ci % self.n_down == 0, ( + f"n_ci ({self.n_ci}) and n_down ({self.n_down}) must be cross-divisible" + ) + assert self.batch_global % self.n_ci == 0 and self.batch_global % self.n_down == 0 + + @property + def ci_is_coarse(self) -> bool: + return self.n_ci <= self.n_down + + @property + def fanout(self) -> int: + """Number of finer-pool slices nested in one coarser-pool slice.""" + return self.n_down // self.n_ci if self.ci_is_coarse else self.n_ci // self.n_down + + @property + def b_ci(self) -> int: + return self.batch_global // self.n_ci @property - def k_ppgd_per_ci(self) -> int: - """How many PPGD batch shards fit inside one CI batch shard. + def b_down(self) -> int: + return self.batch_global // self.n_down - Validator guarantees N_ci | N_ppgd, so this is an integer. + # ── Pairing: which ranks talk to which ── + + def ci_slices_for_down_slice(self, down_slice_idx: int) -> tuple[int, ...]: + """CI slice idxs whose batches overlap downstream slice `down_slice_idx`. + + Coarse-CI regime: exactly one CI slice (the one containing it). + Fine-CI regime: the `fanout` CI slices nested in it. """ - assert self.n_ppgd % self.n_ci == 0 - return self.n_ppgd // self.n_ci + if self.ci_is_coarse: + return (down_slice_idx // self.fanout,) + return tuple(range(down_slice_idx * self.fanout, (down_slice_idx + 1) * self.fanout)) - 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 down_slices_for_ci_slice(self, ci_slice_idx: int) -> tuple[int, ...]: + """Downstream slice idxs whose batches overlap CI slice `ci_slice_idx`. + + Coarse-CI regime: the `fanout` downstream slices nested in it. + Fine-CI regime: exactly one downstream slice (the one containing it). + """ + if self.ci_is_coarse: + return tuple(range(ci_slice_idx * self.fanout, (ci_slice_idx + 1) * self.fanout)) + return (ci_slice_idx // self.fanout,) - 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 + # ── Sub-slices: the overlap, expressed in each rank's local batch tensor ── - 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`. + def overlap_within_ci(self, ci_slice_idx: int, down_slice_idx: int) -> slice: + """The CI↔downstream overlap as a sub-slice of CI rank `ci_slice_idx`'s + local ``[B_ci, ...]`` tensor. - Assumes the caller is the CI rank with - `ci_slice_idx == ci_slice_of_lw_block_rank(block_rank_idx)`. + The overlap spans the FINER pool's shard (``min(b_ci, b_down)`` rows). """ - 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)) + return self._overlap(ci_slice_idx, down_slice_idx, base_is_ci=True) + + def overlap_within_down(self, ci_slice_idx: int, down_slice_idx: int) -> slice: + """The CI↔downstream overlap as a sub-slice of downstream rank + `down_slice_idx`'s local ``[B_down, ...]`` tensor.""" + return self._overlap(ci_slice_idx, down_slice_idx, base_is_ci=False) + + def _overlap(self, ci_slice_idx: int, down_slice_idx: int, *, base_is_ci: bool) -> slice: + ci_global_start = ci_slice_idx * self.b_ci + down_global_start = down_slice_idx * self.b_down + overlap_global_start = max(ci_global_start, down_global_start) + overlap_len = min(self.b_ci, self.b_down) + base_start = ci_global_start if base_is_ci else down_global_start + local_start = overlap_global_start - base_start + assert local_start >= 0, f"non-overlapping slices: ci={ci_slice_idx} down={down_slice_idx}" + return slice(local_start, local_start + overlap_len) def build_world( @@ -535,27 +616,30 @@ def i_lead_site(self, site: str) -> bool: 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 → LW: for each LW rank whose batch shard overlaps my CI slice, + isend the overlapping sub-slice of each of that LW rank's owned sites. ``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]``. + CI fn is global). Values have shape ``[B_local_ci, S, C_s]``. Each send + carries ``overlap_len = min(b_ci, b_lw)`` rows — the whole overlap of + this CI rank with the target LW rank, in either fan direction. Returned buffers must be kept alive until ``work.wait()`` completes. """ assert self.my_pool == "ci" and self.my_ci_slice_idx is not None + edge = self.world.ci_lw_edge + my_down_slices = edge.down_slices_for_ci_slice(self.my_ci_slice_idx) works: list[dist.Work] = [] buffers: list[Tensor] = [] - my_lw_block_ranks = self.world.lw_block_ranks_for_ci_slice(self.my_ci_slice_idx) with _time_nccl_op("async_send_ci_to_layerwise"): for bg in self.world.layerwise_block_groups: - for block_rank_idx in my_lw_block_ranks: + for block_rank_idx in my_down_slices: target_lw_rank = bg.ranks[block_rank_idx] - sub = self.world.lw_sub_slice_within_ci(block_rank_idx) + sub = edge.overlap_within_ci(self.my_ci_slice_idx, block_rank_idx) # Coalesce all of this block's owned-sites into one packed # send per (block, block_rank). Layout (must match recv): - # for each site in bg.owned_sites order, b_lw * seq_len * C_s - # contiguous _WIRE_DTYPE elements. + # for each site in bg.owned_sites order, overlap_len * seq_len + # * C_s contiguous _WIRE_DTYPE elements. parts = [ ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() for site in bg.owned_sites @@ -572,20 +656,21 @@ def async_send_ci_to_layerwise( 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).""" + """CI → PPGD: for each PPGD rank whose batch shard overlaps my CI slice, + isend the full-model CI overlap sub-slice (all sites).""" assert self.my_pool == "ci" and self.my_ci_slice_idx is not None + edge = self.world.ci_ppgd_edge + my_down_slices = edge.down_slices_for_ci_slice(self.my_ci_slice_idx) works: list[dist.Work] = [] buffers: list[Tensor] = [] - my_ppgd_slice_idxs = self.world.ppgd_slice_idxs_for_ci_slice(self.my_ci_slice_idx) with _time_nccl_op("async_send_ci_to_ppgd"): - for ppgd_slice_idx in my_ppgd_slice_idxs: + for ppgd_slice_idx in my_down_slices: target_ppgd_rank = self.world.ppgd_ranks[ppgd_slice_idx] - sub = self.world.ppgd_sub_slice_within_ci(ppgd_slice_idx) - # Coalesce all 96 sites into one packed send per PPGD target. + sub = edge.overlap_within_ci(self.my_ci_slice_idx, ppgd_slice_idx) + # Coalesce all sites into one packed send per PPGD target. # Layout (must match recv): for each site in self.world.all_sites - # order, b_pp * seq_len * C_s contiguous _WIRE_DTYPE elements. + # order, overlap_len * seq_len * C_s contiguous _WIRE_DTYPE elements. parts = [ ci_full[site][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() for site in self.world.all_sites @@ -608,16 +693,17 @@ def send_ci_eval_to_ppgd(self, ci: CIOutputs) -> None: Pack layout per send (must match ``recv_ci_eval_from_ci_pool``): three contiguous blocks in order (lower_leaky, upper_leaky, pre_sigmoid). Each - block has, for each site in ``self.world.all_sites`` order, ``b_pp * - seq_len * C_s`` contiguous ``_WIRE_DTYPE`` elements. + block has, for each site in ``self.world.all_sites`` order, + ``overlap_len * 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) + edge = self.world.ci_ppgd_edge + my_down_slices = edge.down_slices_for_ci_slice(self.my_ci_slice_idx) with _time_nccl_op("send_ci_eval_to_ppgd"): - for ppgd_slice_idx in my_ppgd_slice_idxs: + for ppgd_slice_idx in my_down_slices: target = self.world.ppgd_ranks[ppgd_slice_idx] - sub = self.world.ppgd_sub_slice_within_ci(ppgd_slice_idx) + sub = edge.overlap_within_ci(self.my_ci_slice_idx, ppgd_slice_idx) parts: list[Tensor] = [] for d in (ci.lower_leaky, ci.upper_leaky, ci.pre_sigmoid): parts.extend( @@ -633,31 +719,38 @@ def recv_g_ci_from_layerwise( 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. + """CI ← LW: recv per-site CI grads, coalesced per overlapping LW rank. 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. + ``send_g_ci_to_ci_pool``); we receive one packed buf per overlapping LW + rank. Each carries ``overlap_len = min(b_ci, b_lw)`` rows — the overlap + of this CI slice with that LW rank. Pack layout (must match sender): for + each site in the LW block's owned sites, ``overlap_len * seq_len * c_s`` + contiguous ``_WIRE_DTYPE`` elements. + + When CI is finer than LW (one LW rank spans several CI slices), each LW + rank still sends this CI rank exactly its own overlap; the recv fills + the whole ``[b_ci]`` dest. When CI is coarser, several LW ranks each fill + a disjoint sub-slice of the dest. """ 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 + edge = self.world.ci_lw_edge + my_down_slices = edge.down_slices_for_ci_slice(self.my_ci_slice_idx) + overlap_len = min(edge.b_ci, edge.b_down) # 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, ...]]] = [] + pending: list[tuple[int, Tensor, dist.Work, tuple[str, ...]]] = [] with _time_nccl_op("recv_g_ci_from_layerwise:post_irecvs"): - for bg_idx, bg in enumerate(self.world.layerwise_block_groups): + for bg in 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: + packed_numel = sum(overlap_len * seq_len * site_to_c[s] for s in owned) + for block_rank_idx in my_down_slices: src = bg.ranks[block_rank_idx] buf = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) w = dist.irecv(buf, src=src, group=self.world.cross_pool_p2p_group) assert w is not None - pending.append((bg_idx, block_rank_idx, buf, w, owned)) + pending.append((block_rank_idx, buf, w, owned)) # Wait + stitch. Allocate one fp32 dest per site, copy each piece in place. out: dict[str, Tensor] = {} @@ -666,14 +759,14 @@ def recv_g_ci_from_layerwise( c_s = site_to_c[site] out[site] = torch.empty(b_ci, seq_len, c_s, device=device, dtype=torch.float32) with _time_nccl_op("recv_g_ci_from_layerwise:wait"): - for _bg_idx, block_rank_idx, buf, w, owned in pending: + for block_rank_idx, buf, w, owned in pending: w.wait() - sub = self.world.lw_sub_slice_within_ci(block_rank_idx) + sub = edge.overlap_within_ci(self.my_ci_slice_idx, 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) + n = overlap_len * seq_len * c_s + site_view = buf[offset : offset + n].view(overlap_len, seq_len, c_s) out[site][sub].copy_(site_view.to(torch.float32)) offset += n return out @@ -686,26 +779,27 @@ def recv_g_ci_from_ppgd( ) -> 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. + One packed irecv per overlapping PPGD source (instead of one per (site, + source)), matching the coalesced ``send_g_ci_to_ci_pool_ppgd``. Each + carries ``overlap_len = min(b_ci, b_pp)`` rows — the overlap of this CI + slice with that PPGD rank. 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. + ``self.world.all_sites`` order, ``overlap_len * 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 + edge = self.world.ci_ppgd_edge + my_down_slices = edge.down_slices_for_ci_slice(self.my_ci_slice_idx) + overlap_len = min(edge.b_ci, edge.b_down) # 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} + site_numels = {s: overlap_len * seq_len * site_to_c[s] for s in self.world.all_sites} packed_numel = sum(site_numels.values()) pending: list[tuple[int, Tensor, dist.Work]] = [] with _time_nccl_op("recv_g_ci_from_ppgd:post_irecvs"): - for ppgd_slice_idx in my_ppgd_slice_idxs: + for ppgd_slice_idx in my_down_slices: src = self.world.ppgd_ranks[ppgd_slice_idx] packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) w = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) @@ -720,12 +814,12 @@ def recv_g_ci_from_ppgd( with _time_nccl_op("recv_g_ci_from_ppgd:wait"): for ppgd_slice_idx, packed, w in pending: w.wait() - sub = self.world.ppgd_sub_slice_within_ci(ppgd_slice_idx) + sub = edge.overlap_within_ci(self.my_ci_slice_idx, 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) + buf = packed[offset : offset + n].view(overlap_len, seq_len, c_s) out[site][sub].copy_(buf.to(torch.float32)) offset += n return out @@ -764,53 +858,66 @@ def async_recv_ci_from_ci_pool( seq_len: int, device: torch.device, ) -> PendingCiRecv: - """LW ← CI: irecv one coalesced packet of CI values for all of this - LW rank's owned sites, from the CI rank whose slice contains my LW - batch shard. - - Layout (must match ``async_send_ci_to_layerwise``): for each site in - ``self.my_owned_sites`` order, ``b_lw * seq_len * C_s`` contiguous - ``_WIRE_DTYPE`` elements. Caller calls ``wait_and_unpack()`` to get - per-site ``[b_lw, seq_len, C_s]`` views (no copy). + """LW ← CI: irecv CI values for all of this LW rank's owned sites, from + every CI rank whose slice overlaps my LW batch shard. + + Coarse-CI: one packet from one CI rank, spanning my whole ``[b_lw]`` + slice. Fine-CI: ``fanout`` packets, each from a different CI rank filling + a disjoint sub-slice. Layout per packet (must match + ``async_send_ci_to_layerwise``): for each site in ``self.my_owned_sites`` + order, ``overlap_len * seq_len * C_s`` contiguous ``_WIRE_DTYPE`` + elements. Caller calls ``wait_and_unpack()`` to get per-site + ``[b_lw, seq_len, C_s]`` tensors. """ 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] + edge = self.world.ci_lw_edge + my_idx = self.my_within_block_idx b_lw = self.world.batch_local_lw - packed_numel = sum(b_lw * seq_len * site_to_c[s] for s in self.my_owned_sites) - packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + packets: list[_CiRecvPacket] = [] with _time_nccl_op("async_recv_ci_from_ci_pool"): - work = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) - assert work is not None + for ci_slice_idx in edge.ci_slices_for_down_slice(my_idx): + src = self.world.ci_ranks[ci_slice_idx] + row_slice = edge.overlap_within_down(ci_slice_idx, my_idx) + overlap_len = row_slice.stop - row_slice.start + packed_numel = sum( + overlap_len * seq_len * site_to_c[s] for s in self.my_owned_sites + ) + packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + work = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) + assert work is not None + packets.append(_CiRecvPacket(packed=packed, work=work, row_slice=row_slice)) return PendingCiRecv( - packed=packed, - work=work, + packets=tuple(packets), sites=self.my_owned_sites, site_to_c=site_to_c, - b=b_lw, + b_down=b_lw, seq_len=seq_len, + device=device, ) 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. + """LW → CI: send per-owned-site CI grads to every CI rank whose slice + overlaps my LW batch shard. - 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)). + Coarse-CI: one packed send (my whole ``[b_lw]`` grad) to one CI rank. + Fine-CI: one packed send per overlapping CI rank, each carrying that CI + rank's overlap sub-slice. Coalesces this rank's owned sites into one + packed send per destination (vs one isend per site). """ 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) + edge = self.world.ci_lw_edge + my_idx = self.my_within_block_idx with _time_nccl_op("send_g_ci_to_ci_pool"): - dist.send(packed, dst=dst, group=self.world.cross_pool_p2p_group) + for ci_slice_idx in edge.ci_slices_for_down_slice(my_idx): + dst = self.world.ci_ranks[ci_slice_idx] + sub = edge.overlap_within_down(ci_slice_idx, my_idx) + parts = [ + g_ci_owned[s][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() + for s in self.my_owned_sites + ] + packed = torch.cat(parts) + dist.send(packed, dst=dst, group=self.world.cross_pool_p2p_group) def recv_g_vu_from_ppgd( self, @@ -929,31 +1036,42 @@ def async_recv_ci_from_ci_pool_ppgd( seq_len: int, device: torch.device, ) -> PendingCiRecv: - """PPGD ← CI: irecv one coalesced packet of full-model CI values from - the CI rank that owns my slice. - - Layout (must match ``async_send_ci_to_ppgd``): for each site in - ``self.world.all_sites`` order, ``b_pp * seq_len * C_s`` contiguous - ``_WIRE_DTYPE`` elements. Caller calls ``wait_and_unpack()`` to get - per-site ``[b_pp, seq_len, C_s]`` views (no copy). + """PPGD ← CI: irecv full-model CI values from every CI rank whose slice + overlaps my PPGD batch shard. + + Coarse-CI: one packet from one CI rank, spanning my whole ``[b_pp]`` + slice. Fine-CI: ``fanout`` packets, each from a different CI rank filling + a disjoint sub-slice. Layout per packet (must match + ``async_send_ci_to_ppgd``): for each site in ``self.world.all_sites`` + order, ``overlap_len * seq_len * C_s`` contiguous ``_WIRE_DTYPE`` + elements. Caller calls ``wait_and_unpack()`` to get per-site + ``[b_pp, seq_len, C_s]`` tensors. """ 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] + edge = self.world.ci_ppgd_edge + my_idx = self.my_ppgd_slice_idx b_pp = self.world.batch_local_ppgd - packed_numel = sum(b_pp * seq_len * site_to_c[s] for s in self.world.all_sites) - packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + packets: list[_CiRecvPacket] = [] with _time_nccl_op("async_recv_ci_from_ci_pool_ppgd"): - work = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) - assert work is not None + for ci_slice_idx in edge.ci_slices_for_down_slice(my_idx): + src = self.world.ci_ranks[ci_slice_idx] + row_slice = edge.overlap_within_down(ci_slice_idx, my_idx) + overlap_len = row_slice.stop - row_slice.start + packed_numel = sum( + overlap_len * seq_len * site_to_c[s] for s in self.world.all_sites + ) + packed = torch.empty(packed_numel, device=device, dtype=_WIRE_DTYPE) + work = dist.irecv(packed, src=src, group=self.world.cross_pool_p2p_group) + assert work is not None + packets.append(_CiRecvPacket(packed=packed, work=work, row_slice=row_slice)) return PendingCiRecv( - packed=packed, - work=work, + packets=tuple(packets), sites=self.world.all_sites, site_to_c=site_to_c, - b=b_pp, + b_down=b_pp, seq_len=seq_len, + device=device, ) def recv_ci_eval_from_ci_pool( @@ -962,56 +1080,72 @@ def recv_ci_eval_from_ci_pool( seq_len: int, device: torch.device, ) -> CIOutputs: - """PPGD ← CI eval: synchronous recv of full ``CIOutputs`` from the CI - rank that owns my slice. - - Pack layout (must match ``send_ci_eval_to_ppgd``): three contiguous - blocks in order (lower_leaky, upper_leaky, pre_sigmoid). Each block has, - for each site in ``self.world.all_sites`` order, ``b_pp * seq_len * - C_s`` contiguous ``_WIRE_DTYPE`` elements. Returned dicts are no-copy - views into the packed buffer. + """PPGD ← CI eval: synchronous recv of full ``CIOutputs`` from every CI + rank whose slice overlaps my PPGD batch shard. + + Coarse-CI: one recv from one CI rank, spanning my whole ``[b_pp]`` slice. + Fine-CI: one recv per overlapping CI rank, each filling a disjoint + sub-slice. Pack layout per recv (must match ``send_ci_eval_to_ppgd``): + three contiguous blocks in order (lower_leaky, upper_leaky, pre_sigmoid). + Each block has, for each site in ``self.world.all_sites`` order, + ``overlap_len * seq_len * C_s`` contiguous ``_WIRE_DTYPE`` elements. """ 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] + edge = self.world.ci_ppgd_edge + my_idx = self.my_ppgd_slice_idx b_pp = self.world.batch_local_ppgd - per_block_numel = sum(b_pp * seq_len * site_to_c[s] for s in self.world.all_sites) - packed = torch.empty(3 * per_block_numel, device=device, dtype=_WIRE_DTYPE) + out: list[dict[str, Tensor]] = [ + { + s: torch.empty(b_pp, seq_len, site_to_c[s], device=device, dtype=_WIRE_DTYPE) + for s in self.world.all_sites + } + for _ in range(3) + ] with _time_nccl_op("recv_ci_eval_from_ci_pool"): - dist.recv(packed, src=src, group=self.world.cross_pool_p2p_group) - - out: list[dict[str, Tensor]] = [{}, {}, {}] - offset = 0 - for block_idx in range(3): - for site in self.world.all_sites: - c_s = site_to_c[site] - numel = b_pp * seq_len * c_s - out[block_idx][site] = packed[offset : offset + numel].view(b_pp, seq_len, c_s) - offset += numel - assert offset == packed.numel(), f"unpack mismatch: {offset} of {packed.numel()}" + for ci_slice_idx in edge.ci_slices_for_down_slice(my_idx): + src = self.world.ci_ranks[ci_slice_idx] + row_slice = edge.overlap_within_down(ci_slice_idx, my_idx) + overlap_len = row_slice.stop - row_slice.start + per_block_numel = sum( + overlap_len * seq_len * site_to_c[s] for s in self.world.all_sites + ) + packed = torch.empty(3 * per_block_numel, device=device, dtype=_WIRE_DTYPE) + dist.recv(packed, src=src, group=self.world.cross_pool_p2p_group) + offset = 0 + for block_idx in range(3): + for site in self.world.all_sites: + c_s = site_to_c[site] + numel = overlap_len * seq_len * c_s + view = packed[offset : offset + numel].view(overlap_len, seq_len, c_s) + out[block_idx][site][row_slice].copy_(view) + offset += numel + assert offset == packed.numel(), f"unpack mismatch: {offset} of {packed.numel()}" return CIOutputs(lower_leaky=out[0], upper_leaky=out[1], pre_sigmoid=out[2]) def send_g_ci_to_ci_pool_ppgd(self, g_ci_full: dict[str, Tensor]) -> None: - """PPGD → CI: send full-model CI grads (PPGD batch slice) to the CI - rank that owns my slice. - - Coalesces all 96-ish sites into a single packed buffer per - send. Per-site isends launch ~10ms of NCCL overhead each, so at - scale (96 sites × N_PPGD ranks) this phase was ~1 s of pure NCCL - launch latency on every step. Single packed send replaces that with - one NCCL op (NCCL handles big tensors efficiently). + """PPGD → CI: send full-model CI grads to every CI rank whose slice + overlaps my PPGD batch shard. + + Coarse-CI: one packed send (my whole ``[b_pp]`` grad) to one CI rank. + Fine-CI: one packed send per overlapping CI rank, each carrying that CI + rank's overlap sub-slice. Coalesces all sites into a single packed buffer + per destination — per-site isends launch ~10ms of NCCL overhead each, so + at scale this would be ~1 s of pure launch latency every step. """ 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) + edge = self.world.ci_ppgd_edge + my_idx = self.my_ppgd_slice_idx with _time_nccl_op("send_g_ci_to_ci_pool_ppgd"): - dist.send(packed, dst=dst, group=self.world.cross_pool_p2p_group) + for ci_slice_idx in edge.ci_slices_for_down_slice(my_idx): + dst = self.world.ci_ranks[ci_slice_idx] + sub = edge.overlap_within_down(ci_slice_idx, my_idx) + parts = [ + g_ci_full[s][sub].detach().to(_WIRE_DTYPE).contiguous().flatten() + for s in self.world.all_sites + ] + packed = torch.cat(parts) + dist.send(packed, dst=dst, group=self.world.cross_pool_p2p_group) def send_g_vu_to_layerwise( self, diff --git a/scripts/hetero_topology_equiv/compare.py b/scripts/hetero_topology_equiv/compare.py new file mode 100644 index 000000000..a1f0adabe --- /dev/null +++ b/scripts/hetero_topology_equiv/compare.py @@ -0,0 +1,127 @@ +"""Compare single-pool vs heterogeneous-3-pool loss curves for the CI-fn-grad diagnostic. + +Reads ``metrics.jsonl`` from each run, aligns the four loss terms per step, and +reports per-seed and seed-averaged final-value ratios (3pool / single-pool). + +The headline question: do imp and/or pgd diverge (3-pool under-optimizing them) +while stoch tracks? That pattern confirms the suspected CI-fn-grad scaling bug; +roughly-tracking curves refute it. + +Single-pool logs ``train/loss/``; 3-pool logs ``train/loss/``. +We map them onto a common term name. +""" + +import json +import sys + +from param_decomp_lab.infra.settings import PARAM_DECOMP_OUT_DIR + +DECOMP = PARAM_DECOMP_OUT_DIR / "decompositions" + +# term -> (single-pool key, 3-pool key) +TERM_KEYS = { + "faith": ("train/loss/FaithfulnessLoss", "train/loss/faith"), + "stoch": ("train/loss/StochasticReconLayerwiseLoss", "train/loss/stoch"), + "imp": ("train/loss/ImportanceMinimalityLoss", "train/loss/imp"), + "ppgd": ("train/loss/PersistentPGDReconLoss", "train/loss/ppgd"), +} + + +def load_curves(run_id: str, key_idx: int) -> dict[str, dict[int, float]]: + """Return ``{term: {step: value}}`` for one run. ``key_idx`` selects the + single-pool (0) or 3-pool (1) key from ``TERM_KEYS``.""" + path = DECOMP / run_id / "metrics.jsonl" + assert path.exists(), f"missing {path}" + rows = [json.loads(line) for line in path.open()] + out: dict[str, dict[int, float]] = {t: {} for t in TERM_KEYS} + for r in rows: + step = r.get("step") + if step is None: + continue + for term, keys in TERM_KEYS.items(): + k = keys[key_idx] + if k in r: + out[term][step] = r[k] + return out + + +def aligned_final(curve: dict[int, float]) -> tuple[int, float]: + last_step = max(curve) + return last_step, curve[last_step] + + +def main(seeds: list[int]) -> None: + per_seed_ratios: dict[str, list[float]] = {t: [] for t in TERM_KEYS} + + for seed in seeds: + sp_id = f"eq-1p-s{seed}" + tp_id = f"eq-3p-s{seed}" + sp = load_curves(sp_id, 0) + tp = load_curves(tp_id, 1) + print(f"\n===== seed {seed} =====") + print(f"{'term':>6} | {'step':>5} | {'single':>14} | {'3pool':>14} | {'3p/1p':>10}") + print("-" * 64) + for term in TERM_KEYS: + if not sp[term] or not tp[term]: + print(f"{term:>6} | (missing data: sp={len(sp[term])} tp={len(tp[term])})") + continue + s_step, s_val = aligned_final(sp[term]) + t_step, t_val = aligned_final(tp[term]) + ratio = t_val / s_val if s_val != 0 else float("nan") + per_seed_ratios[term].append(ratio) + print( + f"{term:>6} | {min(s_step, t_step):>5} | " + f"{s_val:>14.6g} | {t_val:>14.6g} | {ratio:>10.4f}" + ) + + print("\n===== seed-averaged final-value ratio (3pool / single-pool) =====") + print(f"{'term':>6} | {'mean ratio':>12} | {'n seeds':>8}") + print("-" * 34) + for term, ratios in per_seed_ratios.items(): + if not ratios: + print(f"{term:>6} | {'(none)':>12} |") + continue + mean = sum(ratios) / len(ratios) + print(f"{term:>6} | {mean:>12.4f} | {len(ratios):>8}") + + # Step-averaged ratio (less sensitive to single-step noise than the final step) + # plus an early-vs-late drift check: if a grad-scaling bug accumulates, the + # 3p/1p ratio should drift systematically across training, not just jitter. + print("\n===== step-averaged 3p/1p ratio + early/late drift (per seed) =====") + print( + f"{'seed':>4} | {'term':>6} | {'mean(all)':>10} | {'early':>8} | {'late':>8} | {'drift':>8}" + ) + print("-" * 60) + for seed in seeds: + sp = load_curves(f"eq-1p-s{seed}", 0) + tp = load_curves(f"eq-3p-s{seed}", 1) + for term in TERM_KEYS: + steps = sorted(set(sp[term]) & set(tp[term])) + if not steps: + continue + ratios = [tp[term][st] / sp[term][st] for st in steps if sp[term][st] != 0] + half = len(ratios) // 2 + early = sum(ratios[:half]) / max(half, 1) + late = sum(ratios[half:]) / max(len(ratios) - half, 1) + mean_all = sum(ratios) / len(ratios) + print( + f"{seed:>4} | {term:>6} | {mean_all:>10.4f} | " + f"{early:>8.4f} | {late:>8.4f} | {late - early:>+8.4f}" + ) + + # Per-step trajectory for the two diagnostic terms (imp, pgd) and stoch. + print("\n===== per-step trajectories (seed 0) =====") + sp0 = load_curves(f"eq-1p-s{seeds[0]}", 0) + tp0 = load_curves(f"eq-3p-s{seeds[0]}", 1) + for term in ("stoch", "imp", "ppgd"): + print(f"\n-- {term} -- step: 1p -> 3p (ratio)") + steps = sorted(set(sp0[term]) & set(tp0[term])) + for st in steps: + s, t = sp0[term][st], tp0[term][st] + r = t / s if s != 0 else float("nan") + print(f" step {st:>4}: {s:>12.6g} -> {t:>12.6g} ({r:.4f})") + + +if __name__ == "__main__": + seeds = [int(x) for x in sys.argv[1:]] or [0, 1, 2] + main(seeds) diff --git a/scripts/hetero_topology_equiv/singlepool.yaml b/scripts/hetero_topology_equiv/singlepool.yaml new file mode 100644 index 000000000..27ad9df4a --- /dev/null +++ b/scripts/hetero_topology_equiv/singlepool.yaml @@ -0,0 +1,107 @@ +# Single-pool reference for the CI-fn-grad equivalence diagnostic. +# Identical pd/runtime/target/data/cadence to threepool_hetero.yaml but +# `three_pool: null` so it runs param_decomp.optimize.Trainer on 1 GPU. Same +# seed/losses/coeffs/steps/batch — the clean reference the 3-pool is compared to. +pd: + seed: 0 + n_mask_samples: 1 + ci_config: + mode: global + fn_type: global_shared_transformer + simple_transformer_ci_cfg: + d_model: 256 + n_blocks: 2 + mlp_hidden_dim: + - 1024 + attn_config: + n_heads: 4 + max_len: 128 + rope_base: 10000.0 + sampling: continuous + sigmoid_type: leaky_hard + decomposition_targets: + - module_pattern: h.0.attn.q_proj + C: 64 + - module_pattern: h.0.attn.k_proj + C: 64 + - module_pattern: h.1.attn.q_proj + C: 64 + - module_pattern: h.1.attn.k_proj + C: 64 + identity_decomposition_targets: null + use_delta_component: true + batch_size: 8 + steps: 100 + components_optimizer: + lr_schedule: + start_val: 5.0e-05 + warmup_pct: 0.0 + final_val_frac: 0.1 + fn_type: cosine + grad_clip_norm: 0.01 + ci_fn_optimizer: + lr_schedule: + start_val: 1.0e-04 + warmup_pct: 0.0 + final_val_frac: 0.1 + fn_type: cosine + faithfulness_warmup_steps: 0 + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - type: FaithfulnessLoss + coeff: 1.0e+08 + - type: StochasticReconLayerwiseLoss + coeff: 50.0 + - type: ImportanceMinimalityLoss + coeff: 4.0e-05 + pnorm: 2.0 + beta: 0.5 + p_anneal_start_frac: 0.0 + p_anneal_final_p: 0.3 + p_anneal_end_frac: 1.0 + eps: 1.0e-12 + - type: PersistentPGDReconLoss + coeff: 1.0 + optimizer: + type: adam + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + start_val: 0.01 + warmup_pct: 0.025 + final_val_frac: 1.0 + fn_type: constant + scope: + type: per_batch_per_position + use_sigmoid_parameterization: false + n_warmup_steps: 2 + n_samples: 1 +cadence: + train_log_every: 5 + save_every: null +eval: null +runtime: + autocast_bf16: true + device: cuda + dp: null +target: + spec: + kind: hf_weights_in_vendored + model_class: param_decomp_lab.experiments.lm.pretrain.models.gpt2_simple.GPT2Simple + model_name: openai-community/gpt2 + output_extract: 0 + activation_checkpointing: false +data: + tokenizer_name: openai-community/gpt2 + max_seq_len: 128 + dataset_name: apollo-research/Skylion007-openwebtext-tokenizer-gpt2 + column_name: input_ids + train_split: train + eval_split: train + is_tokenized: true + streaming: true + buffer_size: 1000 + shuffle_each_epoch: true +three_pool: null diff --git a/scripts/hetero_topology_equiv/submit_threepool_topo.sh b/scripts/hetero_topology_equiv/submit_threepool_topo.sh new file mode 100755 index 000000000..b16ecd04e --- /dev/null +++ b/scripts/hetero_topology_equiv/submit_threepool_topo.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Submit one 8-GPU single-node 3-pool run for a given topology yaml. +# Usage: submit_threepool_topo.sh +set -euo pipefail +YAML="$1" +RUN_ID="$2" +PORT="$3" +WORKTREE="/mnt/home/oli/param-decomp/.claude/worktrees/agent-ab6947e4c08bc68ae" + +sbatch < n_per_block=2: old `n_per_block % n_ci` (2%4) FAILED -> needs relaxation. + # n_ci=4 > n_ppgd=2: old `n_ppgd % n_ci` (2%4) FAILED -> needs relaxation. + # Both edges run the inverted (CI-fine) routing path. batch_size=8 % lcm(4,2,2)=4. + ci_ranks: [2, 3, 4, 5] + ppgd_ranks: [6, 7] + layerwise_block_groups: + - ranks: [0, 1] + owned_sites: + - h.0.attn.q_proj + - h.0.attn.k_proj + - h.1.attn.q_proj + - h.1.attn.k_proj diff --git a/scripts/hetero_topology_equiv/threepool_inverted_s1.yaml b/scripts/hetero_topology_equiv/threepool_inverted_s1.yaml new file mode 100644 index 000000000..0585353af --- /dev/null +++ b/scripts/hetero_topology_equiv/threepool_inverted_s1.yaml @@ -0,0 +1,134 @@ +# Heterogeneous-topology 3-pool config for the CI-fn-grad equivalence diagnostic. +# +# Topology (8 ranks, single node): +# LW: 1 block × 4 ranks = ranks [0,1,2,3] (owns all 4 sites) +# CI: ranks [4,5] (2 ranks) +# PPGD: ranks [6,7] (2 ranks) +# Rank 0 must be the LW block-0 leader (3-pool convention; reductions ship +# CI/PPGD losses to rank 0), so LW comes first. +# n_ci=2 | n_per_block=4 (4%2=0 OK); n_ci=2 | n_ppgd=2 (2%2=0 OK). +# DELIBERATELY NON-SQUARE: n_per_block=4 != n_ci=2 — this is what exposes the +# suspected CI-fn-grad scaling bug (cancels only when n_ci==n_per_block==n_ppgd). +# batch_size=8 divisible by lcm(2,4,2)=4. +# +# Decomposes 4 small attn projection sites of gpt2. Short run, dense logging. +pd: + seed: 1 + n_mask_samples: 1 + ci_config: + mode: global + fn_type: global_shared_transformer + simple_transformer_ci_cfg: + d_model: 256 + n_blocks: 2 + mlp_hidden_dim: + - 1024 + attn_config: + n_heads: 4 + max_len: 128 + rope_base: 10000.0 + sampling: continuous + sigmoid_type: leaky_hard + decomposition_targets: + - module_pattern: h.0.attn.q_proj + C: 64 + - module_pattern: h.0.attn.k_proj + C: 64 + - module_pattern: h.1.attn.q_proj + C: 64 + - module_pattern: h.1.attn.k_proj + C: 64 + identity_decomposition_targets: null + use_delta_component: true + batch_size: 8 + steps: 100 + components_optimizer: + lr_schedule: + start_val: 5.0e-05 + warmup_pct: 0.0 + final_val_frac: 0.1 + fn_type: cosine + grad_clip_norm: 0.01 + ci_fn_optimizer: + lr_schedule: + start_val: 1.0e-04 + warmup_pct: 0.0 + final_val_frac: 0.1 + fn_type: cosine + faithfulness_warmup_steps: 0 + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - type: FaithfulnessLoss + coeff: 1.0e+08 + - type: StochasticReconLayerwiseLoss + coeff: 50.0 + - type: ImportanceMinimalityLoss + coeff: 4.0e-05 + pnorm: 2.0 + beta: 0.5 + p_anneal_start_frac: 0.0 + p_anneal_final_p: 0.3 + p_anneal_end_frac: 1.0 + eps: 1.0e-12 + - type: PersistentPGDReconLoss + coeff: 1.0 + optimizer: + type: adam + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + start_val: 0.01 + warmup_pct: 0.025 + final_val_frac: 1.0 + fn_type: constant + scope: + type: per_batch_per_position + use_sigmoid_parameterization: false + n_warmup_steps: 2 + n_samples: 1 +cadence: + train_log_every: 5 + save_every: null +eval: null +runtime: + autocast_bf16: true + device: cuda + dp: null +target: + spec: + kind: hf_weights_in_vendored + model_class: param_decomp_lab.experiments.lm.pretrain.models.gpt2_simple.GPT2Simple + model_name: openai-community/gpt2 + output_extract: 0 + activation_checkpointing: false +data: + tokenizer_name: openai-community/gpt2 + max_seq_len: 128 + dataset_name: apollo-research/Skylion007-openwebtext-tokenizer-gpt2 + column_name: input_ids + train_split: train + eval_split: train + is_tokenized: true + streaming: true + buffer_size: 1000 + shuffle_each_epoch: true +three_pool: + use_fused_kl: true + # RELAXATION-REQUIRING topology (CI finer than both downstream pools): + # LW: 1 block x 2 ranks = ranks [0,1] (owns all 4 sites) + # CI: ranks [2,3,4,5] (4 ranks) + # PPGD: ranks [6,7] (2 ranks) + # n_ci=4 > n_per_block=2: old `n_per_block % n_ci` (2%4) FAILED -> needs relaxation. + # n_ci=4 > n_ppgd=2: old `n_ppgd % n_ci` (2%4) FAILED -> needs relaxation. + # Both edges run the inverted (CI-fine) routing path. batch_size=8 % lcm(4,2,2)=4. + ci_ranks: [2, 3, 4, 5] + ppgd_ranks: [6, 7] + layerwise_block_groups: + - ranks: [0, 1] + owned_sites: + - h.0.attn.q_proj + - h.0.attn.k_proj + - h.1.attn.q_proj + - h.1.attn.k_proj diff --git a/scripts/hetero_topology_equiv/threepool_legal.yaml b/scripts/hetero_topology_equiv/threepool_legal.yaml new file mode 100644 index 000000000..7abfc041e --- /dev/null +++ b/scripts/hetero_topology_equiv/threepool_legal.yaml @@ -0,0 +1,132 @@ +# Heterogeneous-topology 3-pool config for the CI-fn-grad equivalence diagnostic. +# +# Topology (8 ranks, single node): +# LW: 1 block × 4 ranks = ranks [0,1,2,3] (owns all 4 sites) +# CI: ranks [4,5] (2 ranks) +# PPGD: ranks [6,7] (2 ranks) +# Rank 0 must be the LW block-0 leader (3-pool convention; reductions ship +# CI/PPGD losses to rank 0), so LW comes first. +# n_ci=2 | n_per_block=4 (4%2=0 OK); n_ci=2 | n_ppgd=2 (2%2=0 OK). +# DELIBERATELY NON-SQUARE: n_per_block=4 != n_ci=2 — this is what exposes the +# suspected CI-fn-grad scaling bug (cancels only when n_ci==n_per_block==n_ppgd). +# batch_size=8 divisible by lcm(2,4,2)=4. +# +# Decomposes 4 small attn projection sites of gpt2. Short run, dense logging. +pd: + seed: 0 + n_mask_samples: 1 + ci_config: + mode: global + fn_type: global_shared_transformer + simple_transformer_ci_cfg: + d_model: 256 + n_blocks: 2 + mlp_hidden_dim: + - 1024 + attn_config: + n_heads: 4 + max_len: 128 + rope_base: 10000.0 + sampling: continuous + sigmoid_type: leaky_hard + decomposition_targets: + - module_pattern: h.0.attn.q_proj + C: 64 + - module_pattern: h.0.attn.k_proj + C: 64 + - module_pattern: h.1.attn.q_proj + C: 64 + - module_pattern: h.1.attn.k_proj + C: 64 + identity_decomposition_targets: null + use_delta_component: true + batch_size: 8 + steps: 100 + components_optimizer: + lr_schedule: + start_val: 5.0e-05 + warmup_pct: 0.0 + final_val_frac: 0.1 + fn_type: cosine + grad_clip_norm: 0.01 + ci_fn_optimizer: + lr_schedule: + start_val: 1.0e-04 + warmup_pct: 0.0 + final_val_frac: 0.1 + fn_type: cosine + faithfulness_warmup_steps: 0 + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - type: FaithfulnessLoss + coeff: 1.0e+08 + - type: StochasticReconLayerwiseLoss + coeff: 50.0 + - type: ImportanceMinimalityLoss + coeff: 4.0e-05 + pnorm: 2.0 + beta: 0.5 + p_anneal_start_frac: 0.0 + p_anneal_final_p: 0.3 + p_anneal_end_frac: 1.0 + eps: 1.0e-12 + - type: PersistentPGDReconLoss + coeff: 1.0 + optimizer: + type: adam + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + start_val: 0.01 + warmup_pct: 0.025 + final_val_frac: 1.0 + fn_type: constant + scope: + type: per_batch_per_position + use_sigmoid_parameterization: false + n_warmup_steps: 2 + n_samples: 1 +cadence: + train_log_every: 5 + save_every: null +eval: null +runtime: + autocast_bf16: true + device: cuda + dp: null +target: + spec: + kind: hf_weights_in_vendored + model_class: param_decomp_lab.experiments.lm.pretrain.models.gpt2_simple.GPT2Simple + model_name: openai-community/gpt2 + output_extract: 0 + activation_checkpointing: false +data: + tokenizer_name: openai-community/gpt2 + max_seq_len: 128 + dataset_name: apollo-research/Skylion007-openwebtext-tokenizer-gpt2 + column_name: input_ids + train_split: train + eval_split: train + is_tokenized: true + streaming: true + buffer_size: 1000 + shuffle_each_epoch: true +three_pool: + use_fused_kl: true + # CURRENT-LEGAL topology (CI coarser; legacy coarse-CI routing path): + # LW: 1 block x 4 ranks = ranks [0,1,2,3] (owns all 4 sites) + # CI: ranks [4,5] (2 ranks) + # PPGD: ranks [6,7] (2 ranks) + # n_ci=2 | n_per_block=4 (forward); n_ci=2 | n_ppgd=2 (square). No relaxation needed. + ci_ranks: [4, 5] + ppgd_ranks: [6, 7] + layerwise_block_groups: + - ranks: [0, 1, 2, 3] + owned_sites: + - h.0.attn.q_proj + - h.0.attn.k_proj + - h.1.attn.q_proj + - h.1.attn.k_proj diff --git a/scripts/validate_nccl_event_timing.py b/scripts/validate_nccl_event_timing.py index 23b5cd530..08548dd30 100644 --- a/scripts/validate_nccl_event_timing.py +++ b/scripts/validate_nccl_event_timing.py @@ -16,7 +16,6 @@ import torch import torch.distributed as dist - from param_decomp.three_pool import layout as L PAYLOAD_NUMEL = 64 * 1024 * 1024 # 256 MB fp32 — big enough to time transfer