From d5665ae5c30d8183d86683ae4df34b91ae5d64e1 Mon Sep 17 00:00:00 2001 From: oli Date: Fri, 3 Jul 2026 12:47:20 +0000 Subject: [PATCH 1/2] feat(eval): per-site faithfulness observability (rel/abs Frobenius) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The global FaithfulnessLoss (Σ_s‖Δ_s‖²/Σ_s numel) hides catastrophically broken small-norm sites (lore 2026-07-03--unmasked-kl-vs-faith-no-eval-bug-4-culprit-sites: 4 of 224 sites carried ~all of the 6.6 unmasked KL while faith sat at ~1e-8). Add per_site_faith_eval.make_per_site_faith_step: a jitted (model, components) -> {site: ‖Δ_s‖²_F} reusing model.weight_deltas — the SAME per-site deltas the global FaithfulnessLoss reduces, so Σ_s/Σ_s numel == FaithfulnessLoss exactly. ‖W_s‖²_F (the constant relative denominator) comes from the same step on a zeroed vu (Δ=W−0), cached once in the LM eval_fn closure. The LM in-loop eval logs, at eval cadence (always-on, cheap — one faith-loss-sized materialization per eval, which already runs every train step): - eval/faith/rel_frob_top{1..8} — 8 largest rel-Frobenius, RANK-indexed (stable line keys) - eval/faith/abs_frob_max - eval/faith/rel_frob_bar_chart — all-site, site-labeled (rank 0) Observability only — relative Frobenius is per-site WEIGHT-space error, NOT behavioral sensitivity (the culprit sites above were behaviorally inert on rel-Frobenius). The KL-under-delta-ablation probe is the diagnostic that ranks behavioral sensitivity; deferred pending a compute/design decision. Co-Authored-By: Claude Fable 5 --- param_decomp/README.md | 1 + param_decomp/per_site_faith_eval.py | 63 ++++++++++++++++ .../tests/test_per_site_faith_eval.py | 72 +++++++++++++++++++ param_decomp_lab/experiments/lm/run.py | 40 +++++++++++ 4 files changed, 176 insertions(+) create mode 100644 param_decomp/per_site_faith_eval.py create mode 100644 param_decomp/tests/test_per_site_faith_eval.py diff --git a/param_decomp/README.md b/param_decomp/README.md index 507f256fb..ba4c4012d 100644 --- a/param_decomp/README.md +++ b/param_decomp/README.md @@ -27,6 +27,7 @@ root with sibling packages `pretrain/` (the in-house target-LM pretrainer) and | `ci_fn.py` | shared-transformer CI fn over ordered site specs; the two leaky-hard squashings (SPEC §4.6, S5/S6) | | `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 | +| `per_site_faith_eval.py` | per-site faithfulness observability: a jitted step returning per-site `‖Δ_s‖²_F` (the same per-site deltas `FaithfulnessLoss` reduces away; `Σ_s / Σ_s numel == FaithfulnessLoss`), plus `per_site_faith_scalars` (rank-indexed rel-Frobenius top-k). The LM eval logs `eval/faith/{rel_frob_top*,abs_frob_max,rel_frob_bar_chart}`. Observability only — weight-space error, NOT behavioral sensitivity | | `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/*` | | `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 | diff --git a/param_decomp/per_site_faith_eval.py b/param_decomp/per_site_faith_eval.py new file mode 100644 index 000000000..c5f460857 --- /dev/null +++ b/param_decomp/per_site_faith_eval.py @@ -0,0 +1,63 @@ +"""Per-site faithfulness observability — the per-site Frobenius of the weight deltas. + +The global `FaithfulnessLoss` (`losses.faithfulness_loss`) is `Σ_s ‖Δ_s‖² / Σ_s numel` +over the per-site deltas `Δ_s = W_s − V_s@U_s` — an aggregate that HIDES catastrophically +broken small-norm sites: a handful of sites can carry nearly all of the unmasked-KL damage +while the mean sits at ~1e-8 (lore `2026-07-03--unmasked-kl-vs-faith-no-eval-bug-4-culprit-sites`). +This exposes the per-site breakdown so those sites are visible at eval cadence. + +Scope: relative Frobenius `‖Δ_s‖_F / ‖W_s‖_F` measures per-site WEIGHT-space error, which +is NOT the same as behavioral sensitivity — the worst rel-Frobenius sites need not be the +ones that dominate KL (in the run above they were behaviorally inert). This is pure +observability, not a culprit detector; the KL-under-delta-ablation probe is what ranks +behavioral sensitivity. +""" + +import jax.numpy as jnp +from jaxtyping import Array, Float + +from param_decomp.components import DecompVU +from param_decomp.jit_util import filter_jit +from param_decomp.lm import DecomposedModel + + +def make_per_site_faith_step( + lm: DecomposedModel, compiler_options: dict[str, bool | int | str] | None = None +): + """Build the `filter_jit`'d `per_site_faith(model, components) -> {site: ‖Δ_s‖²_F}` + (fp32 squared Frobenius per site). Reuses `model.weight_deltas` — the SAME per-site + `Δ_s = W_s − V_s@U_s` the global `FaithfulnessLoss` reduces — so + `Σ_s result[s] / Σ_s numel == FaithfulnessLoss` exactly. + + `model` (the frozen-weight-bearing `DecomposedModel`) is the jit ARG — array leaves + traced, never baked (the HLO-baking rule). Each `Δ_s` is materialized then reduced to a + scalar, so XLA frees it after its reduction: peak memory is ~one delta, not the whole + model. Passing a zeroed `components` yields `‖W_s‖²_F` (`Δ_s = W_s − 0`) — the constant + denominator for relative Frobenius. + """ + site_names = lm.site_names # static, read off the closed-over model (HLO-baking rule) + + def per_site_faith(model: DecomposedModel, components: DecompVU) -> dict[str, Float[Array, ""]]: + deltas = model.weight_deltas(components) + return {site: (deltas[site].astype(jnp.float32) ** 2).sum() for site in site_names} + + return filter_jit(per_site_faith, compiler_options=compiler_options) + + +PER_SITE_FAITH_TOP_K = 8 + + +def per_site_faith_scalars( + rel_frob: dict[str, float], abs_frob: dict[str, float], top_k: int +) -> dict[str, float]: + """Stable-keyed scalar summaries of the per-site Frobenius breakdown: the `top_k` + LARGEST relative-Frobenius values (rank-indexed `1..k`, so the keys stay fixed as the + ranking reshuffles over training), plus the largest absolute `‖Δ_s‖_F`. Which SITE holds + each rank rides the bar chart, not these keys.""" + ranked_rel = sorted(rel_frob.values(), reverse=True) + out = { + f"eval/faith/rel_frob_top{rank}": ranked_rel[rank - 1] + for rank in range(1, min(top_k, len(ranked_rel)) + 1) + } + out["eval/faith/abs_frob_max"] = max(abs_frob.values()) + return out diff --git a/param_decomp/tests/test_per_site_faith_eval.py b/param_decomp/tests/test_per_site_faith_eval.py new file mode 100644 index 000000000..36b232737 --- /dev/null +++ b/param_decomp/tests/test_per_site_faith_eval.py @@ -0,0 +1,72 @@ +"""CPU tests for per-site faithfulness observability. + +The load-bearing invariant: the per-site `‖Δ_s‖²_F` this step emits are the SAME per-site +deltas the global `FaithfulnessLoss` reduces, so `Σ_s result / Σ_s numel == FaithfulnessLoss` +exactly. Also checks the zeroed-vu `‖W_s‖²_F` denominator and the rank-indexed scalar helper. +""" + +import jax +import jax.numpy as jnp + +from param_decomp.components import init_decomp_vu +from param_decomp.losses import faithfulness_loss +from param_decomp.per_site_faith_eval import make_per_site_faith_step, per_site_faith_scalars +from param_decomp.targets.llama8b import llama_site_specs, mlp_family_site_cs +from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm + + +def _tiny_lm_and_vu(): + cfg = _tiny_cfg() + sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) + lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) + vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) + return lm, vu + + +def test_per_site_delta_sq_matches_weight_deltas(): + lm, vu = _tiny_lm_and_vu() + delta_sq = make_per_site_faith_step(lm)(lm, vu) + deltas = lm.weight_deltas(vu) + assert set(delta_sq) == set(deltas) + for site, delta in deltas.items(): + expected = (delta.astype(jnp.float32) ** 2).sum() + assert jnp.allclose(delta_sq[site], expected, rtol=1e-6) + + +def test_sum_matches_global_faithfulness_loss(): + lm, vu = _tiny_lm_and_vu() + delta_sq = make_per_site_faith_step(lm)(lm, vu) + deltas = lm.weight_deltas(vu) + total_numel = sum(d.size for d in deltas.values()) + reconstructed = sum(float(v) for v in delta_sq.values()) / total_numel + assert jnp.allclose(reconstructed, faithfulness_loss(deltas), rtol=1e-6) + + +def test_zeroed_vu_gives_frozen_weight_norms(): + lm, vu = _tiny_lm_and_vu() + step = make_per_site_faith_step(lm) + zero_vu = jax.tree.map(jnp.zeros_like, vu) + w_sq = step(lm, zero_vu) + for site, w in lm.weight_deltas(zero_vu).items(): # Δ = W − 0 = W + assert jnp.allclose(w_sq[site], (w.astype(jnp.float32) ** 2).sum(), rtol=1e-6) + assert float(w_sq[site]) > 0.0 + + +def test_per_site_faith_scalars_ranks_and_keys(): + rel_frob = {"a": 0.1, "b": 0.9, "c": 0.3, "d": 0.5} + abs_frob = {"a": 2.0, "b": 9.0, "c": 3.0, "d": 5.0} + out = per_site_faith_scalars(rel_frob, abs_frob, top_k=3) + assert out["eval/faith/rel_frob_top1"] == 0.9 + assert out["eval/faith/rel_frob_top2"] == 0.5 + assert out["eval/faith/rel_frob_top3"] == 0.3 + assert "eval/faith/rel_frob_top4" not in out + assert out["eval/faith/abs_frob_max"] == 9.0 + + +def test_per_site_faith_scalars_top_k_exceeds_site_count(): + rel_frob = {"a": 0.2, "b": 0.7} + abs_frob = {"a": 1.0, "b": 4.0} + out = per_site_faith_scalars(rel_frob, abs_frob, top_k=8) + assert out["eval/faith/rel_frob_top1"] == 0.7 + assert out["eval/faith/rel_frob_top2"] == 0.2 + assert "eval/faith/rel_frob_top3" not in out diff --git a/param_decomp_lab/experiments/lm/run.py b/param_decomp_lab/experiments/lm/run.py index 36bdc508a..d8aa427f9 100644 --- a/param_decomp_lab/experiments/lm/run.py +++ b/param_decomp_lab/experiments/lm/run.py @@ -70,6 +70,11 @@ ) from param_decomp.lm import DecomposedModel from param_decomp.log import setup_logger +from param_decomp.per_site_faith_eval import ( + PER_SITE_FAITH_TOP_K, + make_per_site_faith_step, + per_site_faith_scalars, +) from param_decomp.run import ( BackgroundRenderer, install_sigterm_flag, @@ -402,8 +407,14 @@ def _make_lm_eval_fn( position_ci_step = make_position_ci_step(lm, co) if want_position_ci else None arithmetic_eval = _make_arithmetic_eval(eval, built.target, lm, mesh, n_proc, is_main, co) + site_names = lm.site_names + per_site_faith_step = make_per_site_faith_step(lm, co) + # ‖W_s‖²_F is constant over training — computed once (lazily, from the first eval's own + # components pytree so shape/sharding are exact) via the same step on a zeroed vu. + frozen_weight_sq: dict[str, float] | None = None def eval_fn(state: TrainState, now_step: int) -> "LogRecord": + nonlocal frozen_weight_sq eval_pass_index = now_step // eval.every # uniform-average of per-batch scalars; mean-safe vs torch's accumulate-then- # compute() ONLY because every emitted key is a per-batch reduction that torch also @@ -428,6 +439,19 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord": eval_record: dict[str, LogValue] = { f"eval/{k}": float(v) / eval.n_steps for k, v in metric_sums.items() } + + # Per-site faithfulness observability: the per-site ‖Δ_s‖²_F the global + # FaithfulnessLoss reduces away. Collective (an all-reduce per site), so it runs in + # lockstep on every rank; the wandb bar chart is rank-0-only below. + delta_sq = {s: float(v) for s, v in per_site_faith_step(lm, state.components).items()} + if frozen_weight_sq is None: + zero_vu = jax.tree.map(jnp.zeros_like, state.components) + frozen_weight_sq = {s: float(v) for s, v in per_site_faith_step(lm, zero_vu).items()} + assert all(w > 0.0 for w in frozen_weight_sq.values()), "zero-norm frozen site" + abs_frob = {s: delta_sq[s] ** 0.5 for s in site_names} + rel_frob = {s: abs_frob[s] / frozen_weight_sq[s] ** 0.5 for s in site_names} + eval_record |= per_site_faith_scalars(rel_frob, abs_frob, PER_SITE_FAITH_TOP_K) + for class_name, attn_step in attn_steps.items(): # token-weighted (Σ sum_kl / Σ n), NOT the uniform per-batch average above — KL # is summed over distributions, divided by their count. @@ -515,11 +539,27 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord": "l0", title=f"L0_{eval.l0_ci_alive_threshold}", ) + # All-site per-site relative Frobenius, site-labeled — the rich per-site view the + # rank-indexed rel_frob_top* scalars can't carry (their keys are stable but drop + # site identity). Descending so the worst sites lead. + eval_record["eval/faith/rel_frob_bar_chart"] = wandb.plot.bar( + wandb.Table( + columns=["site", "rel_frob"], + data=[ + [s, v] + for s, v in sorted(rel_frob.items(), key=lambda kv: kv[1], reverse=True) + ], + ), + "site", + "rel_frob", + title="per-site ‖Δ‖_F / ‖W‖_F", + ) if is_main: headline = { k: eval_record[f"eval/{k}"] for k in ("ce_kl/kl_ci_masked", "ce_kl/ce_difference_ci_masked") } + headline["faith/rel_frob_top1"] = eval_record["eval/faith/rel_frob_top1"] print(f"[eval @ {now_step}] {headline}", flush=True) return eval_record From 57057595da37dd597960b6e2f8fc9559c00e74f7 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Tue, 28 Jul 2026 14:40:52 +0000 Subject: [PATCH 2/2] =?UTF-8?q?fix(test):=20update=20per-site-faith=20test?= =?UTF-8?q?=20to=20glu=5Fsite=5Fspecs=20(llama8b=E2=86=92glu=5Ftransformer?= =?UTF-8?q?=20split)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit llama_site_specs moved/renamed to glu_site_specs and mlp_family_site_cs moved to targets.glu_transformer since this branch's original base — the test's imports were stale after rebasing onto current feature/jax. Crew-Address: agent/6pmp --- param_decomp/tests/test_per_site_faith_eval.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/param_decomp/tests/test_per_site_faith_eval.py b/param_decomp/tests/test_per_site_faith_eval.py index 36b232737..ade1c35ef 100644 --- a/param_decomp/tests/test_per_site_faith_eval.py +++ b/param_decomp/tests/test_per_site_faith_eval.py @@ -11,13 +11,13 @@ from param_decomp.components import init_decomp_vu from param_decomp.losses import faithfulness_loss from param_decomp.per_site_faith_eval import make_per_site_faith_step, per_site_faith_scalars -from param_decomp.targets.llama8b import llama_site_specs, mlp_family_site_cs +from param_decomp.targets.glu_transformer import glu_site_specs, mlp_family_site_cs from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm def _tiny_lm_and_vu(): cfg = _tiny_cfg() - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) return lm, vu