perf(llama8b): batch weight_deltas per kind (224 site matmul groups → 7)

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

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

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,30 @@ 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
 
 

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

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],

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