From 5c4880d099d0c0d71d5333dc39d601f33317c642 Mon Sep 17 00:00:00 2001 From: oli Date: Fri, 3 Jul 2026 14:01:57 +0000 Subject: [PATCH 1/2] fix(invariance_check): feed int tokens, not a float residual batch (SPEC D4) The check predated the S18 amendment that made the model input token ids (Int[B,T] -> embed_tokens indexing); its float32 'resid' batch crashed clean_output with 'Indexer must have integer or boolean type'. Draw random int tokens like the llama8b tests do. Same failure family as the slow-eval hidden_acts crash fixed by #919. Verified: XLA_FLAGS=--xla_force_host_platform_device_count=4 passes (worst rel 1.02e-04 within tol). Co-Authored-By: Claude Fable 5 --- param_decomp/experiments/invariance_check.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/param_decomp/experiments/invariance_check.py b/param_decomp/experiments/invariance_check.py index e30bf712f..3431a326a 100644 --- a/param_decomp/experiments/invariance_check.py +++ b/param_decomp/experiments/invariance_check.py @@ -77,11 +77,11 @@ def _run(steps: int, sharded: bool) -> list[dict[str, float]]: src = init_persistent_sources( lm.site_names, tuple(s.C for s in lm.sites), (1, seq), jnp.float32, random.PRNGKey(3) ) - resid = random.normal(random.PRNGKey(4), (gbatch, seq, cfg.n_embd)) * 0.5 + tokens = random.randint(random.PRNGKey(4), (gbatch, seq), 0, cfg.vocab_size) mesh = hsdp_mesh() if sharded else None if mesh is not None: - resid = shard_batch(resid, mesh, batch_axis=0) + tokens = shard_batch(tokens, mesh, batch_axis=0) ppgd_cfg = PersistentPGDReconLossConfig( coeff=0.5, @@ -131,7 +131,7 @@ def _run(steps: int, sharded: bool) -> list[dict[str, float]]: out = [] for i in range(steps): - state, m = step(lm, state, resid, random.PRNGKey(100 + i)) + state, m = step(lm, state, tokens, random.PRNGKey(100 + i)) out.append({k: float(v) for k, v in m.items()}) return out From 227e8ef9a9234457d1fe9d90f66c95769c97fa3d Mon Sep 17 00:00:00 2001 From: Oliver Clive-Griffin Date: Wed, 15 Jul 2026 04:54:22 +0000 Subject: [PATCH 2/2] refactor(components): owner-partitioned stacked persistence for the V/U masters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DecompVU persists as same-shape stacks — one (Vs [g,d_in,C], Us [g,C,d_out]) pair per shape group + a static site->(shape,slot) index — instead of 2*n_sites per-site leaves. Placement (hybrid HSDP rule): stack axis /replicate (whole matrices owned per node-group: zero cross-node weight collectives per step, muon NS node-local), matrix d dims /fsdp, C /tp — total /N, same memory as the retired intra-matrix ZeRO-1; per-group fallback to intra-matrix sharding when a stack doesn't tile replicate (site subsets). SPEC D4 amended 2026-07-15 (Oli-requested). - llama8b _stack_per_kind_vu: contiguous-slot fast path = a static slice of the resting stack (the per-step 32-way restack disappears for full-model runs); _per_kind_dims reads shapes statically off site_slots. - init_decomp_vu_placed: one jit, 2*n_shapes outputs (the stack/unstack fan-out and its transient copy are gone). - muon: the V/U tree is all-3D now, so BOTH optimizer groups label leaves via stacked_muon_dimension_numbers (renamed from chunk_stacked_*; optax's default 2D rule would silently Adam every V/U leaf). - decomp_vu_from_sites: explicit per-site-dict constructor (toys/tests). - cherry-pick f63dcd5b1 (invariance_check int tokens — the fix was stranded on its bridge branch; harness needed here). Validated: full suite 535 passed + multidevice 5 passed; invariance_check at 4 sim devices OK (worst rel 9.6e-6, reassociation-only — the D4 anchor for the layout change). Known follow-ups: one-off checkpoint layout migration for pre-change runs; grad_norms per-leaf wandb keys now path through stacks. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01SwEM9ZetQw9W4yjH8okwaH --- param_decomp/CLAUDE.md | 10 +- param_decomp/README.md | 1 + param_decomp/SPEC.md | 2 +- param_decomp/components.py | 221 +++++++++++------- param_decomp/run.py | 2 +- param_decomp/run_state.py | 18 +- param_decomp/targets/llama8b.py | 47 ++-- param_decomp/targets/llama8b_sharding.py | 28 +-- .../tests/equivalence/jax_equivalence.py | 6 +- .../stacked_parity/test_stacked_parity.py | 6 +- param_decomp/tests/test_arithmetic_eval.py | 2 +- param_decomp/tests/test_checkpoint.py | 6 +- param_decomp/tests/test_generic_model_io.py | 4 +- param_decomp/tests/test_llama8b.py | 6 +- param_decomp/tests/test_llama_simple_mlp.py | 4 +- param_decomp/tests/test_no_bake_invariant.py | 6 +- param_decomp/tests/test_optim_torch_parity.py | 10 +- param_decomp/tests/test_sharding.py | 38 ++- param_decomp_lab/experiments/lm/run.py | 2 +- param_decomp_lab/experiments/resid_mlp/run.py | 2 +- .../experiments/resid_mlp/test_resid_mlp.py | 2 +- param_decomp_lab/experiments/tms/run.py | 2 +- param_decomp_lab/experiments/tms/test_tms.py | 2 +- param_decomp_lab/tests/test_toy_uv_eval.py | 12 +- 24 files changed, 246 insertions(+), 193 deletions(-) diff --git a/param_decomp/CLAUDE.md b/param_decomp/CLAUDE.md index 0914b2a4d..c59784856 100644 --- a/param_decomp/CLAUDE.md +++ b/param_decomp/CLAUDE.md @@ -151,9 +151,13 @@ site-slot). NOTE: this is the pure-HSDP backup branch — the mesh is `(replicat NO tensor-parallel / Megatron-C axis (`fsdp` = the 8 intra-node NVLink GPUs, `replicate` = across nodes). The CI output C axis is NEVER sharded, so the per-site heads are a layout convenience here (they were load-bearing under the prior TP layout, which sliced a tp-sharded -glued-ΣC head mid-site). **ZeRO-1 ÷N**: the trainable V/U + CI-fn fp32 masters AND their Adam -m/v shard ÷N over the FULL mesh (`("replicate","fsdp")` on V's d_in / U's d_out / the CI fn's -d_model) — the dominant optimizer-state memory scales 1/N, not the fixed 1/fsdp. The bf16 +glued-ΣC head mid-site). **Persistence layouts (÷N)**: the trainable V/U masters AND their +optimizer moments persist as same-shape STACKS (`DecompVU.stacks`, owner-partitioned: stack +axis ÷`replicate` — whole matrices owned per node-group, zero cross-node weight collectives, +muon NS node-local — matrix d dims ÷`fsdp`, C ÷`tp`; SPEC D4 amendment 2026-07-15; per-group +fallback to intra-matrix data sharding when a stack doesn't tile `replicate`). The CI-fn +masters + moments keep intra-matrix ZeRO-1 (`("replicate","fsdp")` on d_model). Either way +the dominant optimizer-state memory scales 1/N, not the fixed 1/fsdp. The bf16 COMPUTE weights are reconstructed to the `fsdp`-sharded (÷fsdp) layout ONCE per step in ENTRY (the cross-`replicate` gather, off the hot path — `llama8b._reconstruct_compute_weights` / `ci_fn._reconstruct_ci_compute_weights` pin `P(None,"fsdp",...)` BEFORE the per-layer / diff --git a/param_decomp/README.md b/param_decomp/README.md index 4d4a11de3..5b8785300 100644 --- a/param_decomp/README.md +++ b/param_decomp/README.md @@ -28,6 +28,7 @@ root with sibling packages `pretrain/` (the in-house target-LM pretrainer) and | `checkpoint.py` | orbax sharded save/resume of `TrainState` (adversary sources + moments included, no full-gather on the loop, SPEC S22) | | `eval.py` | in-loop eval pass: the six CE/KL masking variants + per-site CI-L0 in one jitted step, logged under the torch `EvalLoop` keys (`eval/ce_kl/*`, `eval/l0/*`) — enabled by the optional `eval:` config block | | `slow_eval.py` | LIBRARY for the in-loop slow (plot) tier (SPEC S28, in-loop only — no offline CLI): the `CIHistograms` / `ComponentActivationDensity` / `CIMeanPerComponent` reductions + renders, the config-gated `PermutedCIPlots` / `IdentityCIError` (off the `(T, C)` position CI), the `UVPlots` figure (`render_uv_figure` / `plot_uv_matrices`, shared by the LM in-loop naive-gather path and the toy `toy_uv_eval` cheap path), and the hidden-acts recon scalars. Torch-free numpy/matplotlib; logged under `slow_eval/figures/*` | +| `components.py` | the decomposition representation: `DecompVU` — V/U masters persisted as same-shape STACKS (owner-partitioned, SPEC D4 amendment 2026-07-15), `site(name)` per-site views, `site_out` the decomposed-linear primitive | | `run_state.py` | optimizer + initial-`TrainState` construction from an `ExperimentConfig` (orbax restores onto this reference) | | `tools/` | `convert_llama_simple_mlp_checkpoint.py` (torch venv) — one-off `.pt` → safetensors conversion of the pile pretrain checkpoint; `migrate_c49k_checkpoint.py` — one-off remap of the frozen C49k clone's orbax `TrainState` (legacy `components.{Vg..Ud}` `(1,*,*)` + flat `sources.`) onto the current layout (site-keyed `components.vu`, `sources..`) so a fine-tune can `restore_latest` it | | `sharding.py` | generic GSPMD helpers (`init_distributed`, `dp_mesh`, `replicate`, `shard_batch`) | diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index 000e75398..6ea1421f4 100644 --- a/param_decomp/SPEC.md +++ b/param_decomp/SPEC.md @@ -341,7 +341,7 @@ is global-batch math. This section is the whole answer to "now shard it": | D1 | Per-shard means + averaged grads must compose to §4's global-batch values for faith/stoch/ppgd (uniform shards). For the *shared-scope source* gradient this means: AVG the per-replica source-grads (each is ∂(local-shard mean)/∂sources) — torch's `reduce_source_grads`; under GSPMD it falls out of autodiff of the global-mean loss. Getting this reduction wrong (e.g. SUM over independent per-position sources) was a real historical bug. | | D2 | Imp-min requires the exact global per-component sums *inside* the `log2` (S8) — the one term where mean-of-shard-results ≠ global result (Jensen). The reduction must also be autograd-aware so gradient reaches each shard's CI values. The eval-path imp-min reuses the one `importance_minimality_terms` impl (there is no separate non-autograd eval reduce as in torch), so the value is device-count invariant by D4; guarded directly at 1 vs N devices by `tests/test_imp_min_global_reduction.py`. | | D3 | Shared PPGD sources stay replica-identical per S16: identical init (broadcast or identical seeding), updates computed from the D1 gradient, identical optimizer steps. | -| D4 | Validation property: with global batch + seed fixed, the metric trajectory is invariant to device count up to floating-point reassociation (cross-shard reduction order; observed rel ≤ ~1e-5 on the tiny-target harness, `experiments/invariance_check.py`). JAX's counter-based RNG makes even the stochastic draws identical across layouts. | +| D4 | Validation property: with global batch + seed fixed, the metric trajectory is invariant to device count up to floating-point reassociation (cross-shard reduction order; observed rel ≤ ~1e-5 on the tiny-target harness, `experiments/invariance_check.py`). JAX's counter-based RNG makes even the stochastic draws identical across layouts. **AMENDED 2026-07-15** (Oli-requested, owner-partitioned persistence): the trainable V/U masters (and their optimizer moments) persist as same-shape STACKS — one `(Vs [g, d_in, C], Us [g, C, d_out])` pair per shape group (`components.DecompVU.stacks` + static `site_slots`), placed by the hybrid HSDP rule: stack axis ÷`replicate` (whole matrices owned per node-group — zero cross-node weight collectives per step; muon NS node-local), matrix d dims ÷`fsdp`, C ÷`tp` (total ÷N, same memory as the retired per-site intra-matrix ZeRO-1); a group whose stack length does not tile `replicate` falls back per-group to intra-matrix data sharding behind the stack axis. Same math, different layout — covered by THIS invariant's reassociation tolerance (re-validated on the harness at 4 sim devices). Checkpoint trees change shape (2·n_shapes stack leaves, not 2·n_sites): pre-change checkpoints need a one-off layout migration to restore at tip. Muon leaf labeling: the V/U tree is now all-3D, so both groups use `stacked_muon_dimension_numbers` (optax's default 2D rule would silently Adam every V/U leaf). | --- diff --git a/param_decomp/components.py b/param_decomp/components.py index 491a875ee..daa78af8b 100644 --- a/param_decomp/components.py +++ b/param_decomp/components.py @@ -1,14 +1,17 @@ """The decomposition representation, shared by every target (LM and toy alike). `SiteC` / `SiteSpec` are the per-site shape primitives (config-level name+C, and the -shape-carrying spec); `DecompVU` is the trainable per-site V/U master pytree; -`init_decomp_vu` seeds it; `site_out` is the one decomposed-linear primitive (SPEC §4.1, -`((x@V)*m)@U + (x@Δ)*d`). These are domain-neutral — they depend only on the site shapes -and the V/U/W arrays — so they live here rather than inside `lm.py` (whose `DecomposedModel` -Protocol references `DecompVU`/`SiteSpec`) or any one target. +shape-carrying spec); `DecompVU` is the trainable master pytree, persisted as same-shape +STACKS (owner-partitioned layout); `init_decomp_vu` seeds it; `site_out` is the one +decomposed-linear primitive (SPEC §4.1, `((x@V)*m)@U + (x@Δ)*d`). These are domain-neutral +— they depend only on the site shapes and the V/U/W arrays — so they live here rather than +inside `lm.py` (whose `DecomposedModel` Protocol references `DecompVU`/`SiteSpec`) or any +one target. """ +from collections.abc import Iterator from dataclasses import dataclass +from functools import cache from typing import Generic, TypeVar import equinox as eqx @@ -37,12 +40,6 @@ class SiteSpec: C: int -# The V/U leaf type: `Array` for the real fp32 masters (the default — so bare `DecompVU` -# means `DecompVU[Array]` and no call site needs the parameter), or `NamedSharding` for the -# same-structure placement tree `.shardings` returns for `jax.jit(out_shardings=...)`. -VULeaf = TypeVar("VULeaf", default=Array) - - _FP8_E4M3_MAX = 448.0 # largest finite magnitude of float8_e4m3fn @@ -70,58 +67,15 @@ def dequantize_fp8(q: Array, scale: Array) -> Array: return q.astype(jnp.bfloat16) * scale.astype(jnp.bfloat16) -class DecompVU(eqx.Module, Generic[VULeaf]): - """Per-decomposed-site V `(d_in, C_s)` / U `(C_s, d_out)`, keyed by site name. The leaves - are fp32 master Arrays (`DecompVU[Array]`), or `NamedSharding`s in the placement tree - returned by `.shardings` (`DecompVU[NamedSharding]`) — same pytree structure, sharding - leaves, for `jax.jit(out_shardings=...)`.""" - - vu: dict[str, tuple[VULeaf, VULeaf]] - - def site(self, name: str) -> tuple[VULeaf, VULeaf]: - return self.vu[name] - - def shardings(self: "DecompVU[Array]", mesh: "Mesh") -> "DecompVU[NamedSharding]": - """True ÷N ZeRO-1 PERSISTENCE layout for the STORED masters, split across the data and - TP axes: V `(d_in, C)` shards d_in over `("replicate","fsdp")` and C over `tp`; U - `(C, d_out)` shards C over `tp` and d_out over `("replicate","fsdp")`. So both still - shard ÷(replicate·fsdp·tp) = ÷N total (C now carries the `tp` factor — the Megatron-C - axis). The fp32 masters + their Adam m/v inherit this; the dominant memory term stays - ÷N. `tp = 1` ⇒ C unsharded, identical to the pure-HSDP layout. - - COMPUTE re-pins d to `fsdp` only (C stays on `tp`): the bf16 compute weights are - reconstructed to `P(None, "fsdp", "tp")` ONCE per step in ENTRY (the ÷N→÷fsdp gather - across `replicate`, off the hot path), so the per-layer scan body gathers only the - `fsdp`-sharded d on NVLink — and only HALF the weight, since C is ÷tp. Asserts d tiles - the data axes and C tiles `tp`.""" - data = ("replicate", "fsdp") - shard_V = NamedSharding(mesh, P(data, "tp")) # d_in ÷(rep·fsdp), C ÷tp → ÷N - shard_U = NamedSharding(mesh, P("tp", data)) # C ÷tp, d_out ÷(rep·fsdp) → ÷N - n_data = mesh.shape["replicate"] * mesh.shape["fsdp"] - n_tp = mesh.shape["tp"] - placed: dict[str, tuple[NamedSharding, NamedSharding]] = {} - for name, (V, U) in self.vu.items(): - assert V.shape[0] % n_data == 0, f"DecompVU[{name}].V.d_in {V.shape[0]} not ÷ {n_data}" - assert V.shape[1] % n_tp == 0, f"DecompVU[{name}].V.C {V.shape[1]} not ÷ tp={n_tp}" - assert U.shape[1] % n_data == 0, f"DecompVU[{name}].U.d_out {U.shape[1]} not ÷ {n_data}" - placed[name] = (shard_V, shard_U) - return DecompVU(vu=placed) - - -def init_decomp_vu(sites: tuple[SiteSpec, ...], key: Array) -> DecompVU: - """Small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5) per site; the - weight-delta channel carries the faithfulness residual at init (before - faithfulness warmup).""" - keys = jax.random.split(key, 2 * len(sites)) - vu: dict[str, tuple[Array, Array]] = {} - for site_idx, spec in enumerate(sites): - V = jax.random.normal(keys[2 * site_idx], (spec.d_in, spec.C)) * spec.d_in**-0.5 - U = jax.random.normal(keys[2 * site_idx + 1], (spec.C, spec.d_out)) * spec.C**-0.5 - vu[spec.name] = (V, U) - return DecompVU(vu=vu) +VUShape = tuple[int, int, int] # (d_in, d_out, C) +# site name -> (shape group, slot on the group's stack axis); static, canonical site order +SiteSlots = tuple[tuple[str, VUShape, int], ...] -VUShape = tuple[int, int, int] # (d_in, d_out, C) +# The V/U leaf type: `Array` for the real fp32 masters (the default — so bare `DecompVU` +# means `DecompVU[Array]` and no call site needs the parameter), or `NamedSharding` for the +# same-structure placement tree `.shardings` returns for `jax.jit(out_shardings=...)`. +VULeaf = TypeVar("VULeaf", default=Array) def vu_shape_groups(sites: tuple[SiteSpec, ...]) -> dict[VUShape, tuple[SiteSpec, ...]]: @@ -132,14 +86,93 @@ def vu_shape_groups(sites: tuple[SiteSpec, ...]) -> dict[VUShape, tuple[SiteSpec return {shape: tuple(specs) for shape, specs in groups.items()} +def site_slots_for(sites: tuple[SiteSpec, ...]) -> SiteSlots: + """The canonical site→(shape, slot) mapping: slots follow site order within each shape + group (`vu_shape_groups` preserves it), entries follow overall site order.""" + by_name: dict[str, tuple[VUShape, int]] = {} + for shape, specs in vu_shape_groups(sites).items(): + for slot, spec in enumerate(specs): + by_name[spec.name] = (shape, slot) + return tuple((spec.name, *by_name[spec.name]) for spec in sites) + + +@cache +def _slot_index(site_slots: SiteSlots) -> dict[str, tuple[VUShape, int]]: + return {name: (shape, slot) for name, shape, slot in site_slots} + + +class DecompVU(eqx.Module, Generic[VULeaf]): + """The trainable V/U masters, persisted as same-shape STACKS — the owner-partitioned + layout: one `(Vs [g, d_in, C], Us [g, C, d_out])` pair per `(d_in, d_out, C)` shape + group, with `site_slots` (static) mapping each site name to its slot on the stack axis. + + Why stacks: per-matrix ownership is only expressible under SPMD by sharding a STACK + axis (a per-site leaf cannot live wholly on one rank), the llama8b scan consumes + per-kind stacks natively, and the checkpoint/init trees shrink from `2·n_sites` leaves + to `2·n_shapes`. Per-site access is `site(name)` — a static slice of the stack. + + Leaves are fp32 master Arrays (`DecompVU[Array]`) or `NamedSharding`s in the placement + tree `.shardings` returns (`DecompVU[NamedSharding]` — same pytree structure).""" + + stacks: dict[VUShape, tuple[VULeaf, VULeaf]] + site_slots: SiteSlots = eqx.field(static=True) + + def site(self: "DecompVU[Array]", name: str) -> tuple[Array, Array]: + shape, slot = _slot_index(self.site_slots)[name] + Vs, Us = self.stacks[shape] + return Vs[slot], Us[slot] + + @property + def site_names(self) -> tuple[str, ...]: + return tuple(name for name, _, _ in self.site_slots) + + def sites_items(self: "DecompVU[Array]") -> Iterator[tuple[str, tuple[Array, Array]]]: + """`(name, (V, U))` in canonical site order — the per-site view of the stacks.""" + for name, _, _ in self.site_slots: + yield name, self.site(name) + + def shardings(self: "DecompVU[Array]", mesh: Mesh) -> "DecompVU[NamedSharding]": + """Owner-partitioned ÷N persistence (hybrid HSDP rule): the STACK axis shards over + `replicate` (whole matrices owned per node-group — the cross-node axis carries ZERO + per-step weight collectives; muon NS on the masters is node-local), matrix d dims + over `fsdp` (NVLink-cheap), C over `tp`. Total ÷(replicate·fsdp·tp) = ÷N — the same + memory as the retired intra-matrix ZeRO-1 layout. + + A group whose stack length does not tile `replicate` (site subsets, e.g. an + L18-only decomposition at multi-node dp) falls back PER-GROUP to intra-matrix data + sharding (`P(None, ("replicate","fsdp"), …)`) — the old layout behind a leading + stack axis. At `replicate == 1` (single node) the two rules coincide.""" + n_rep = mesh.shape["replicate"] + n_fsdp = mesh.shape["fsdp"] + n_tp = mesh.shape["tp"] + placed: dict[VUShape, tuple[NamedSharding, NamedSharding]] = {} + for (d_in, d_out, c), (Vs, _) in self.stacks.items(): + g = Vs.shape[0] + assert c % n_tp == 0, f"C {c} not ÷ tp={n_tp}" + if g % n_rep == 0: + assert d_in % n_fsdp == 0, f"d_in {d_in} not ÷ fsdp={n_fsdp}" + assert d_out % n_fsdp == 0, f"d_out {d_out} not ÷ fsdp={n_fsdp}" + shard_V = NamedSharding(mesh, P("replicate", "fsdp", "tp")) + shard_U = NamedSharding(mesh, P("replicate", "tp", "fsdp")) + else: + data = ("replicate", "fsdp") + n_data = n_rep * n_fsdp + assert d_in % n_data == 0, f"d_in {d_in} not ÷ {n_data}" + assert d_out % n_data == 0, f"d_out {d_out} not ÷ {n_data}" + shard_V = NamedSharding(mesh, P(None, data, "tp")) + shard_U = NamedSharding(mesh, P(None, "tp", data)) + placed[(d_in, d_out, c)] = (shard_V, shard_U) + return DecompVU(stacks=placed, site_slots=self.site_slots) + + def init_decomp_vu_stacked( sites: tuple[SiteSpec, ...], key: Array ) -> dict[VUShape, tuple[Array, Array]]: - """`init_decomp_vu` computed as same-shape stacks: `{(d_in, d_out, C): - (V [n, d_in, C], U [n, C, d_out])}`, vmapped over the SAME per-site keys — so - `unstack_decomp_vu` recovers `init_decomp_vu`'s output BIT-IDENTICALLY (pinned by - `test_sharding`). Under jit the graph has 2×n_shapes sharded outputs instead of - 2×n_sites, which cuts the production init compile ~10× (448 outputs → 14 at 32L).""" + """Seeded stack init: `{(d_in, d_out, C): (V [g, d_in, C], U [g, C, d_out])}`, vmapped + over per-site keys drawn in site order — each site's slice is BIT-IDENTICAL to the + retired per-site init's draw (pinned by `test_sharding`). Under jit the graph has + 2×n_shapes sharded outputs instead of 2×n_sites, which cuts the production init compile + ~10× (448 outputs → 14 at 32L).""" keys = jax.random.split(key, 2 * len(sites)) site_index = {spec.name: idx for idx, spec in enumerate(sites)} stacked: dict[VUShape, tuple[Array, Array]] = {} @@ -151,17 +184,29 @@ def init_decomp_vu_stacked( return stacked -def unstack_decomp_vu( - sites: tuple[SiteSpec, ...], stacked: dict[VUShape, tuple[Array, Array]] -) -> DecompVU: - """Slice `init_decomp_vu_stacked`'s per-shape stacks back into the per-site DecompVU.""" - vu: dict[str, tuple[Array, Array]] = {} - for shape, specs in vu_shape_groups(sites).items(): - Vs, Us = stacked[shape] - assert Vs.shape[0] == len(specs), (Vs.shape, len(specs)) - for j, spec in enumerate(specs): - vu[spec.name] = (Vs[j], Us[j]) - return DecompVU(vu={spec.name: vu[spec.name] for spec in sites}) +def decomp_vu_from_sites(vu: dict[str, tuple[Array, Array]]) -> DecompVU: + """Build the stacked `DecompVU` from a per-site `{name: (V, U)}` dict (site order = + dict order). The explicit-arrays constructor for toys and tests; the trainer inits + directly in the stacked layout (`init_decomp_vu`).""" + sites = tuple( + SiteSpec(name=name, d_in=V.shape[0], d_out=U.shape[1], C=V.shape[1]) + for name, (V, U) in vu.items() + ) + stacks = { + shape: ( + jnp.stack([vu[s.name][0] for s in specs]), + jnp.stack([vu[s.name][1] for s in specs]), + ) + for shape, specs in vu_shape_groups(sites).items() + } + return DecompVU(stacks=stacks, site_slots=site_slots_for(sites)) + + +def init_decomp_vu(sites: tuple[SiteSpec, ...], key: Array) -> DecompVU: + """Small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5) per site, built directly in + the stacked persistence layout; the weight-delta channel carries the faithfulness + residual at init (before faithfulness warmup).""" + return DecompVU(stacks=init_decomp_vu_stacked(sites, key), site_slots=site_slots_for(sites)) def site_out( @@ -181,15 +226,15 @@ def site_out( # Pin the decomposed matmuls DATA-PARALLEL over the FULL mesh: the d_in/d_out-space # activation `x` stays batch-sharded over `('replicate', 'fsdp')`, feature-replicated, and # the component-space activation `x@V` stays batch-sharded, C-REPLICATED (no TP axis). - # This forces the `fsdp`-sharded V/U masters to be GATHERED for compute and their grads - # reduce-scattered back on `fsdp` (symmetric FSDP on NVLink — the intended layout). - # WITHOUT pinning `x`, the weight-grad backward is free to instead shard `x`'s feature - # dim on `fsdp` and REPLICATE the batch (the forward gathers V, the backward does not), - # which GSPMD can't reshard cheaply -> involuntary full rematerialization -> OOM. Pinning - # the activations (not V/U) keeps the weights as plain matmul args. Guarded so it's a - # no-op off-mesh (CPU tests / single device); `run.py` sets the global mesh. waist is - # `[*leading, d]`, leading = (batch, *position): pin batch over the full mesh (positions + - # feature replicated for `x`; C replicated for `x@V`). + # This forces the sharded V/U masters to be GATHERED for compute and their grads + # reduced back to the master layout (symmetric — the intended layout). WITHOUT pinning + # `x`, the weight-grad backward is free to instead shard `x`'s feature dim and REPLICATE + # the batch (the forward gathers V, the backward does not), which GSPMD can't reshard + # cheaply -> involuntary full rematerialization -> OOM. Pinning the activations (not + # V/U) keeps the weights as plain matmul args. Guarded so it's a no-op off-mesh (CPU + # tests / single device); `run.py` sets the global mesh. waist is `[*leading, d]`, + # leading = (batch, *position): pin batch over the full mesh (positions + feature + # replicated for `x`; C replicated for `x@V`). on_mesh = not jax.sharding.get_abstract_mesh().empty batch_axes = ("replicate", "fsdp") if on_mesh: @@ -209,8 +254,8 @@ def site_out( if delta_mask is not None: # `(x @ Δ.T)` for `Δ = W − (V@U).T`, expanded to activation space as # `x@W.T − (x@V)@U` so the `[d_out, d_in]` weight delta is NEVER formed. Under - # FSDP that delta would mix V's dp-sharded d_in with U's dp-sharded d_out (two dims - # demanding the dp axis) and force a replicate-then-repartition reshard; the + # FSDP that delta would mix V's sharded d_in with U's sharded d_out (two dims + # demanding the data axis) and force a replicate-then-repartition reshard; the # activation-space form is all activation×weight matmuls that shard cleanly. # (Still a bf16-rounding DIVERGENCE vs the fp32 oracle delta — accepted; the # faithfulness loss uses the fp32 `weight_deltas`, SPEC N2, not this path.) diff --git a/param_decomp/run.py b/param_decomp/run.py index 52da863ae..9ad21086d 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -560,7 +560,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float: jax.block_until_ready(_ident(tree)) return (time.perf_counter() - _b0) / n - _vu = state.components.vu + _vu = dict(state.components.sites_items()) _by_kind: dict[str, list[tuple[jax.Array, jax.Array]]] = _collections.defaultdict(list) for _name, _VU in _vu.items(): _by_kind[_name.split(".")[-1]].append(_VU) diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index ad4aa355a..1a4b088ed 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -78,12 +78,14 @@ def update( return optax.GradientTransformation(init, update) -def chunk_stacked_muon_dimension_numbers(params: optax.Params) -> optax.Params: - """Muon leaf labeling for the chunkwise CI fn's per-chunk stacks (`ci_fn.py`): every - 3D leaf is a `[n_chunks, d_in, d_out]` stack of `x @ W` matrices — orthogonalize the - trailing two axes, chunk axis batched — and everything else (`[n_chunks, d]` bias - stacks) takes the Adam fallback. optax's default rule (2D → muon) would do the - REVERSE on this tree: NS-orthogonalize the bias stacks and Adam the matrices.""" +def stacked_muon_dimension_numbers(params: optax.Params) -> optax.Params: + """Muon leaf labeling for matrix-STACK trees: every 3D leaf is a `[stack, a, b]` stack + of matrices — orthogonalize the trailing two axes, stack axis batched — and everything + else (e.g. the CI fn's `[n_chunks, d]` bias stacks) takes the Adam fallback. Covers + BOTH optimizer groups now: the chunkwise CI fn's per-chunk stacks (`ci_fn.py`) and the + owner-partitioned V/U shape-group stacks (`components.py` — all leaves 3D, so the + fallback never fires there). optax's default rule (2D → muon) would Adam every V/U + leaf silently and, on the CI tree, NS-orthogonalize the bias stacks instead.""" dims = optax.contrib.MuonDimensionNumbers(reduction_axis=-2, output_axis=-1) return jax.tree.map(lambda leaf: dims if leaf.ndim == 3 else None, params) @@ -144,10 +146,10 @@ def build_optimizers(pd: PDConfig, mesh: Mesh | None): sched_vu = optax_schedule(pd.components_optimizer.lr_schedule, pd.steps) sched_ci = optax_schedule(pd.ci_fn_optimizer.lr_schedule, pd.steps) opt_vu = _optimizer_with_clip( - pd.components_optimizer, sched_vu, muon_dimension_numbers=None, mesh=mesh + pd.components_optimizer, sched_vu, stacked_muon_dimension_numbers, mesh=mesh ) ci_muon_dim_nums = ( - chunk_stacked_muon_dimension_numbers + stacked_muon_dimension_numbers if isinstance(pd.ci_config, ChunkwiseTransformerCiConfig) else None ) diff --git a/param_decomp/targets/llama8b.py b/param_decomp/targets/llama8b.py index 1849df0f3..bc562e8ec 100644 --- a/param_decomp/targets/llama8b.py +++ b/param_decomp/targets/llama8b.py @@ -300,9 +300,9 @@ def _per_kind_dims(components: DecompVU) -> dict[str, tuple[int, int, int]]: uniformity, the precondition for the layer-`lax.scan` masked forward (it stacks each kind across layers, so every layer's matrix of that kind must be the same shape).""" kind_dims: dict[str, tuple[int, int, int]] = {} - for name, (V, U) in components.vu.items(): + for name, (d_in, d_out, c), _slot in components.site_slots: kind = parse_site_name(name)[1] - dims = (V.shape[0], V.shape[1], U.shape[1]) + dims = (d_in, c, d_out) assert kind_dims.setdefault(kind, dims) == dims, ( f"per-kind dims must be uniform across layers for the scan masked forward: " f"{kind} {dims} != {kind_dims[kind]}" @@ -318,22 +318,41 @@ def _stack_per_kind_vu(components: DecompVU, n_layers: int) -> dict[str, dict[st forward in a step, so they are built ONCE via `prepare_compute_weights` and shared. Per-kind dims (d_in, C, d_out) must be uniform across layers (asserted in `_per_kind_dims`).""" kind_dims = _per_kind_dims(components) - vu_dt = next(iter(components.vu.values()))[0].dtype + slot_of = {name: (shape, slot) for name, shape, slot in components.site_slots} + vu_dt = next(iter(components.stacks.values()))[0].dtype per_kind: dict[str, dict[str, Array]] = {} for kind, (d_in, C, d_out) in kind_dims.items(): names = [site_name(layer, kind) for layer in range(n_layers)] - Vs = jnp.stack( - [ - components.vu[n][0] if n in components.vu else jnp.zeros((d_in, C), vu_dt) - for n in names - ] - ) - Us = jnp.stack( - [ - components.vu[n][1] if n in components.vu else jnp.zeros((C, d_out), vu_dt) - for n in names - ] + present = [slot_of[n] for n in names if n in slot_of] + shapes = {shape for shape, _ in present} + slots = [slot for _, slot in present] + contiguous = ( + len(names) == len(present) + and len(shapes) == 1 + and slots == list(range(slots[0], slots[0] + len(slots))) ) + if contiguous: + # Full-kind fast path: one shape group, consecutive slots — the per-kind scan + # stack is a STATIC SLICE of the resting stack (no restack, no zero-fill; the + # owner->compute reshard downstream is the only data movement). + (shape,) = shapes + Vs_all, Us_all = components.stacks[shape] + lo = slots[0] + Vs = Vs_all[lo : lo + len(slots)] + Us = Us_all[lo : lo + len(slots)] + else: + Vs = jnp.stack( + [ + components.site(n)[0] if n in slot_of else jnp.zeros((d_in, C), vu_dt) + for n in names + ] + ) + Us = jnp.stack( + [ + components.site(n)[1] if n in slot_of else jnp.zeros((C, d_out), vu_dt) + for n in names + ] + ) per_kind[kind] = {"V": Vs, "U": Us} return per_kind diff --git a/param_decomp/targets/llama8b_sharding.py b/param_decomp/targets/llama8b_sharding.py index ad446c04d..2ff4cb873 100644 --- a/param_decomp/targets/llama8b_sharding.py +++ b/param_decomp/targets/llama8b_sharding.py @@ -59,9 +59,6 @@ DecompVU, SiteSpec, init_decomp_vu, - init_decomp_vu_stacked, - unstack_decomp_vu, - vu_shape_groups, ) from param_decomp.configs import BSCScope, SCScope from param_decomp.sharding import hsdp_mesh, place_via_shardings @@ -85,25 +82,12 @@ def place_target(tgt: LlamaDecomposedModel, mesh: Mesh) -> LlamaDecomposedModel: def init_decomp_vu_placed(sites: tuple[SiteSpec, ...], key: PRNGKeyArray, mesh: Mesh) -> DecompVU: - """Seeded per-site V/U init placed by `DecompVU.shardings`, bit-identical to - `init_decomp_vu` (pinned by `test_sharding`) but compiled in two cheap stages: the RNG - runs vmap-STACKED per V/U shape (2×n_shapes sharded outputs instead of 2×n_sites — the - SPMD/layout pass over a 448-output RNG graph was a multi-minute compile at 32L), then a - trivial slice jit fans the stacks out to the per-site layout. The stack axis is - unsharded (each trailing spec is the per-site spec behind a leading None). The stacked - copy cannot be buffer-donated into the split outputs, so init peak briefly holds ONE - extra shard-local copy of the params (÷N; a few GB/rank at production dp32, against - an init-time-empty HBM), freed when the fan-out returns.""" - per_site_shardings = eqx.filter_eval_shape(partial(init_decomp_vu, sites), key).shardings(mesh) - stacked_shardings = { - shape: tuple( - NamedSharding(mesh, P(None, *per_site_shardings.vu[specs[0].name][i].spec)) - for i in range(2) - ) - for shape, specs in vu_shape_groups(sites).items() - } - stacked = jax.jit(partial(init_decomp_vu_stacked, sites), out_shardings=stacked_shardings)(key) - return jax.jit(partial(unstack_decomp_vu, sites), out_shardings=per_site_shardings)(stacked) + """Seeded V/U init placed by `DecompVU.shardings` (the owner-partitioned stack layout), + values bit-identical to the retired per-site init (pinned by `test_sharding`). One jit, + 2×n_shapes sharded outputs — the persistence layout IS the stacked layout, so the old + two-stage stack-then-unstack fan-out (and its transient extra copy) is gone.""" + placement = eqx.filter_eval_shape(partial(init_decomp_vu, sites), key).shardings(mesh) + return jax.jit(partial(init_decomp_vu, sites), out_shardings=placement)(key) def init_ci_fn_placed( diff --git a/param_decomp/tests/equivalence/jax_equivalence.py b/param_decomp/tests/equivalence/jax_equivalence.py index cada369c6..9a0cdabc9 100644 --- a/param_decomp/tests/equivalence/jax_equivalence.py +++ b/param_decomp/tests/equivalence/jax_equivalence.py @@ -28,7 +28,7 @@ jax.config.update("jax_enable_x64", False) from param_decomp.adversary import source_masks # noqa: E402 -from param_decomp.components import DecompVU, site_out # noqa: E402 +from param_decomp.components import DecompVU, decomp_vu_from_sites, site_out # noqa: E402 from param_decomp.losses import ( # noqa: E402 faithfulness_loss, importance_minimality_terms, @@ -104,8 +104,8 @@ def _build(f: dict[str, np.ndarray]): ] # inv_freq unused (attn zeroed); a dummy valid-shaped array. inv_freq = jnp.ones((d // 4,), jnp.float32) - vu = DecompVU( - vu={ + vu = decomp_vu_from_sites( + { site_name(i, kind): (a(f"V{kind[0]}_{i}"), a(f"U{kind[0]}_{i}")) for i in range(n_layers) for kind in MLP_KINDS diff --git a/param_decomp/tests/stacked_parity/test_stacked_parity.py b/param_decomp/tests/stacked_parity/test_stacked_parity.py index 4d05ec4d6..0d9049d76 100644 --- a/param_decomp/tests/stacked_parity/test_stacked_parity.py +++ b/param_decomp/tests/stacked_parity/test_stacked_parity.py @@ -33,7 +33,7 @@ init_persistent_sources, init_sources_adam_state, ) -from param_decomp.components import DecompVU +from param_decomp.components import DecompVU, decomp_vu_from_sites from param_decomp.configs import ( AdamPGDConfig, ChunkwiseSubsetReconLossConfig, @@ -120,7 +120,9 @@ def a(key: str) -> jnp.ndarray: layers=layers, norm=a("tgt::norm"), lm_head=a("tgt::lm_head"), inv_freq=llama3_inv_freq(cfg), cfg=cfg, sites=sites, ) # fmt: skip - vu = DecompVU(vu={s.name: (a(f"vu::V::{s.name}"), a(f"vu::U::{s.name}")) for s in sites}) + vu = decomp_vu_from_sites( + {s.name: (a(f"vu::V::{s.name}"), a(f"vu::U::{s.name}")) for s in sites} + ) return f, lm, vu, a("resid") diff --git a/param_decomp/tests/test_arithmetic_eval.py b/param_decomp/tests/test_arithmetic_eval.py index e3fc1a52f..b0a387ae9 100644 --- a/param_decomp/tests/test_arithmetic_eval.py +++ b/param_decomp/tests/test_arithmetic_eval.py @@ -80,7 +80,7 @@ def test_grid_step_ci_xv_and_masked_max_match_hand_rolled(): np.testing.assert_allclose(ci, ci_exp, rtol=1e-4, atol=1e-4) xv = np.asarray(xv_grids[site]) assert xv.shape == (n_pad, C) and np.all(np.isfinite(xv)) - _, u = vu.vu[site] + _, u = vu.site(site) out_got = np.asarray(outputs[site])[:, ANSWER_POSITION, :].astype(np.float32) np.testing.assert_allclose(out_got, xv @ np.asarray(u, np.float32), atol=1e-2) # max CI is over the REAL rows only — the garbage tail must not decide liveness diff --git a/param_decomp/tests/test_checkpoint.py b/param_decomp/tests/test_checkpoint.py index 556663954..63a9b8001 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -41,7 +41,7 @@ from param_decomp.lm import DecomposedModel from param_decomp.muon_stacked import stacked_muon from param_decomp.recon import build_loss_terms -from param_decomp.run_state import chunk_stacked_muon_dimension_numbers +from param_decomp.run_state import stacked_muon_dimension_numbers from param_decomp.schedule import ScheduleConfig from param_decomp.sharding import hsdp_mesh from param_decomp.targets.llama8b import ( @@ -124,7 +124,7 @@ def muon_impl(dim_nums: "Callable[[optax.Params], optax.Params] | None"): inner_vu = muon_impl(None) if muon_components else optax.adamw(1e-3, weight_decay=0.0) opt_vu = optax.chain(optax.clip_by_global_norm(0.01), inner_vu) opt_ci = ( - muon_impl(chunk_stacked_muon_dimension_numbers) + muon_impl(stacked_muon_dimension_numbers) if muon_ci_fn else optax.adamw(1e-3, weight_decay=0.0) ) @@ -210,7 +210,7 @@ def test_muon_roundtrip_and_exact_resume(tmp_path: Path): def test_muon_ci_fn_roundtrip_and_exact_resume(tmp_path: Path): """SPEC S20 amendment (2026-07-11): same guarantee with muon on BOTH groups, the ci-fn - partitioned by `chunk_stacked_muon_dimension_numbers` (3D chunk stacks muon'd, 2D bias + partitioned by `stacked_muon_dimension_numbers` (3D chunk stacks muon'd, 2D bias stacks in the Adam-fallback mask).""" _roundtrip_and_exact_resume(tmp_path, muon_components=True, muon_ci_fn=True) diff --git a/param_decomp/tests/test_generic_model_io.py b/param_decomp/tests/test_generic_model_io.py index 7640df7b3..c951e3920 100644 --- a/param_decomp/tests/test_generic_model_io.py +++ b/param_decomp/tests/test_generic_model_io.py @@ -32,7 +32,7 @@ ChunkwiseTransformerCIArch, build_ci_fn, ) -from param_decomp.components import DecompVU, SiteSpec +from param_decomp.components import DecompVU, SiteSpec, decomp_vu_from_sites from param_decomp.configs import ( FaithfulnessLossConfig, ImportanceMinimalityLossConfig, @@ -181,7 +181,7 @@ def _synthetic_lm(key: jax.Array) -> SyntheticDecomposedModel: def _synthetic_vu(key: jax.Array) -> DecompVU: V = random.normal(random.fold_in(key, 3), (D, C)) * 0.1 U = random.normal(random.fold_in(key, 4), (C, D)) * 0.1 - return DecompVU(vu={SITE: (V, U)}) + return decomp_vu_from_sites({SITE: (V, U)}) def _synthetic_inputs(key: jax.Array) -> dict[str, Array]: diff --git a/param_decomp/tests/test_llama8b.py b/param_decomp/tests/test_llama8b.py index 00fb594b2..489df57af 100644 --- a/param_decomp/tests/test_llama8b.py +++ b/param_decomp/tests/test_llama8b.py @@ -208,7 +208,7 @@ def test_masked_component_activations_pre_mask_and_matches_outputs(): for s in names: assert acts[s].shape == (b, t, C) assert jnp.all(jnp.isfinite(acts[s])) - _, u = vu.vu[s] + _, u = vu.site(s) expected = acts[s].astype(jnp.float32) @ u.astype(jnp.float32) assert jnp.allclose(outputs[s].astype(jnp.float32), expected, atol=1e-2), s @@ -443,7 +443,7 @@ def test_step_trains_and_has_vpd_signature(site_cs: tuple[SiteC, ...]): assert losses[-1]["p_imp"] < 2.0 # fp32 masters preserved through updates (SPEC N1). assert isinstance(state.components, DecompVU) - for V, U in state.components.vu.values(): + for _, (V, U) in state.components.sites_items(): assert V.dtype == jnp.float32 and U.dtype == jnp.float32 assert isinstance(state.ci_fn, ChunkwiseTransformerCIFn) assert state.ci_fn.chunks.in_proj_w.dtype == jnp.float32 @@ -479,7 +479,7 @@ def test_decomp_vu_shapes_fp32(): assert V_v.shape == (d, 12) and U_v.shape == (12, kvd) assert V_d.shape == (di, 8) and U_d.shape == (8, d) assert isinstance(vu, DecompVU) - assert all(a.dtype == jnp.float32 for pair in vu.vu.values() for a in pair) + assert all(a.dtype == jnp.float32 for pair in vu.stacks.values() for a in pair) def test_fresh_pgd_adversary_step(): diff --git a/param_decomp/tests/test_llama_simple_mlp.py b/param_decomp/tests/test_llama_simple_mlp.py index 567e524ea..634c852a8 100644 --- a/param_decomp/tests/test_llama_simple_mlp.py +++ b/param_decomp/tests/test_llama_simple_mlp.py @@ -434,7 +434,7 @@ def test_step_trains_and_has_vpd_signature(): assert losses[-1]["p_imp"] < 2.0 # fp32 masters preserved through updates (SPEC N1). assert isinstance(state.components, DecompVU) - for V, U in state.components.vu.values(): + for _, (V, U) in state.components.sites_items(): assert V.dtype == jnp.float32 and U.dtype == jnp.float32 assert isinstance(state.ci_fn, ChunkwiseTransformerCIFn) assert state.ci_fn.chunks.in_proj_w.dtype == jnp.float32 @@ -471,7 +471,7 @@ def test_decomp_vu_shapes_fp32(): assert V_v.shape == (d, 12) and U_v.shape == (12, kvd) assert V_fc.shape == (d, 8) and U_fc.shape == (8, di) assert V_dn.shape == (di, 16) and U_dn.shape == (16, d) - assert all(a.dtype == jnp.float32 for pair in vu.vu.values() for a in pair) + assert all(a.dtype == jnp.float32 for pair in vu.stacks.values() for a in pair) _REAL_CACHE_DIR = Path("/mnt/data/artifacts/mechanisms/param-decomp/pretrain_cache/spd-t-9d2b8f02") diff --git a/param_decomp/tests/test_no_bake_invariant.py b/param_decomp/tests/test_no_bake_invariant.py index 92f809081..8569a9d5b 100644 --- a/param_decomp/tests/test_no_bake_invariant.py +++ b/param_decomp/tests/test_no_bake_invariant.py @@ -11,7 +11,7 @@ from jax import random from param_decomp.ci_fn import Chunk, ChunkwiseTransformerCIArch, build_ci_fn -from param_decomp.components import DecompVU, SiteSpec +from param_decomp.components import SiteSpec, decomp_vu_from_sites from param_decomp.configs import ( FaithfulnessLossConfig, ImportanceMinimalityLossConfig, @@ -39,8 +39,8 @@ def _build_step_and_args(): leading_axes=("sequence",), ) assert lm.W.size == FROZEN_W_SIZE - components = DecompVU( - vu={ + components = decomp_vu_from_sites( + { SITE: ( random.normal(random.fold_in(key, 3), (D, C)) * 0.1, random.normal(random.fold_in(key, 4), (C, D)) * 0.1, diff --git a/param_decomp/tests/test_optim_torch_parity.py b/param_decomp/tests/test_optim_torch_parity.py index c0d6f4072..6fa0a5615 100644 --- a/param_decomp/tests/test_optim_torch_parity.py +++ b/param_decomp/tests/test_optim_torch_parity.py @@ -17,9 +17,9 @@ from param_decomp.configs import AdamWOptimizerConfig, AnyOptimizerConfig, MuonOptimizerConfig from param_decomp.run_state import ( _optimizer_with_clip, - chunk_stacked_muon_dimension_numbers, clip_by_global_norm_with_eps, optax_schedule, + stacked_muon_dimension_numbers, ) from param_decomp.schedule import ScheduleConfig @@ -145,7 +145,7 @@ def test_muon_orthogonalizes_2d_leaves_and_adam_falls_back_elsewhere(): def test_muon_chunk_stacked_dimension_numbers_orthogonalize_3d_and_adam_2d_bias_stacks(): - """SPEC S20 amendment (2026-07-11): under `chunk_stacked_muon_dimension_numbers` a 3D + """SPEC S20 amendment (2026-07-11): under `stacked_muon_dimension_numbers` a 3D `[n_chunks, d_in, d_out]` matrix stack is NS-orthogonalized per chunk slice, while a 2D `[n_chunks, d]` bias stack takes the Adam fallback — the reverse of optax's default 2D rule, which on the chunkwise CI-fn tree would orthogonalize the bias stacks.""" @@ -158,7 +158,7 @@ def test_muon_chunk_stacked_dimension_numbers_orthogonalize_3d_and_adam_2d_bias_ ) lr = 1e-3 opt = _optimizer_with_clip( - muon_cfg, lambda count: jnp.float32(lr), chunk_stacked_muon_dimension_numbers, mesh=None + muon_cfg, lambda count: jnp.float32(lr), stacked_muon_dimension_numbers, mesh=None ) key = jax.random.key(0) n_chunks = 3 @@ -205,7 +205,7 @@ def test_stacked_muon_update_matches_optax_muon(): name: jax.random.normal(jax.random.fold_in(key, i), p.shape) for i, (name, p) in enumerate(params.items()) } - dim_nums = chunk_stacked_muon_dimension_numbers + dim_nums = stacked_muon_dimension_numbers from typing import Literal @@ -250,7 +250,7 @@ def test_stacked_muon_sharded_matches_unsharded(): name: jax.random.normal(jax.random.fold_in(key, i), p.shape) for i, (name, p) in enumerate(params.items()) } - dim_nums = chunk_stacked_muon_dimension_numbers + dim_nums = stacked_muon_dimension_numbers def one_update(mesh: "Mesh | None"): opt = stacked_muon( diff --git a/param_decomp/tests/test_sharding.py b/param_decomp/tests/test_sharding.py index cf4346435..1764e3452 100644 --- a/param_decomp/tests/test_sharding.py +++ b/param_decomp/tests/test_sharding.py @@ -95,35 +95,25 @@ def test_jitted_sharded_inits_match_eager_values(): from param_decomp.components import vu_shape_groups assert max(len(g) for g in vu_shape_groups(sites).values()) >= 2 - # Placement is MODEL-OWNED and true ÷N ZeRO-1, split across the data + TP axes: V shards - # d_in over the data axes and C over `tp` (`P(("replicate","fsdp"), "tp")`); U shards C - # over `tp` and d_out over the data axes (`P("tp", ("replicate","fsdp"))`). C now carries - # the Megatron-C `tp` factor, so master + Adam still shard ÷(replicate·fsdp·tp) = ÷N. At - # tp=1 the tp axis is size 1 (C effectively unsharded). - full = ("replicate", "fsdp") + # Placement is MODEL-OWNED, owner-partitioned ÷N (hybrid HSDP rule): the STACK axis + # shards over `replicate` (whole matrices owned per node-group), matrix d dims over + # `fsdp`, C over `tp` — total ÷(replicate·fsdp·tp) = ÷N. On the sim mesh replicate is 1 + # (every stack length tiles it), so every group takes the stack rule. vu_placed = init_decomp_vu_placed(sites, jax.random.PRNGKey(1), mesh) vu_eager = init_decomp_vu(sites, jax.random.PRNGKey(1)) - for spec in sites: - V, U = vu_placed.site(spec.name) - assert isinstance(V.sharding, NamedSharding) and isinstance(U.sharding, NamedSharding) - assert V.sharding.spec == P(full, "tp"), (spec.name, V.sharding.spec) - assert U.sharding.spec == P("tp", full), (spec.name, U.sharding.spec) - # The placed init runs vmap-stacked per shape group (compile-time optimization) and - # must be BIT-identical to the jitted per-site `init_decomp_vu` it replaced (vmap over - # the same per-site keys) — the trajectory anchor. The unjitted eager reference differs - # from EITHER jitted path by ~1 ULP (XLA fusion), so it only gets allclose. + for Vs, Us in vu_placed.stacks.values(): + assert isinstance(Vs.sharding, NamedSharding) and isinstance(Us.sharding, NamedSharding) + assert Vs.sharding.spec == P("replicate", "fsdp", "tp"), Vs.sharding.spec + assert Us.sharding.spec == P("replicate", "tp", "fsdp"), Us.sharding.spec + # The placed init must be BIT-identical to the same init jitted WITHOUT placement + # (threefry is partitionable: `out_shardings` cannot perturb the stream) — the + # trajectory anchor. The unjitted eager reference differs from either jitted path by + # ~1 ULP (XLA fusion), so it only gets allclose. from functools import partial - import equinox as eqx - - per_site_shardings = eqx.filter_eval_shape( - partial(init_decomp_vu, sites), jax.random.PRNGKey(1) - ).shardings(mesh) - vu_jitted_per_site = jax.jit(partial(init_decomp_vu, sites), out_shardings=per_site_shardings)( - jax.random.PRNGKey(1) - ) + vu_jitted_unplaced = jax.jit(partial(init_decomp_vu, sites))(jax.random.PRNGKey(1)) for got, want in zip( - jax.tree.leaves(vu_placed), jax.tree.leaves(vu_jitted_per_site), strict=True + jax.tree.leaves(vu_placed), jax.tree.leaves(vu_jitted_unplaced), strict=True ): assert got.shape == want.shape and got.dtype == want.dtype assert jnp.array_equal(jnp.asarray(got), jnp.asarray(want)) diff --git a/param_decomp_lab/experiments/lm/run.py b/param_decomp_lab/experiments/lm/run.py index fb6ea9170..1d9a4fa7b 100644 --- a/param_decomp_lab/experiments/lm/run.py +++ b/param_decomp_lab/experiments/lm/run.py @@ -481,7 +481,7 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord": if perm_spec.want_uv_plots: components = { name: (np.asarray(V), np.asarray(U)) - for name, (V, U) in state.components.vu.items() + for name, (V, U) in state.components.sites_items() } slow_renderer.submit( partial( diff --git a/param_decomp_lab/experiments/resid_mlp/run.py b/param_decomp_lab/experiments/resid_mlp/run.py index 526296395..c4df7dc2f 100644 --- a/param_decomp_lab/experiments/resid_mlp/run.py +++ b/param_decomp_lab/experiments/resid_mlp/run.py @@ -158,7 +158,7 @@ def eval_fn(state: TrainState, now_step: int) -> dict[str, float]: ci_lower, ci_upper = single_feature_ci(lm, state.ci_fn) toy_uv_eval.log_uv_figure( uv_spec, - state.components.vu, + dict(state.components.sites_items()), ci_upper, now_step, wandb_active=built.run.wandb is not None, diff --git a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py index 3545eff3c..a237dcbc2 100644 --- a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py +++ b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py @@ -277,7 +277,7 @@ def test_step_trains_positionless_no_persistent_sources(): assert int(state.step) == 6 assert state.adversaries == {} # no persistent sources for the stochastic configs assert isinstance(state.components, DecompVU) - for V, U in state.components.vu.values(): + for _, (V, U) in state.components.sites_items(): assert V.dtype == jnp.float32 and U.dtype == jnp.float32 diff --git a/param_decomp_lab/experiments/tms/run.py b/param_decomp_lab/experiments/tms/run.py index 56e60a80e..835b9163d 100644 --- a/param_decomp_lab/experiments/tms/run.py +++ b/param_decomp_lab/experiments/tms/run.py @@ -141,7 +141,7 @@ def eval_fn(state: TrainState, now_step: int) -> dict[str, float]: ci_lower, ci_upper = single_feature_ci(lm, state.ci_fn) toy_uv_eval.log_uv_figure( uv_spec, - state.components.vu, + dict(state.components.sites_items()), ci_upper, now_step, wandb_active=built.run.wandb is not None, diff --git a/param_decomp_lab/experiments/tms/test_tms.py b/param_decomp_lab/experiments/tms/test_tms.py index 89932c322..a75b3e384 100644 --- a/param_decomp_lab/experiments/tms/test_tms.py +++ b/param_decomp_lab/experiments/tms/test_tms.py @@ -229,7 +229,7 @@ def test_step_trains_positionless_no_persistent_sources(): assert state.adversaries == {} # fp32 masters preserved assert isinstance(state.components, DecompVU) - for V, U in state.components.vu.values(): + for _, (V, U) in state.components.sites_items(): assert V.dtype == jnp.float32 and U.dtype == jnp.float32 diff --git a/param_decomp_lab/tests/test_toy_uv_eval.py b/param_decomp_lab/tests/test_toy_uv_eval.py index e08bf8d49..84758a4c2 100644 --- a/param_decomp_lab/tests/test_toy_uv_eval.py +++ b/param_decomp_lab/tests/test_toy_uv_eval.py @@ -71,7 +71,9 @@ def test_log_uv_figure_renders_png_when_configured(monkeypatch: pytest.MonkeyPat fake = _FakeWandb() monkeypatch.setitem(sys.modules, "wandb", fake) - toy_uv_eval.log_uv_figure(spec, vu.vu, probe_upper, now_step=42, wandb_active=True) + toy_uv_eval.log_uv_figure( + spec, dict(vu.sites_items()), probe_upper, now_step=42, wandb_active=True + ) assert len(fake.logged) == 1 payload, step = fake.logged[0] @@ -86,12 +88,16 @@ def test_log_uv_figure_noop_when_unconfigured_or_wandb_off(monkeypatch: pytest.M # config does not name UVPlots -> no-op no_uv = toy_uv_eval.toy_uv_spec(lm, _raw([])) - toy_uv_eval.log_uv_figure(no_uv, vu.vu, probe_upper, now_step=42, wandb_active=True) + toy_uv_eval.log_uv_figure( + no_uv, dict(vu.sites_items()), probe_upper, now_step=42, wandb_active=True + ) assert fake.logged == [] # configured but wandb off -> no-op with_uv = toy_uv_eval.toy_uv_spec( lm, _raw([{"type": "UVPlots", "identity_patterns": None, "dense_patterns": None}]) ) - toy_uv_eval.log_uv_figure(with_uv, vu.vu, probe_upper, now_step=42, wandb_active=False) + toy_uv_eval.log_uv_figure( + with_uv, dict(vu.sites_items()), probe_upper, now_step=42, wandb_active=False + ) assert fake.logged == []