From 162c64b9ce5e706226c523e777a1bde3c8f8bb2c Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Tue, 14 Jul 2026 04:17:28 +0000 Subject: [PATCH] refactor(logging): canonical wandb keys at the source, namespace-validated sink MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sink stacked three naming regimes: a short-key remap table (_METRIC_KEYS), startswith-based train/ prefixing, and verbatim pass-through — held together by a torch-overlay-parity constraint that retired with the torch trainer. Now every producer emits full canonical keys and MetricsSink.log asserts each key lives under the declared namespace tree (_WANDB_KEY_NAMESPACES), so a new metric joins the tree deliberately instead of drifting to the wandb top level. Emitted key strings are byte-identical to before — panels and overlays on existing runs are unaffected. (Only exception: a cfg.name override on the faith/imp terms now flows to the key like it always did for recon terms, instead of being ignored.) Also deletes dead lab code: infra/wandb.py's init_wandb (zero callers; still defined the retired 3-pool async slow_eval/step axis) and try_wandb (zero callers). Crew-Address: task/tidy-wandb-logging-structure Co-Authored-By: Claude Fable 5 --- param_decomp/run.py | 69 ++++++++-------- .../stacked_parity/test_stacked_parity.py | 18 +++-- param_decomp/tests/test_generic_model_io.py | 4 +- param_decomp/tests/test_llama8b.py | 19 +++-- param_decomp/tests/test_llama_simple_mlp.py | 2 +- param_decomp/tests/test_recon_log_keys.py | 79 ++++++++++--------- param_decomp/train.py | 36 ++++----- .../experiments/resid_mlp/test_resid_mlp.py | 2 +- param_decomp_lab/experiments/tms/test_tms.py | 2 +- param_decomp_lab/infra/wandb.py | 72 ----------------- 10 files changed, 124 insertions(+), 179 deletions(-) diff --git a/param_decomp/run.py b/param_decomp/run.py index 17ebca0eb..224b5fb95 100644 --- a/param_decomp/run.py +++ b/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/`, -# `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/`) 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/`) + "train/schedules/", # p_imp | gamma_imp, lr/{components,ci_fn,src[/]} + "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}_ + "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/, 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))) diff --git a/param_decomp/tests/stacked_parity/test_stacked_parity.py b/param_decomp/tests/stacked_parity/test_stacked_parity.py index 4d05ec4d6..53db5048e 100644 --- a/param_decomp/tests/stacked_parity/test_stacked_parity.py +++ b/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]: diff --git a/param_decomp/tests/test_generic_model_io.py b/param_decomp/tests/test_generic_model_io.py index 7640df7b3..962725f2b 100644 --- a/param_decomp/tests/test_generic_model_io.py +++ b/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" diff --git a/param_decomp/tests/test_llama8b.py b/param_decomp/tests/test_llama8b.py index 00fb594b2..c577831d4 100644 --- a/param_decomp/tests/test_llama8b.py +++ b/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,13 +548,18 @@ 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() @@ -562,9 +567,9 @@ def run_step(n_ascent_steps: int) -> tuple[TrainState, dict[str, jax.Array]]: 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(): diff --git a/param_decomp/tests/test_llama_simple_mlp.py b/param_decomp/tests/test_llama_simple_mlp.py index 567e524ea..69e56bed7 100644 --- a/param_decomp/tests/test_llama_simple_mlp.py +++ b/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(): diff --git a/param_decomp/tests/test_recon_log_keys.py b/param_decomp/tests/test_recon_log_keys.py index 5025b976f..64e15abd9 100644 --- a/param_decomp/tests/test_recon_log_keys.py +++ b/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/` - (`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/` 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/` 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/`; - 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/` 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}] diff --git a/param_decomp/train.py b/param_decomp/train.py index 6694d2157..f7184f5e8 100644 --- a/param_decomp/train.py +++ b/param_decomp/train.py @@ -76,24 +76,23 @@ class TrainState: def _grad_norm_metrics(components_grad: DecompVU, ci_fn_grad: Any) -> dict[str, Array]: - """Pre-clip gradient L2 norms, matching the torch `component_grad_norms` families: - per-leaf `grad_norms/components` / `grad_norms/ci_fns` (paths are this - pytree's own — e.g. `.vu['layers.18.mlp.gate_proj'][0]` for the per-site Llama - layout, vs torch's per-site names) and the overlay-critical - `grad_norms/summary/{components,ci_fns,total}`.""" + """Pre-clip gradient L2 norms: per-leaf `train/grad_norms/components` / + `train/grad_norms/ci_fns` (paths are this pytree's own — e.g. + `.vu['layers.18.mlp.gate_proj'][0]` for the per-site Llama layout) and the + overlay-critical `train/grad_norms/summary/{components,ci_fns,total}`.""" out: dict[str, Array] = {} def family(grad_tree: Any, prefix: str) -> Array: sum_sq = jnp.zeros((), jnp.float32) for path, leaf in jax.tree_util.tree_flatten_with_path(grad_tree)[0]: leaf_sum_sq = jnp.sum(leaf.astype(jnp.float32) ** 2) - out[f"grad_norms/{prefix}{jax.tree_util.keystr(path)}"] = jnp.sqrt(leaf_sum_sq) + out[f"train/grad_norms/{prefix}{jax.tree_util.keystr(path)}"] = jnp.sqrt(leaf_sum_sq) sum_sq = sum_sq + leaf_sum_sq - out[f"grad_norms/summary/{prefix}"] = jnp.sqrt(sum_sq) + out[f"train/grad_norms/summary/{prefix}"] = jnp.sqrt(sum_sq) return sum_sq total_sq = family(components_grad, "components") + family(ci_fn_grad, "ci_fns") - out["grad_norms/summary/total"] = jnp.sqrt(total_sq) + out["train/grad_norms/summary/total"] = jnp.sqrt(total_sq) return out @@ -132,11 +131,12 @@ def make_train_step( imp_min = imp_term.cfg imp_coeff = imp_term.coeff freq_coeff = imp_min.frequency.coeff if imp_min.frequency is not None else 0.0 - # Log the imp-min loss + its annealed param under penalty-kind-specific keys: the param - # is `p` for L_p / `gamma` for smooth-L0, and the loss carries the penalty's class name. + # Every term logs under `train/loss/` (`cfg.name` override or the config + # `type` literal). The imp-min annealed param key is penalty-kind-specific: `p` for + # L_p, `gamma` for smooth-L0. is_smooth_l0 = isinstance(imp_min, SmoothL0ImportanceMinimalityLossConfig) - imp_loss_key = "imp_smooth_l0" if is_smooth_l0 else "imp" - imp_min_param_key = "gamma_imp" if is_smooth_l0 else "p_imp" + imp_loss_key = f"train/loss/{imp_term.name}" + imp_min_param_key = f"train/schedules/{'gamma_imp' if is_smooth_l0 else 'p_imp'}" def batch_sharded(x: Array) -> Array: return batch_shard_leading(x, mesh) @@ -509,21 +509,21 @@ def pre_built_fwd( step=state.step + 1, ) metrics = { - "total": total_loss, - "faith": faith_loss, + "train/loss/total": total_loss, + f"train/loss/{faith_term.name}": faith_loss, imp_loss_key: imp_lp, - "freq": imp_freq, + "train/loss/FrequencyMinimalityLoss": imp_freq, imp_min_param_key: imp_min_param, - **{f"loss/{t.name}": v for t, v in zip(recon_terms, term_losses, strict=True)}, + **{f"train/loss/{t.name}": v for t, v in zip(recon_terms, term_losses, strict=True)}, **grad_norm_metrics, } source_lrs = { k: adv.source_lr(step_f32, total_steps) for k, adv in state.adversaries.items() } if len(source_lrs) == 1: - metrics["src_lr"] = next(iter(source_lrs.values())) + metrics["train/schedules/lr/src"] = next(iter(source_lrs.values())) else: - metrics |= {f"schedules/lr/src/{k}": v for k, v in source_lrs.items()} + metrics |= {f"train/schedules/lr/src/{k}": v for k, v in source_lrs.items()} return new_state, metrics return filter_jit(step, donate="all-except-first", compiler_options=compiler_options) diff --git a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py index 3545eff3c..0dbef1cb8 100644 --- a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py +++ b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py @@ -400,7 +400,7 @@ def test_end_to_end_pretrain_decompose_recovers_identity(): jax.random.fold_in(data_key, i), 2048, 5, 0.05, "at_least_zero_active" ) state, m = step(lm, state, x @ target.W_E, jax.random.fold_in(jax.random.PRNGKey(321), i)) - totals.append(float(m["total"])) + totals.append(float(m["train/loss/total"])) assert totals[-1] < totals[0], (totals[0], totals[-1]) ci_lower = single_feature_ci(lm, state.ci_fn, n_features=5) diff --git a/param_decomp_lab/experiments/tms/test_tms.py b/param_decomp_lab/experiments/tms/test_tms.py index 89932c322..0fe4c92b2 100644 --- a/param_decomp_lab/experiments/tms/test_tms.py +++ b/param_decomp_lab/experiments/tms/test_tms.py @@ -342,7 +342,7 @@ def test_end_to_end_pretrain_decompose_recovers_identity(): jax.random.fold_in(data_key, i), 2048, 5, 0.05, "at_least_zero_active" ) state, m = step(lm, state, x, jax.random.fold_in(jax.random.PRNGKey(321), i)) - totals.append(float(m["total"])) + totals.append(float(m["train/loss/total"])) assert totals[-1] < totals[0], (totals[0], totals[-1]) ci_lower = single_feature_ci(lm, state.ci_fn, n_features=5) diff --git a/param_decomp_lab/infra/wandb.py b/param_decomp_lab/infra/wandb.py index 51afa02fd..79671e745 100644 --- a/param_decomp_lab/infra/wandb.py +++ b/param_decomp_lab/infra/wandb.py @@ -1,17 +1,11 @@ import os import re -from collections.abc import Callable from pathlib import Path -from typing import Any import wandb -import wandb.errors from dotenv import load_dotenv from wandb.apis.public import File, Run -from param_decomp.log import logger -from param_decomp_lab.infra.settings import REPO_ROOT - # Regex patterns for parsing W&B run references. PD run IDs are formatted as # `p-<8 hex chars>` (see `RUN_TYPE_ABBREVIATIONS`). Legacy `s-…` IDs predate the # Run refactor; they still resolve when given as a full `entity/project/runs/id` path. @@ -121,69 +115,3 @@ def download_wandb_file(run: Run, wandb_run_dir: Path, file_name: str) -> Path: assert isinstance(file_on_wandb, File) file_on_wandb.download(exist_ok=True, replace=False, root=str(wandb_run_dir)) return wandb_run_dir / file_name - - -def init_wandb( - project: str, - run_id: str, - config_dict: dict[str, Any], - *, - resume: bool, - entity: str | None = None, - name: str | None = None, - tags: list[str] | None = None, - group: str | None = None, - view_meta: dict[str, Any] | None = None, -) -> None: - """Initialise W&B and log `config_dict` (already rendered to a flat-ish dict by the - caller — see `eval_metrics.wandb_config_dict`). - - `entity` falls back to `get_wandb_entity()`; `view_meta` is merged under a - `view_meta/` prefix so the UI can group runs by researcher-facing axes. `resume=True` - continues the existing wandb run `run_id` (continuous curves across a SLURM requeue); - `resume=False` creates a fresh run. - """ - wandb.init( - id=run_id, - project=project, - entity=entity or get_wandb_entity(), - name=name, - tags=tags, - group=group, - resume="allow" if resume else None, - ) - assert wandb.run is not None - wandb.run.log_code(root=str(REPO_ROOT / "param_decomp")) - - # Slow eval keys ride a dedicated `slow_eval/step` axis. The single-pool path - # logs them in-train (monotonic); the 3-pool path logs them retroactively from - # the async job (non-monotonic on the default axis). Defining the axis here lets - # both share the same panels. The async job redefines it too — idempotent. - wandb.define_metric("slow_eval/step") - wandb.define_metric("slow_eval/*", step_metric="slow_eval/step") - - wandb.config.update(config_dict) - - if view_meta: - wandb.config.update({f"view_meta/{k}": v for k, v in view_meta.items()}) - - -_n_try_wandb_comm_errors = 0 - - -# this exists to stop infra issues from crashing training runs -def try_wandb[**P, T](wandb_fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T | None: - """Call `wandb_fn`, warning and returning `None` on a wandb `CommError`. - - `CommError` is chosen to catch issues communicating with the wandb server but not - legitimate logging errors (e.g. not passing a dict to `wandb.log`, or the wrong - arguments to `wandb.save`). - """ - global _n_try_wandb_comm_errors - try: - return wandb_fn(*args, **kwargs) - except wandb.errors.CommError as e: - _n_try_wandb_comm_errors += 1 - logger.error( - f"wandb communication error, skipping log (total comm errors: {_n_try_wandb_comm_errors}): {e}" - )