From 5a1a29a67abf7e3bfe4b8cc3460d7d627a62afd4 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sat, 11 Jul 2026 00:09:28 +0000 Subject: [PATCH 1/2] =?UTF-8?q?perf(llama8b):=20batch=20weight=5Fdeltas=20?= =?UTF-8?q?per=20kind=20=E2=80=94=20224=20site=20matmul=20groups=20?= =?UTF-8?q?=E2=86=92=207?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit llama8b's weight_deltas Python-looped over all 224 sites emitting a per-site fp32 W − (V@U).T matmul group into the loss graph (train step faith term + faith-warmup step). Per-kind dims are uniform across layers (already asserted by _per_kind_dims — the scan forward's precondition) and the frozen weights are already layer-stacked, so compute one batched fp32 matmul per KIND and return {kind: [n_sites, d_out, d_in]} stacks (rows layer-ascending; per_kind_delta_index maps site → (kind, row)). SPEC N2/S17 pin the math (fp32 deltas, global mean of squared entries), not the grouping: faithfulness_loss was already grouping-agnostic, its annotation and the DecomposedModel.weight_deltas contract now say so. Toy targets and llama_simple_mlp keep per-site returns (a valid grouping; their site counts are small). Parity (test_weight_deltas_match_per_site_reference, 3 site sets incl. heterogeneous C and a partial layer range): per-site delta values and V/U faithfulness grads BIT-IDENTICAL to the per-site loop; the loss scalar differs only by fp32 summation regrouping (224 → 7 partial sums, rtol 1e-6) and feeds no gradient, so trajectories are unchanged. Isolated-region compile A/B (faith-warmup step AOT at production 8B shapes, 224 sites, CPU backend): optimized-HLO dots 672 → 21, lines 41590 → 27506. Crew-Address: task/batch-weight-deltas-per-kind Co-Authored-By: Claude Fable 5 --- param_decomp/lm.py | 7 +- param_decomp/losses.py | 7 +- param_decomp/targets/llama8b.py | 38 ++++++++-- .../stacked_parity/test_stacked_parity.py | 7 +- param_decomp/tests/test_llama8b.py | 70 +++++++++++++++++-- 5 files changed, 111 insertions(+), 18 deletions(-) diff --git a/param_decomp/lm.py b/param_decomp/lm.py index 833dd0929..99443d9c6 100644 --- a/param_decomp/lm.py +++ b/param_decomp/lm.py @@ -186,8 +186,11 @@ def masked_site_outputs( the recon grid, which stays KL-on-final-logits.""" ... - def weight_deltas(self, vu: DecompVU) -> dict[str, Float[Array, "d_out d_in"]]: - """fp32 `W − V@U` per site from the fp32 master `vu` (SPEC N2).""" + def weight_deltas(self, vu: DecompVU) -> dict[str, Float[Array, "..."]]: + """fp32 `W − V@U` deltas from the fp32 master `vu` (SPEC N2), grouped as the target + chooses — per-site `(d_out, d_in)` entries, or per-kind stacks with a leading site + axis (llama8b). Consumed only by `faithfulness_loss` (S17), whose global mean is + grouping-invariant.""" ... diff --git a/param_decomp/losses.py b/param_decomp/losses.py index 3b261a8e8..5e73d41c8 100644 --- a/param_decomp/losses.py +++ b/param_decomp/losses.py @@ -67,9 +67,10 @@ def kl_per_position( @jaxtyped(typechecker=beartype) -def faithfulness_loss(weight_deltas: dict[str, Float[Array, "_ _"]]) -> Float[Array, ""]: - """`Σ_s ‖Δ_s‖² / Σ_s numel` over fp32 deltas (SPEC S17). Each `Δ_s` is `(d_out, d_in)`; - dims are per-site (anonymous, not bound across sites).""" +def faithfulness_loss(weight_deltas: dict[str, Float[Array, "..."]]) -> Float[Array, ""]: + """`Σ_s ‖Δ_s‖² / Σ_s numel` over fp32 deltas (SPEC S17). Values are per-site `(d_out, + d_in)` deltas or stacks of them with a leading site axis (`DecomposedModel.weight_deltas` + picks the grouping); the global mean is grouping-invariant.""" numerator = sum( ((delta.astype(jnp.float32) ** 2).sum() for delta in weight_deltas.values()), start=jnp.zeros((), jnp.float32), diff --git a/param_decomp/targets/llama8b.py b/param_decomp/targets/llama8b.py index 1849df0f3..68e88c72e 100644 --- a/param_decomp/targets/llama8b.py +++ b/param_decomp/targets/llama8b.py @@ -310,6 +310,18 @@ def _per_kind_dims(components: DecompVU) -> dict[str, tuple[int, int, int]]: return kind_dims +def per_kind_delta_index(sites: tuple[SiteSpec, ...]) -> dict[str, tuple[str, int]]: + """Site name -> `(kind, row)` into the per-kind stacks `weight_deltas` returns (rows are + layer-ascending within a kind — canonical site order restricted to the kind).""" + index: dict[str, tuple[str, int]] = {} + rows: dict[str, int] = {} + for spec in sites: + kind = parse_site_name(spec.name)[1] + index[spec.name] = (kind, rows.get(kind, 0)) + rows[kind] = rows.get(kind, 0) + 1 + return index + + def _stack_per_kind_vu(components: DecompVU, n_layers: int) -> dict[str, dict[str, Array]]: """Per decomposed KIND, the layer-stacked `(V, U)` arrays — the MASK-INDEPENDENT part of the scan inputs (a leading layer axis, one homogeneous body across layers). Mask/live/ @@ -923,15 +935,29 @@ def masked_component_activations( return collect_activations def weight_deltas(self, vu: DecompVU) -> dict[str, Array]: - """fp32 `W − V@U` per site from fp32 masters (SPEC N2; faithfulness input).""" - out: dict[str, Array] = {} + """fp32 `W − V@U` from fp32 masters (SPEC N2; faithfulness input), stacked per KIND: + `{kind: [n_sites, d_out, d_in]}`, rows layer-ascending (`per_kind_delta_index`). One + batched matmul per kind instead of a matmul group per site (224 sites → 7 groups in + the loss graph — compile time; S17's mean is grouping-invariant). Per-kind dims must + be uniform across layers (`_per_kind_dims`), as the scan masked forward already + requires.""" + _per_kind_dims(vu) + layers_by_kind: dict[str, list[int]] = {} for spec in self.sites: layer, kind = parse_site_name(spec.name) - W = _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind) - V, U = vu.site(spec.name) - out[spec.name] = ( - W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T + layers_by_kind.setdefault(kind, []).append(layer) + out: dict[str, Array] = {} + for kind, layers in layers_by_kind.items(): + all_layers = _frozen_site_weight(self.stacked, kind) + W = ( + all_layers + if layers == list(range(self.n_layer)) + else all_layers[jnp.asarray(layers)] ) + V = jnp.stack([vu.site(site_name(layer, kind))[0] for layer in layers]) + U = jnp.stack([vu.site(site_name(layer, kind))[1] for layer in layers]) + VU = V.astype(jnp.float32) @ U.astype(jnp.float32) + out[kind] = W.astype(jnp.float32) - VU.transpose(0, 2, 1) return out diff --git a/param_decomp/tests/stacked_parity/test_stacked_parity.py b/param_decomp/tests/stacked_parity/test_stacked_parity.py index 4d05ec4d6..4cbd2aed5 100644 --- a/param_decomp/tests/stacked_parity/test_stacked_parity.py +++ b/param_decomp/tests/stacked_parity/test_stacked_parity.py @@ -53,6 +53,7 @@ build_decomposed_lm, llama_site_specs, mlp_family_site_cs, + per_kind_delta_index, ) from param_decomp.tests.test_llama8b import _tiny_cfg from param_decomp.train import TrainState, make_train_step @@ -158,9 +159,9 @@ def test_site_inputs_and_weight_deltas_match(): site_inputs = lm.read_activations(resid, lm.site_names) for name in lm.site_names: _assert_close(site_inputs[name], f[f"out::site_input::{name}"], f"site_input {name}") - deltas = lm.weight_deltas(vu) - for name in lm.site_names: - _assert_close(deltas[name], f[f"out::wd::{name}"], f"weight_delta {name}") + deltas = lm.weight_deltas(vu) # per-kind stacks; fixtures are keyed per site + for name, (kind, row) in per_kind_delta_index(lm.sites).items(): + _assert_close(deltas[kind][row], f[f"out::wd::{name}"], f"weight_delta {name}") @_PENDING_REGEN diff --git a/param_decomp/tests/test_llama8b.py b/param_decomp/tests/test_llama8b.py index 00fb594b2..f62a2e288 100644 --- a/param_decomp/tests/test_llama8b.py +++ b/param_decomp/tests/test_llama8b.py @@ -35,17 +35,21 @@ UniformKSubsetRoutingConfig, ) from param_decomp.lm import DecomposedModel +from param_decomp.losses import faithfulness_loss from param_decomp.recon import build_loss_terms from param_decomp.schedule import ScheduleConfig from param_decomp.targets.llama8b import ( + KIND_ORDER, FrozenAttn, LlamaDecomposedModel, LlamaLayer, + _frozen_site_weight, build_decomposed_lm, canonical_site_cs, llama_site_specs, mlp_family_site_cs, parse_site_name, + per_kind_delta_index, site_name, ) from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step @@ -254,8 +258,9 @@ def test_clean_path_and_masked_identity(first: int, last: int): assert set(site_in) == set(names) deltas = lm.weight_deltas(vu) d, di = cfg.n_embd, cfg.n_intermediate - assert deltas[names[0]].shape == (di, d) # gate: (d_out, d_in) - assert deltas[names[2]].shape == (d, di) # down + n_sites = last - first + 1 + assert deltas["gate"].shape == (n_sites, di, d) # (n_sites, d_out, d_in) + assert deltas["down"].shape == (n_sites, d, di) assert all(v.dtype == jnp.float32 for v in deltas.values()) @@ -323,8 +328,8 @@ def test_attention_sites_clean_and_masked_identity(): deltas = lm.weight_deltas(vu) qd, kvd = cfg.n_head * cfg.head_dim, cfg.n_kv_head * cfg.head_dim - assert deltas[q_site].shape == (qd, cfg.n_embd) - assert deltas["layers.4.self_attn.v_proj"].shape == (kvd, cfg.n_embd) + assert deltas["q"].shape == (1, qd, cfg.n_embd) + assert deltas["v"].shape == (1, kvd, cfg.n_embd) def test_o_site_masks_attention_output(): @@ -353,6 +358,63 @@ def test_o_site_masks_attention_output(): assert site_in[o_site].shape == (b, t, cfg.n_head * cfg.head_dim) +@pytest.mark.parametrize( + "site_cs", + [ + canonical_site_cs( + tuple( + SiteC(site_name(layer, kind), 8) + for layer in range(_tiny_cfg().n_layer) + for kind in KIND_ORDER + ) + ), + mlp_family_site_cs(3, 6, 8), + _QVDOWN_SITE_CS, + ], + ids=["all_kinds_all_layers", "mlp_l3_6", "qv_down_l4"], +) +def test_weight_deltas_match_per_site_reference(site_cs: tuple[SiteC, ...]): + """The per-kind batched `weight_deltas` reproduces the per-site fp32 `W − (V@U).T` loop + bit-identically (SPEC N2 — same math, batched). The faithfulness loss over the stacks + matches the per-site loss to fp32 summation-regrouping tolerance (per-site partial sums + become per-kind reduces; S17's mean is grouping-invariant in exact math); its V/U grads + have no cross-site accumulation, so they must stay bit-identical.""" + cfg = _tiny_cfg() + sites = llama_site_specs(cfg, site_cs) + lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) + vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) + + def per_site_reference(vu_: DecompVU) -> dict[str, jax.Array]: + out: dict[str, jax.Array] = {} + for spec in lm.sites: + layer, kind = parse_site_name(spec.name) + W = _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], lm.stacked), kind) + V, U = vu_.site(spec.name) + out[spec.name] = ( + W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T + ) + return out + + deltas = lm.weight_deltas(vu) + reference = per_site_reference(vu) + index = per_kind_delta_index(lm.sites) + assert {kind for kind, _ in index.values()} == set(deltas) + assert all(v.dtype == jnp.float32 for v in deltas.values()) + for name, (kind, row) in index.items(): + assert jnp.array_equal(deltas[kind][row], reference[name]), name + + ref_loss, ref_grad = eqx.filter_value_and_grad( + lambda vu_: faithfulness_loss(per_site_reference(vu_)) + )(vu) + new_loss, new_grad = eqx.filter_value_and_grad( + lambda vu_: faithfulness_loss(lm.weight_deltas(vu_)) + )(vu) + assert jnp.allclose(new_loss, ref_loss, rtol=1e-6), (new_loss, ref_loss) + for name in lm.site_names: + for got, want, which in zip(new_grad.vu[name], ref_grad.vu[name], ("V", "U"), strict=True): + assert jnp.array_equal(got, want), (name, which) + + @pytest.mark.parametrize( "site_cs", [mlp_family_site_cs(4, 4, 8), mlp_family_site_cs(3, 6, 8), _QVDOWN_SITE_CS], From e1ac6de5b1f9c97f63a88be73798ea4933ba2ae0 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sat, 11 Jul 2026 00:53:01 +0000 Subject: [PATCH 2/2] perf(train): remat the faith loss so delta stacks aren't stored for backward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The faith reduce's backward needs the deltas (grad = 2Δ/N); with per-kind stacked weight_deltas XLA stored every [n_sites, d_out, d_in] fp32 stack from forward to backward — +3.3GiB/rank static temp measured at dp32 vs the per-site loop. jax.checkpoint at the two faith call sites (train step + faith warmup) recomputes them in the backward instead: one matmul per kind, numerics identical (SPEC N2/S17 unchanged). Crew-Address: task/batch-weight-deltas-per-kind Co-Authored-By: Claude Fable 5 --- param_decomp/targets/llama8b.py | 1 + param_decomp/train.py | 12 +++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/param_decomp/targets/llama8b.py b/param_decomp/targets/llama8b.py index 68e88c72e..50df1600d 100644 --- a/param_decomp/targets/llama8b.py +++ b/param_decomp/targets/llama8b.py @@ -946,6 +946,7 @@ def weight_deltas(self, vu: DecompVU) -> dict[str, Array]: for spec in self.sites: layer, kind = parse_site_name(spec.name) layers_by_kind.setdefault(kind, []).append(layer) + out: dict[str, Array] = {} for kind, layers in layers_by_kind.items(): all_layers = _frozen_site_weight(self.stacked, kind) diff --git a/param_decomp/train.py b/param_decomp/train.py index 6694d2157..df3fa3597 100644 --- a/param_decomp/train.py +++ b/param_decomp/train.py @@ -378,7 +378,14 @@ def loss_fn( # checkpointed block (mask never held, the memory win); others fall back to building # masks then `masked_output`. Either way the engine holds no per-forward mask stacks. ci_stacked = model.stack_ci(ci.lower) - faith_loss = faithfulness_loss(model.weight_deltas(components)) + # Remat: the loss reduce's backward needs the deltas (grad = 2Δ/N), so a target + # returning stacked deltas would otherwise keep every [n_sites, d_out, d_in] + # fp32 stack alive from forward to backward (≈ a full fp32 model of stored + # temps; +3.3GiB/rank static temp measured at dp32). Recomputing them in the + # backward is one matmul per stack. + faith_loss = jax.checkpoint( + lambda components_: faithfulness_loss(model.weight_deltas(components_)) + )(components) imp_lp, imp_freq = imp_min_terms(ci.upper, imp_min, imp_min_param) term_losses: list[Array] = [] @@ -542,6 +549,9 @@ def make_faith_warmup_step( def warmup_step( model: DecomposedModel, components: DecompVU, opt_state: optax.OptState ) -> tuple[DecompVU, optax.OptState, Array]: + # jax.checkpoint for the same reason as the train step's faith term: don't store + # the delta stacks for the backward, recompute them. + @jax.checkpoint def loss_fn(components_: DecompVU) -> Array: return faithfulness_loss(model.weight_deltas(components_))