refactor(logging): canonical wandb keys at the source, namespace-validated sink

param_decomp/run.py

@@ -129,23 +129,26 @@ def is_mesh_placed(a: object) -> bool:
     return eqx.combine(mesh_placed, fixed)
 
 
-# wandb keys match the torch trainer's (`train_step.py` emits `loss/<instance_key>`,
-# `optimize.py` prefixes `train/`) so a torch-vs-jax run pair overlays on one panel.
-# Recon-term keys arrive from the step already shaped (`loss/<instance_key>`) and are
-# train/-prefixed by the sink; this table maps only the step's fixed scalar keys.
-_METRIC_KEYS = {
-    "total": "train/loss/total",
-    "faith": "train/loss/FaithfulnessLoss",
-    "imp": "train/loss/ImportanceMinimalityLoss",
-    "imp_smooth_l0": "train/loss/SmoothL0ImportanceMinimalityLoss",
-    "freq": "train/loss/FrequencyMinimalityLoss",
-    "p_imp": "train/schedules/p_imp",
-    "gamma_imp": "train/schedules/gamma_imp",
-    "src_lr": "train/schedules/lr/src",
-    "step_time_s": "train/perf/step_time_s",
-    "elapsed_s": "train/perf/elapsed_s",
-    "eta_s": "train/perf/eta_s",
-}
+# The canonical wandb key tree. Every producer emits full keys (no sink-side remapping),
+# and `MetricsSink.log` asserts each key lives under one of these namespaces — a new
+# metric joins the tree deliberately instead of drifting to the wandb top level. The
+# `train/loss/` leaf names are the loss-config `type` literals (torch-era class names,
+# kept so runs overlay across the torch->jax transition). The figure tiers log outside
+# the sink (`render_and_log_slow_eval` -> `slow_eval/`, `render_and_log_arithmetic` ->
+# `eval/arithmetic/`) on rank-0 background threads.
+_WANDB_KEY_NAMESPACES = (
+    "train/loss/",  # total + one key per loss term (`train/loss/<term.name>`)
+    "train/schedules/",  # p_imp | gamma_imp, lr/{components,ci_fn,src[/<state_key>]}
+    "train/perf/",  # step_time_s, elapsed_s, eta_s
+    "train/mem/",  # peak_gb_per_rank
+    "train/grad_norms/",  # summary/{components,ci_fns,total}/{min,max,median} + per-leaf
+    "eval/ce_kl/",  # {kl,ce_difference}_<masking variant>
+    "eval/l0/",  # per-site + per-group L0 at the config threshold, bar_chart
+    "eval/loss/",  # PGDReconLoss (the sweep referee), attn-patterns terms
+    "eval/slow/",  # slow-tier scalars: loss/<hidden-acts terms>, identity_ci_error/*
+    "eval/arithmetic/",  # arithmetic-grid scalars (n_alive etc.)
+    "eval/identity_ci_error/",  # the toys' ground-truth CI eval
+)
 
 
 def _fmt_duration(seconds: float) -> str:
