refactor(components): owner-partitioned stacked persistence for V/U (flagship muon sharding)

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 /

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.<site>`) onto the current layout (site-keyed `components.vu`, `sources.<state_key>.<site>`) so a fine-tune can `restore_latest` it |
 | `sharding.py` | generic GSPMD helpers (`init_distributed`, `dp_mesh`, `replicate`, `shard_batch`) |

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). |
 
 ---