@@ -226,12 +229,8 @@ def for_run(
     def log(self, step: int, record: "LogRecord") -> None:
         if self._jsonl is None:
             return
-        record = {
-            _METRIC_KEYS.get(
-                k, f"train/{k}" if k.startswith(("grad_norms/", "loss/", "schedules/")) else k
-            ): v
-            for k, v in record.items()
-        }  # keys already starting "train/" or "eval/" pass through verbatim
+        off_tree = [k for k in record if not k.startswith(_WANDB_KEY_NAMESPACES)]
+        assert not off_tree, f"metric keys outside the canonical wandb tree: {off_tree}"
         # wandb-only viz objects (e.g. the CI_L0 bar chart) ride alongside the scalars to
         # wandb but are not jsonl/console serializable; split them off.
         scalars = {k: v for k, v in record.items() if isinstance(v, float)}
@@ -589,7 +588,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float:
         _aa2 = time.perf_counter()
         _as3, _am3 = step_fn(lm, _as2, _ab, _atk)
         _aa3 = time.perf_counter()
-        jax.block_until_ready((_as3, _am3["total"]))
+        jax.block_until_ready((_as3, _am3["train/loss/total"]))
         _aa4 = time.perf_counter()
         if is_main:
             print(
@@ -622,7 +621,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float:
         _dev.memory_stats()  # reset peak baseline read
         _ms_state, _ms_metrics = step_fn(lm, state, _mb, _mk)
         _ms_state, _ms_metrics = step_fn(lm, _ms_state, _mb, _mk)
-        jax.block_until_ready((_ms_state, _ms_metrics["total"]))
+        jax.block_until_ready((_ms_state, _ms_metrics["train/loss/total"]))
         _ms = _dev.memory_stats()
         if is_main and _ms is not None:
             print(
@@ -683,7 +682,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float:
             _ts1 = time.perf_counter()
             state, metrics = step_fn(lm, state, batch, random.fold_in(run_key, step))
             _ts2 = time.perf_counter()
-            jax.block_until_ready((state, metrics["total"]))
+            jax.block_until_ready((state, metrics["train/loss/total"]))
             _ts3 = time.perf_counter()
             print(
                 f"PD_TIME step {step}: sample={_ts1 - _ts0:.3f}s dispatch(py)={_ts2 - _ts1:.3f}s "
@@ -695,11 +694,11 @@ def _bench_dispatch(tree: object, n: int = 15) -> float:
             state, metrics = step_fn(lm, state, batch, random.fold_in(run_key, step))
 
         grad_norm_summary_window.append(
-            {k: v for k, v in metrics.items() if k.startswith("grad_norms/summary/")}
+            {k: v for k, v in metrics.items() if k.startswith("train/grad_norms/summary/")}
         )
 
         if _profiling:
-            jax.block_until_ready((state, metrics["total"]))
+            jax.block_until_ready((state, metrics["train/loss/total"]))
             if is_main:
                 print(
                     f"PD_PROFILE: step {step} wall {time.time() - _prof_t0:.3f}s (cumulative)",
@@ -721,22 +720,24 @@ def _bench_dispatch(tree: object, n: int = 15) -> float:
             or (dense is not None and now_step <= dense.until_step and now_step % dense.every == 0)
         )
         if log_now:
-            jax.block_until_ready(metrics["total"])
+            jax.block_until_ready(metrics["train/loss/total"])
             dt = time.time() - window_t0
             per_step = dt / max(now_step - last_logged, 1)
             last_logged = now_step
             record = {
-                k: float(v) for k, v in metrics.items() if not k.startswith("grad_norms/summary/")
+                k: float(v)
+                for k, v in metrics.items()
+                if not k.startswith("train/grad_norms/summary/")
             }
             record.update(_grad_norm_summary_window_stats(grad_norm_summary_window))
             grad_norm_summary_window.clear()
-            for loss_name in ("total", *(k for k in record if k.startswith("loss/"))):
+            for loss_name in (k for k in record if k.startswith("train/loss/")):
                 assert math.isfinite(record[loss_name]), (
                     f"non-finite loss {loss_name!r} at step {now_step}: {record[loss_name]}"
                 )
-            record["step_time_s"] = per_step
-            record["elapsed_s"] = time.time() - loop_t0
-            record["eta_s"] = (pd.steps - now_step) * per_step
+            record["train/perf/step_time_s"] = per_step
+            record["train/perf/elapsed_s"] = time.time() - loop_t0
+            record["train/perf/eta_s"] = (pd.steps - now_step) * per_step
             # the LR this step applied (optax count is the pre-increment `step` == now_step - 1)
             record["train/schedules/lr/components"] = float(jnp.asarray(sched_vu(now_step - 1)))
             record["train/schedules/lr/ci_fn"] = float(jnp.asarray(sched_ci(now_step - 1)))

param_decomp/tests/stacked_parity/test_stacked_parity.py

@@ -75,12 +75,20 @@
     "grad_norms/summary/components", "grad_norms/summary/ci_fns", "grad_norms/summary/total",
 )  # fmt: skip
 METRIC_KEY_BY_FIXTURE_KEY = {
-    "stoch": "loss/ChunkwiseSubsetReconLoss",
-    "ppgd": "loss/PersistentPGDReconLoss",
+    "total": "train/loss/total",
+    "faith": "train/loss/FaithfulnessLoss",
+    "imp": "train/loss/ImportanceMinimalityLoss",
+    "stoch": "train/loss/ChunkwiseSubsetReconLoss",
+    "ppgd": "train/loss/PersistentPGDReconLoss",
+    "p_imp": "train/schedules/p_imp",
+    "src_lr": "train/schedules/lr/src",
+    "grad_norms/summary/components": "train/grad_norms/summary/components",
+    "grad_norms/summary/ci_fns": "train/grad_norms/summary/ci_fns",
+    "grad_norms/summary/total": "train/grad_norms/summary/total",
 }
-"""The fixtures predate the recon-loss-terms unification; their metric keys are the
-old fixed names. The stored arrays themselves are untouched — only the lookup into
-the live metrics dict is remapped."""
+"""The fixtures predate the recon-loss-terms unification and the canonical-key sink;
+their metric keys are the old short names. The stored arrays themselves are untouched —
+only the lookup into the live metrics dict is remapped."""
 
 
 def _load() -> tuple[dict[str, np.ndarray], DecomposedModel, DecompVU, jnp.ndarray]:

param_decomp/tests/test_generic_model_io.py

@@ -275,8 +275,8 @@ def test_train_step_runs_through_generic_target():
     run_key = random.PRNGKey(3)
     for step_idx in range(2):
         state, metrics = step_fn(lm, state, inputs, random.fold_in(run_key, step_idx))
-        assert jnp.isfinite(metrics["total"]), (step_idx, metrics["total"])
-        assert "loss/StochasticReconLoss" in metrics
+        assert jnp.isfinite(metrics["train/loss/total"]), (step_idx, metrics["train/loss/total"])
+        assert "train/loss/StochasticReconLoss" in metrics
 
     assert not jnp.allclose(state.components.site(SITE)[0], V_before), (
         "V did not move — step is a no-op"

param_decomp/tests/test_llama8b.py

@@ -440,7 +440,7 @@ def test_step_trains_and_has_vpd_signature(site_cs: tuple[SiteC, ...]):
     for v in ppgd_adv.sources.values():
         assert float(v.min()) >= 0.0 and float(v.max()) <= 1.0
     # SPEC S9: p annealed below its 2.0 start by step 4 of 100.
-    assert losses[-1]["p_imp"] < 2.0
+    assert losses[-1]["train/schedules/p_imp"] < 2.0
     # fp32 masters preserved through updates (SPEC N1).
     assert isinstance(state.components, DecompVU)
     for V, U in state.components.vu.values():
@@ -548,23 +548,28 @@ def run_step(n_ascent_steps: int) -> tuple[TrainState, dict[str, jax.Array]]:
         return step(lm, make_state(), tokens, jax.random.PRNGKey(100))
 
     state, metrics = run_step(n_ascent_steps=1)
-    assert "loss/PGDReconLoss" in metrics
-    assert "loss/PersistentPGDReconLoss" not in metrics and "src_lr" not in metrics
+    assert "train/loss/PGDReconLoss" in metrics
+    assert "train/loss/PersistentPGDReconLoss" not in metrics
+    assert "train/schedules/lr/src" not in metrics
     assert jnp.isfinite(
         jnp.array(
             [
                 float(metrics[k])
-                for k in ("total", "loss/PGDReconLoss", "loss/ChunkwiseSubsetReconLoss")
+                for k in (
+                    "train/loss/total",
+                    "train/loss/PGDReconLoss",
+                    "train/loss/ChunkwiseSubsetReconLoss",
+                )
             ]
         )
     ).all()
     assert state.adversaries == {}, "fresh adversary carries no persistent sources"
     assert int(state.step) == 1
 
     _, metrics_unascended = run_step(n_ascent_steps=0)
-    assert float(metrics["loss/PGDReconLoss"]) >= float(metrics_unascended["loss/PGDReconLoss"]), (
-        "one sign step from the same init must not weaken the adversary"
-    )
+    assert float(metrics["train/loss/PGDReconLoss"]) >= float(
+        metrics_unascended["train/loss/PGDReconLoss"]
+    ), "one sign step from the same init must not weaken the adversary"
 
 
 def test_chunkwise_ci_init_vmap_matches_unrolled_reference():

param_decomp/tests/test_llama_simple_mlp.py

@@ -431,7 +431,7 @@ def test_step_trains_and_has_vpd_signature():
     for v in ppgd_adv.sources.values():
         assert float(v.min()) >= 0.0 and float(v.max()) <= 1.0
     # SPEC S9: p annealed below its 2.0 start by step 4 of 100.
-    assert losses[-1]["p_imp"] < 2.0
+    assert losses[-1]["train/schedules/p_imp"] < 2.0
     # fp32 masters preserved through updates (SPEC N1).
     assert isinstance(state.components, DecompVU)
     for V, U in state.components.vu.values():

param_decomp/tests/test_recon_log_keys.py

@@ -1,21 +1,17 @@
-"""E11 logged-key parity: every `train/loss/*` key a JAX run emits must equal the
-torch key for the same config, so a torch-vs-jax run pair overlays on one wandb panel.
-
-Two emit paths feed `train/loss/*`:
-
-  * Fixed scalar terms (`faith`, `imp`) — mapped through `run._METRIC_KEYS`.
-  * Recon terms — arrive from the jitted step already shaped `loss/<term.name>`
-    (`train.py`), then `train/`-prefixed by the sink.
-
-`term.name` is set by `recon.build_loss_terms` to `cfg.name or cfg.type`. Torch's
-`Metric.instance_key` is `cfg.name or type(self).__name__`. They agree byte-for-byte
-because torch's `LOSS_METRIC_CLASSES` is keyed by `cls.__name__` and dispatch does
-`LOSS_METRIC_CLASSES[cfg.type](cfg)` — so torch can only run a config whose `type`
-literal equals the class name. This test pins the JAX half (`term.name == cfg.type`);
-the torch half (`cls.__name__ == cfg.type`) was pinned against live torch dispatch
-before push-1 severed the torch import — it now holds by the frozen `type` literals.
+"""Logged-key pins: every loss term logs under `train/loss/<term.name>` where
+`term.name` is `cfg.name or cfg.type`, and `MetricsSink.log` rejects keys outside the
+canonical namespace tree (`run._WANDB_KEY_NAMESPACES`).
+
+The `type` literals are the torch-era class names — frozen when the JAX trainer took
+over so run pairs overlay across the transition, kept now for continuity of the
+existing wandb corpus.
 """
 
+import json
+from pathlib import Path
+
+import pytest
+
 from param_decomp.configs import (
     AdamPGDConfig,
     CIMaskedReconLayerwiseLossConfig,
@@ -34,6 +30,7 @@
     UnmaskedReconLossConfig,
 )
 from param_decomp.recon import ReconLossTerm, build_loss_terms
+from param_decomp.run import MetricsSink
 from param_decomp.schedule import ScheduleConfig
 
 SITE_NAMES = ("h.0.mlp.c_fc", "h.0.mlp.down_proj")
@@ -82,36 +79,42 @@ def _recon_terms(recon_configs: tuple[object, ...]) -> tuple[ReconLossTerm, ...]
 
 def test_recon_term_name_is_config_type_by_default():
     """Each recon term's log-key suffix (`term.name`) is the config `type` literal when
-    no `name` override is set — the JAX half of torch's `instance_key`."""
+    no `name` override is set."""
     emitted = {term.name for term in _recon_terms(RECON_CONFIGS)}
     expected = {cfg.type for cfg in RECON_CONFIGS}
     assert emitted == expected, emitted ^ expected
 
 
 def test_recon_term_name_honors_name_override():
-    """A `name` override on the config flows to `term.name`, exactly as torch's
-    `instance_key` returns `cfg.name` when set."""
+    """A `name` override on the config flows to `term.name` (and so to the
+    `train/loss/<name>` key)."""
     cfg = StochasticReconLossConfig(coeff=1.0, name="StochasticReconLoss_probe")
     (term,) = _recon_terms((cfg,))
     assert term.name == "StochasticReconLoss_probe"
 
 
-def test_fixed_scalar_loss_keys_use_torch_class_names():
-    """`run._METRIC_KEYS` maps the two non-recon scalar terms to `train/loss/<ClassName>`;
-    pin those suffixes to the config `type` literals (== torch class names)."""
-    from param_decomp.run import _METRIC_KEYS
-
-    faith, imp = _non_recon_configs()
-    assert _METRIC_KEYS["faith"] == f"train/loss/{faith.type}"
-    assert _METRIC_KEYS["imp"] == f"train/loss/{imp.type}"
-
-
-def test_no_train_loss_no_beta_key():
-    """Torch's `ImportanceMinimalityLoss_no_beta` diagnostic is emitted only by
-    `Metric.compute` (the eval path), never from the train step. The JAX trainer must
-    not emit a `train/loss/*_no_beta` key — doing so was a logged-key divergence from
-    torch (#647)."""
-    from param_decomp.run import _METRIC_KEYS
-
-    assert not any(k.endswith("_no_beta") for k in _METRIC_KEYS), _METRIC_KEYS
-    assert not any(v.endswith("_no_beta") for v in _METRIC_KEYS.values()), _METRIC_KEYS
+def test_fixed_term_names_are_the_frozen_class_literals():
+    """The faith/imp singleton terms carry the frozen class-name literals — the
+    `train/loss/<term.name>` keys existing panels select on."""
+    faith_cfg, imp_cfg = _non_recon_configs()
+    terms = build_loss_terms(
+        (faith_cfg, imp_cfg, StochasticReconLossConfig(coeff=1.0)),
+        site_names=SITE_NAMES,
+    )
+    assert terms.faith.name == "FaithfulnessLoss"
+    assert terms.imp.name == "ImportanceMinimalityLoss"
+
+
+def test_sink_rejects_keys_outside_the_canonical_tree(tmp_path: Path):
+    """`MetricsSink.log` asserts every key lives under a declared namespace — a bare or
+    typo'd key fails fast instead of drifting to the wandb top level."""
+    jsonl_path = tmp_path / "metrics.jsonl"
+    with jsonl_path.open("a") as jsonl:
+        sink = MetricsSink(jsonl=jsonl, wandb_module=None)
+        sink.log(1, {"train/loss/total": 1.0, "eval/ce_kl/kl_ci_masked": 0.5})
+        with pytest.raises(AssertionError, match="canonical wandb tree"):
+            sink.log(2, {"step_time_s": 0.1})
+        with pytest.raises(AssertionError, match="canonical wandb tree"):
+            sink.log(3, {"train/losss/total": 1.0})
+    logged = [json.loads(line) for line in jsonl_path.read_text().splitlines()]
+    assert logged == [{"step": 1, "train/loss/total": 1.0, "eval/ce_kl/kl_ci_masked": 0.5}]