From bc55bfed1b8a5d1769127443a9e753e9ad9c8c9d Mon Sep 17 00:00:00 2001 From: Antovigo Date: Fri, 14 Aug 2026 22:07:19 +0000 Subject: [PATCH 01/15] feat(metrics): one namespace rule, and core evals on both tPD streams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rule: the data a run OPTIMIZES FOR is unlabelled. Keys become `/[/]/`, the stream segment sitting immediately after `train/`/`eval/` — the only position that works uniformly, since `eval` has several families (`ce_kl/`, `l0/`, `loss/`). A plain run has ONE stream and never emits the segment, so no plain-run or toy key moves; that is enforced by `stream_log_prefix` keying off `context.target_batches is None` and pinned by `test_stream_log_namespace.py`. Naming fixes, all reachable only on a targeted run: - The stream segment sat in different places per side (`train/loss/nontarget/X` vs `eval/nontarget/...`); it now leads both. - One quantity had three spellings: the non-target imp-min loss was keyed by the step's short record key (`imp_smooth_l0`), the target stream spelled it by class name, and the coefficient was `impmin`. `IMP_MIN_METRIC_NAMES` is now shared by both streams so they cannot drift. - Attention-pattern recon, hidden-acts recon, `IdentityCIError` and `WellTemperedness` read the broad corpus but hardcoded `eval/`, which on a targeted run reads as target data. They stay single-stream; only the label is corrected. Core-side metrics shared with the toys take a `log_prefix_for_context` callback, since core cannot import the LM helper. Deliberately NOT fixed, because both are reachable only through plain-run keys: the stray `slow/` segment (present on two of three slow-tier evals), and stream-independent scalars keeping the bare namespace. Evals: `make_lm_evaluation` now emits one operation PER STREAM per metric, so a metric is authored once and a tPD run gets both readouts. The stream set comes from `target_pool_batches_for` — `None` on the plain root collapses it to the single broad stream. `CI_L0` / `PGDReconLoss` / `CIMaskedReconLoss` bind to both; `UnmaskedNoDeltaReconLoss` to the optimized stream alone (it is the non-target pass's own training term, already reported there as a train loss, so an eval of it off-target would restate the objective). `CIMaskedReconLoss` and `UnmaskedNoDeltaReconLoss` become authorable as evals — the `PGDReconLoss` dual-role pattern — each selecting ONE arm of the CE/KL evaluator. That is how a tPD run gets those two numbers without `CEandKLLosses`'s 11-scalar record, at 2 forwards per batch instead of 7. Masks for every arm are still drawn, so a narrowed evaluator's numbers are bit-identical to the full one's. They keep the `ce_kl/kl_` spelling a plain run logs them under; authoring both a narrow metric and `CEandKLLosses` is refused at bind time rather than colliding hours into a run. Ports: `WeightMagnitude` (now sorted by DESCENDING magnitude — x is a rank, not a component id, so the spectrum's knee is readable) and `TwoStreamCIMeanPerComponent` (figure key pluralised to match its config and plot function). Breaking for targeted runs only: every `nontarget/` key is renamed and broad-stream evals move under `eval/nontarget_data/`, so tPD dashboards need updating. No config field is removed, so pinned launch configs keep parsing. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/CLAUDE.md | 33 +++ param_decomp/core/configs.py | 36 +++ param_decomp/core/run.py | 11 +- param_decomp/core/slow_eval.py | 87 +++++++ .../core/tests/test_eval_averaging_parity.py | 2 + .../core/tests/test_well_temperedness_eval.py | 1 + param_decomp/core/train.py | 45 +++- param_decomp/core/well_temperedness_eval.py | 24 +- param_decomp/experiments/CLAUDE.md | 26 ++ param_decomp/experiments/eval_config.py | 8 + .../lm/diagnostic_eval_operations.py | 96 ++++++- param_decomp/experiments/lm/eval.py | 48 +++- param_decomp/experiments/lm/eval_config.py | 12 + param_decomp/experiments/lm/eval_context.py | 5 + .../experiments/lm/eval_operations.py | 237 +++++++++++++----- .../experiments/lm/scalar_eval_operations.py | 119 ++++++++- .../experiments/lm/test_eval_operations.py | 4 + .../lm/test_stream_log_namespace.py | 124 +++++++++ param_decomp/experiments/lm/training.py | 10 +- .../experiments/lm/training_targeted.py | 33 ++- param_decomp/experiments/test_toy_eval.py | 2 + .../experiments/tms/test_targeted_tms.py | 8 +- param_decomp/experiments/toy_eval.py | 3 + param_decomp/tests/test_eval_tier.py | 4 + 24 files changed, 850 insertions(+), 128 deletions(-) create mode 100644 param_decomp/experiments/lm/test_stream_log_namespace.py diff --git a/param_decomp/core/CLAUDE.md b/param_decomp/core/CLAUDE.md index 0df95f4bb..15e5573c6 100644 --- a/param_decomp/core/CLAUDE.md +++ b/param_decomp/core/CLAUDE.md @@ -399,6 +399,39 @@ on one. priority-fusion). Don't chase graph-shrink refactors for compile time without new evidence. +## Metric namespaces: the data a run optimizes for is unlabelled + +Every logged key is `/[/]/` — `tier` is `train`/`eval`, and the +STREAM segment sits immediately after it (the only position that works uniformly, since +`eval` has several families: `ce_kl/`, `l0/`, `loss/`). The one rule: + +- The data the run OPTIMIZES FOR carries no stream segment. +- Anything else carries `nontarget_data/` (`configs.NONTARGET_STREAM`). + +So `train/loss/total` and `eval/l0/...` mean "the data of interest" in BOTH run kinds — a +plain run's corpus, a tPD run's prompt pool — while a targeted run's broad corpus reads +`train/nontarget_data/loss/total` / `eval/nontarget_data/l0/...`. The consequence is +deliberate and is the trap: the bare namespace is comparable across run kinds AS AN +OBJECTIVE, but it is not the same DATA, so a corpus-vs-corpus comparison must read +`eval/nontarget_data/` on the tPD side. + +**A plain run has ONE stream and therefore never emits the segment** — every plain-run and +toy key is exactly what it was before targeted runs existed. That is a COMPATIBILITY +GUARANTEE, not an accident: `scalar_eval_operations.stream_log_prefix` keys off +`context.target_batches is None`, and `experiments/lm/test_stream_log_namespace.py` pins it. +Two consequences when adding a metric: + +- An eval that reads `context.batches` measures the BROAD stream and must take its prefix + from `stream_log_prefix("broad", context)` — never a hardcoded `"eval/"`, which would + claim to be target data on a tPD run. Core-side metrics shared with the toys take a + `log_prefix_for_context` callback instead (`well_temperedness_eval`). +- An eval that reads no batch at all (`WeightMagnitude`, the U/V norm ratios) has no stream + and keeps the bare namespace on both run kinds. + +One quantity gets ONE name across streams: `IMP_MIN_METRIC_NAMES` is shared by the target +stream (expanded in `run._METRIC_KEYS`) and the non-target stream (keyed in +`train.make_targeted_train_step`) so the two cannot drift apart. + ## Gotchas - **Process bring-up is config-derived, NEVER SLURM-sniffing** (`sharding.py`): the LM diff --git a/param_decomp/core/configs.py b/param_decomp/core/configs.py index c82653ca2..b230312d9 100644 --- a/param_decomp/core/configs.py +++ b/param_decomp/core/configs.py @@ -233,6 +233,7 @@ class SmoothL0ImportanceMinimalityLossConfig(LossMetricConfig): class CIMaskedReconLossConfig(LossMetricConfig, HiddenActsReconstructionMixin): + slow: ClassVar[bool] = False type: Literal["CIMaskedReconLoss"] = "CIMaskedReconLoss" @@ -467,6 +468,18 @@ class ComponentActivationDensityConfig(BaseConfig): ci_alive_threshold: float = 0.0 +class WeightMagnitudeConfig(BaseConfig): + """Per-site `‖V_c‖·‖U_c‖` scatter in descending magnitude order, log y. + + Reads the trained V/U alone — no forward pass and no eval batch, so it costs two norm + reductions and a plot. The norms reduce on device; only `C` floats per site are pulled. + Stream-independent by construction: it never touches a batch, so it carries no stream + segment on either run kind.""" + + slow: ClassVar[bool] = True + type: Literal["WeightMagnitude"] = "WeightMagnitude" + + class IdentityCITargetSpec(BaseConfig): """A layer expected to produce an Identity CI pattern over `n_features` features.""" @@ -657,6 +670,7 @@ class UnmaskedNoDeltaReconLossConfig(LossMetricConfig): exists here, and it is non-target-only — the plain and target-pass unions have no member for it.""" + slow: ClassVar[bool] = False type: Literal["UnmaskedNoDeltaReconLoss"] = "UnmaskedNoDeltaReconLoss" @@ -974,6 +988,28 @@ class ResumeProvenance(BaseConfig): """The parent's orbax `ckpts//` checkpoint step to initialize V/U + ci_fn from.""" +NONTARGET_STREAM = "nontarget_data" +"""The log-namespace segment for a tPD run's broad (non-target) stream. + +ONE rule across `train/` and `eval/`: the data a run OPTIMIZES FOR is unlabelled, and +anything else carries its stream, as the segment immediately after the tier — so +`train/loss/X` / `eval/X` mean "the data of interest" in both run kinds, and the corpus +stream of a targeted run reads `train/nontarget_data/loss/X` / `eval/nontarget_data/X`. A +plain run has ONE stream and therefore never emits this segment: every plain-run key is +what it was before targeted runs existed.""" + + +IMP_MIN_METRIC_NAMES: dict[str, str] = { + "imp": "ImportanceMinimalityLoss", + "imp_smooth_l0": "SmoothL0ImportanceMinimalityLoss", + "freq": "FrequencyMinimalityLoss", +} +"""Step-record short key -> logged loss name, for the terms whose record key is not already +the term's class name. Shared so the target stream (expanded by `run._METRIC_KEYS`) and the +non-target stream (keyed in `train.make_targeted_train_step`) cannot drift apart: both +streams must spell the same quantity the same way.""" + + # --------------------------------------------------------------------------- # wandb.config shaping # --------------------------------------------------------------------------- diff --git a/param_decomp/core/run.py b/param_decomp/core/run.py index 6fb082f31..3055c428c 100644 --- a/param_decomp/core/run.py +++ b/param_decomp/core/run.py @@ -50,6 +50,8 @@ from param_decomp.core.ci_fn import CIFnArch from param_decomp.core.components import init_component_stacks from param_decomp.core.configs import ( + IMP_MIN_METRIC_NAMES, + NONTARGET_STREAM, AnyPDConfig, Cadence, NontargetConfig, @@ -221,9 +223,7 @@ def is_mesh_placed(a: object) -> bool: _METRIC_KEYS = { "total": "train/loss/total", "faith": "train/loss/FaithfulnessLoss", - "imp": "train/loss/ImportanceMinimalityLoss", - "imp_smooth_l0": "train/loss/SmoothL0ImportanceMinimalityLoss", - "freq": "train/loss/FrequencyMinimalityLoss", + **{short: f"train/loss/{name}" for short, name in IMP_MIN_METRIC_NAMES.items()}, "p_imp": "train/schedules/p_imp", "gamma_imp": "train/schedules/gamma_imp", "src_lr": "train/schedules/lr/src", @@ -333,7 +333,10 @@ def log(self, step: int, record: "LogRecord") -> None: self._last_committed_step = step record = { _METRIC_KEYS.get( - k, f"train/{k}" if k.startswith(("grad_norms/", "loss/", "schedules/")) else k + k, + f"train/{k}" + if k.startswith(("grad_norms/", "loss/", "schedules/", f"{NONTARGET_STREAM}/")) + else k, ): v for k, v in record.items() } # keys already starting "train/" or "eval/" pass through verbatim diff --git a/param_decomp/core/slow_eval.py b/param_decomp/core/slow_eval.py index 04794f1db..e4e6fb522 100644 --- a/param_decomp/core/slow_eval.py +++ b/param_decomp/core/slow_eval.py @@ -68,6 +68,7 @@ lower_leaky_hard_sigmoid, upper_leaky_hard_sigmoid, ) +from param_decomp.core.components import ComponentStacks from param_decomp.core.configs import ( DenseCITargetSpec, IdentityCIErrorConfig, @@ -578,6 +579,92 @@ def _plot_ci_matrices(matrices: dict[str, np.ndarray], colormap: str, title_pref return _render_figure(fig) +def _component_weight_magnitudes(components: ComponentStacks) -> dict[str, Array]: + return { + name: jnp.linalg.norm(sc.V.astype(jnp.float32), axis=0) + * jnp.linalg.norm(sc.U.astype(jnp.float32), axis=1) + for name, sc in components.sites_items() + } + + +def weight_magnitudes(components: ComponentStacks) -> dict[str, np.ndarray]: + """Per-site `‖V_c‖·‖U_c‖` as host `(C,)` vectors. The norms reduce ON DEVICE, so only + C floats per site cross the boundary — never the V/U matrices themselves.""" + return { + name: np.asarray(value) for name, value in _component_weight_magnitudes(components).items() + } + + +def mean_cis(reductions: dict[str, SiteReduction]) -> dict[str, np.ndarray]: + """Per-site token-weighted mean CI, with the zero-position guard in ONE place.""" + assert all(r.n_positions > 0 for r in reductions.values()) + return {site: r.ci_sums / r.n_positions for site, r in reductions.items()} + + +def plot_weight_magnitudes(magnitudes: dict[str, np.ndarray]) -> bytes: + """Per-site `‖V_c‖·‖U_c‖` in DESCENDING magnitude order, log y. x is a component's rank + within its site, NOT its component id: the readable quantity is the spectrum's shape — + how fast magnitude falls off, and where it knees — which component order destroys.""" + n_rows, n_cols = _grid_dims(len(magnitudes)) + fig = Figure(figsize=(8 * n_cols, 3 * n_rows)) + axs = fig.subplots(n_rows, n_cols, squeeze=False) + flat_axes = axs.T.ravel() + for ax in flat_axes[len(magnitudes) :]: + ax.set_visible(False) + for ax, (name, values) in zip(flat_axes, magnitudes.items(), strict=False): + ax.scatter(range(len(values)), np.sort(values)[::-1], marker="x", s=10) + ax.set_yscale("log") + ax.set_xlabel("Component (descending ‖V‖·‖U‖)") + ax.set_ylabel("‖V‖·‖U‖") + ax.set_title(name, fontsize=10) + fig.tight_layout() + return _render_figure(fig) + + +def plot_mean_component_cis_two_streams( + target_mean_cis: dict[str, np.ndarray], + nontarget_mean_cis: dict[str, np.ndarray], +) -> tuple[bytes, bytes]: + """Both streams' mean CI on one axis per site, ordered by DESCENDING TARGET mean. + + The non-target series is reordered by the same permutation rather than sorted on its + own, so a component's two bars line up vertically — the whole point is reading, per + component, how much on-target importance comes with off-target importance.""" + assert target_mean_cis.keys() == nontarget_mean_cis.keys(), ( + sorted(target_mean_cis), + sorted(nontarget_mean_cis), + ) + n_rows, n_cols = _grid_dims(len(target_mean_cis)) + images: list[bytes] = [] + for log_y in (False, True): + fig = Figure(figsize=(8 * n_cols, 3 * n_rows)) + axs = fig.subplots(n_rows, n_cols, squeeze=False) + flat_axes = axs.T.ravel() + for ax in flat_axes[len(target_mean_cis) :]: + ax.set_visible(False) + for ax, (name, target) in zip(flat_axes, target_mean_cis.items(), strict=False): + order = np.argsort(target)[::-1] + x = range(len(order)) + if log_y: + ax.set_yscale("log") + ax.bar(x, target[order], color="#1f77b4", label="target", width=1.0) + ax.bar( + x, + nontarget_mean_cis[name][order], + color="#d62728", + label="non-target", + width=1.0, + alpha=0.6, + ) + ax.set_xlabel("Component (sorted by target mean CI)") + ax.set_ylabel("mean CI") + ax.set_title(name, fontsize=10) + ax.legend(fontsize=7) + fig.tight_layout() + images.append(_render_figure(fig)) + return images[0], images[1] + + def plot_permuted_ci_heatmaps( position_ci: dict[str, PositionCI], permutation: dict[str, "Literal['identity', 'dense']"] ) -> tuple[bytes, bytes]: diff --git a/param_decomp/core/tests/test_eval_averaging_parity.py b/param_decomp/core/tests/test_eval_averaging_parity.py index ef898f62b..6b02eb038 100644 --- a/param_decomp/core/tests/test_eval_averaging_parity.py +++ b/param_decomp/core/tests/test_eval_averaging_parity.py @@ -132,12 +132,14 @@ def step( jnp.array([0, 0], dtype=jnp.uint32), train_steps=0, eval_steps=2, + stream="broad", ) context = LMEvalContext( state=_state_stub(), # pyright: ignore[reportArgumentType] now_step=0, pass_index=0, batches=(jnp.asarray(1.0), jnp.asarray(3.0)), + target_batches=None, ) assert operation.run(context)["eval/loss/probe/hidden_acts_reconstruction"] == 2.0 diff --git a/param_decomp/core/tests/test_well_temperedness_eval.py b/param_decomp/core/tests/test_well_temperedness_eval.py index 4b5ae9673..666eb1bcc 100644 --- a/param_decomp/core/tests/test_well_temperedness_eval.py +++ b/param_decomp/core/tests/test_well_temperedness_eval.py @@ -46,6 +46,7 @@ def unexpected_render(_ablations: Ablations) -> bytes: mesh=None, compiler_options={}, inputs_for_context=lambda _context: (jnp.zeros((1,)), jax.random.PRNGKey(0)), + log_prefix_for_context=lambda _context: "eval/", figure_rendering=None, ) state = SimpleNamespace(decomposition=SimpleNamespace(components=object(), ci_fn=object())) diff --git a/param_decomp/core/train.py b/param_decomp/core/train.py index 99348b1ad..3d3b6e6db 100644 --- a/param_decomp/core/train.py +++ b/param_decomp/core/train.py @@ -40,7 +40,12 @@ from param_decomp.core.adversary import PersistentAdversary, init_fresh_pgd_sources from param_decomp.core.ci_fn import CI, CIFn, evaluate_ci from param_decomp.core.components import ComponentStacks, VUShape -from param_decomp.core.configs import LossCoeff, SmoothL0ImportanceMinimalityLossConfig +from param_decomp.core.configs import ( + IMP_MIN_METRIC_NAMES, + NONTARGET_STREAM, + LossCoeff, + SmoothL0ImportanceMinimalityLossConfig, +) from param_decomp.core.jit_util import filter_jit from param_decomp.core.losses import ( ReconstructionLoss, @@ -193,12 +198,18 @@ def per_slice_sq(stack: Float[Array, "g a b"]) -> Float[Array, " g"]: def _scheduled_coeff_metrics( - step_f32: Array, total_steps: int, coeffs: dict[str, LossCoeff] + step_f32: Array, total_steps: int, stream_prefix: str, coeffs: dict[str, LossCoeff] ) -> dict[str, Array]: """Per-step values of the SCHEDULED coefficients only — a constant would be log - noise, and a moving coefficient invisible in wandb is a debugging trap.""" + noise, and a moving coefficient invisible in wandb is a debugging trap. + + `stream_prefix` is `""` for the data the run optimizes for and `"nontarget_data/"` for a + targeted run's corpus stream, so the stream segment leads the key exactly as it does on + the loss keys.""" return { - f"schedules/coeff/{name}": scheduled_value_traced(step_f32, total_steps, coeff) + f"{stream_prefix}schedules/coeff/{name}": scheduled_value_traced( + step_f32, total_steps, coeff + ) for name, coeff in coeffs.items() if isinstance(coeff, ScheduleConfig) } @@ -1080,7 +1091,7 @@ def loss_fn( step_f32=step_f32, ) | {"faith": faith_loss} - | _scheduled_coeff_metrics(step_f32, atoms.total_steps, coeff_schedules) + | _scheduled_coeff_metrics(step_f32, atoms.total_steps, "", coeff_schedules) ) return new_state, metrics @@ -1185,13 +1196,18 @@ def make_targeted_train_step[PreparedT]( for term in objective.target.recon if term.hidden_acts_reconstruction is not None }, - "nontarget/impmin": objective.nontarget.impmin_coeff, - **{f"nontarget/{term.name}": term.coeff for term in nt_terms}, } if objective.target.imp.cfg.frequency is not None: coeff_schedules[f"{objective.target.imp.name}/frequency"] = ( objective.target.imp.cfg.frequency.coeff ) + # The non-target imp-min coefficient scales the SAME term the target stream spells by + # name, so it is spelled that way here too — the two streams' coefficients for one + # quantity must be readable as a pair. + nontarget_coeff_schedules: dict[str, LossCoeff] = { + objective.target.imp.name: objective.nontarget.impmin_coeff, + **{term.name: term.coeff for term in nt_terms}, + } def nontarget_draw_loss( model: DecomposedModel[PreparedT], @@ -1360,8 +1376,10 @@ def loss_fn( nt_imp_lp, nt_imp_freq = imp_min_terms(nt_ci.upper, atoms.imp_min, imp_min_param) nt_total = nt_imp_coeff * nt_imp_lp + freq_coeff * nt_imp_freq nt_aux = { - f"loss/nontarget/{atoms.imp_loss_key}": nt_imp_lp, - "loss/nontarget/freq": nt_imp_freq, + # Same spelling as the target stream's keys (which `run._METRIC_KEYS` expands + # from the same map), so one quantity is not two names. + f"{NONTARGET_STREAM}/loss/{IMP_MIN_METRIC_NAMES[atoms.imp_loss_key]}": nt_imp_lp, + f"{NONTARGET_STREAM}/loss/{IMP_MIN_METRIC_NAMES['freq']}": nt_imp_freq, } nt_breakdowns = atoms.grid_losses( nt_terms, @@ -1374,8 +1392,8 @@ def loss_fn( nt_terms, nt_recon_coeffs, nt_breakdowns, strict=True ): nt_total = nt_total + coeff * breakdown.total - nt_aux[f"loss/nontarget/{term.name}"] = breakdown.total - nt_aux["loss/nontarget/total"] = nt_total + nt_aux[f"{NONTARGET_STREAM}/loss/{term.name}"] = breakdown.total + nt_aux[f"{NONTARGET_STREAM}/loss/total"] = nt_total total_loss = total_loss + nt_total reported_total = reported_total + nt_total return total_loss, (reported_total, imp_lp, imp_freq, term_breakdowns, nt_aux) @@ -1445,7 +1463,10 @@ def loss_fn( ) | nt_aux | wd_metrics - | _scheduled_coeff_metrics(step_f32, atoms.total_steps, coeff_schedules) + | _scheduled_coeff_metrics(step_f32, atoms.total_steps, "", coeff_schedules) + | _scheduled_coeff_metrics( + step_f32, atoms.total_steps, f"{NONTARGET_STREAM}/", nontarget_coeff_schedules + ) ) return new_state, metrics diff --git a/param_decomp/core/well_temperedness_eval.py b/param_decomp/core/well_temperedness_eval.py index 3b6fd9043..b60ced38c 100644 --- a/param_decomp/core/well_temperedness_eval.py +++ b/param_decomp/core/well_temperedness_eval.py @@ -32,8 +32,13 @@ well_temperedness_log_entries, ) -_PREFIX = "eval/slow/well_temperedness/" -_FIGURE_KEY = f"{_PREFIX}figures/preactivation_vs_ablation_damage" +_NAMESPACE = "slow/well_temperedness/" +"""The part of the key BELOW the stream namespace. The stream itself is the caller's to +supply: this metric samples the eval distribution, so on a two-stream run it belongs to +whichever stream those batches came from.""" +_FIGURE_STEP_KEY = f"eval/{_NAMESPACE}figure_step" +"""The figure step AXIS, deliberately stream-independent — a second stream must not fork +the axis W&B serializes these renders against (SPEC S28).""" _MAX_FIGURE_LOCATIONS = 48 type FigureRendering = BackgroundRenderer | Literal["synchronous"] | None @@ -98,11 +103,11 @@ def _plot_preactivation_vs_damage(ablations: Ablations) -> bytes: return png_buffer.getvalue() -def _render_deferred(ablations: Ablations, now_step: int) -> DeferredMediaRecord: +def _render_deferred(ablations: Ablations, figure_key: str, now_step: int) -> DeferredMediaRecord: return DeferredMediaRecord( - step_key=f"{_PREFIX}figure_step", + step_key=_FIGURE_STEP_KEY, step=now_step, - media={_FIGURE_KEY: _plot_preactivation_vs_damage(ablations)}, + media={figure_key: _plot_preactivation_vs_damage(ablations)}, ) @@ -114,6 +119,7 @@ def make_well_temperedness_operation[ContextT: EvalInvocation]( mesh: Mesh | None, compiler_options: dict[str, bool | int | str], inputs_for_context: Callable[[ContextT], tuple[Array, PRNGKeyArray]], + log_prefix_for_context: Callable[[ContextT], str], figure_rendering: FigureRendering, ) -> EvalOperation[ContextT]: if figure_rendering is not None: @@ -136,17 +142,19 @@ def run(context: ContextT) -> LogRecord: sampling_key, ) ablations = jax.device_get(device_ablations) + prefix = f"{log_prefix_for_context(context)}{_NAMESPACE}" + figure_key = f"{prefix}figures/preactivation_vs_ablation_damage" log_record: dict[str, float | PNGImage] = { - f"{_PREFIX}{name}": value + f"{prefix}{name}": value for name, value in well_temperedness_log_entries(ablations, site_groups).items() } match figure_rendering: case None: pass case "synchronous": - log_record[_FIGURE_KEY] = PNGImage(_plot_preactivation_vs_damage(ablations)) + log_record[figure_key] = PNGImage(_plot_preactivation_vs_damage(ablations)) case BackgroundRenderer() as renderer: - renderer.submit(partial(_render_deferred, ablations, context.now_step)) + renderer.submit(partial(_render_deferred, ablations, figure_key, context.now_step)) return log_record return EvalOperation(schedule, run) diff --git a/param_decomp/experiments/CLAUDE.md b/param_decomp/experiments/CLAUDE.md index a997b85ce..30e2c4700 100644 --- a/param_decomp/experiments/CLAUDE.md +++ b/param_decomp/experiments/CLAUDE.md @@ -175,6 +175,32 @@ experiments/ └── resid_mlp/ # ResidMLP (CPU): run.py + configs/ + test (target: param_decomp/targets/resid_mlp.py) ``` +## `eval.metrics` on a targeted run — one operation PER STREAM + +`make_lm_evaluation` binds each authored metric to every stream it measures, so a metric is +authored ONCE and a tPD run gets both readouts. The stream set comes from +`target_pool_batches_for`: `None` (the plain root) collapses it to the single broad stream, +so a plain run's operations and keys are untouched. The namespace rule itself is in +`param_decomp/core/CLAUDE.md`. + +| authored metric | streams it binds to | +|---|---| +| `CI_L0`, `PGDReconLoss`, `CIMaskedReconLoss`, `CEandKLLosses` | both on tPD, broad on plain | +| `UnmaskedNoDeltaReconLoss` | the OPTIMIZED stream only | +| attn-patterns / hidden-acts recon, `IdentityCIError`, `WellTemperedness`, the site figures | broad only (labelled `nontarget_data/` on tPD) | +| `WeightMagnitude` | none — reads V/U, no batch | +| `TwoStreamCIMeanPerComponent` | both, in one figure; refuses on a plain run | +| `ArithmeticCIGrid` | none — brings its own probe grid | + +`CIMaskedReconLoss` and `UnmaskedNoDeltaReconLoss` are authorable under `eval.metrics` as +well as `loss_metrics` — the `PGDReconLoss` dual-role pattern (`coeff` null in the eval +seat). As evals they are ONE arm of the CE/KL evaluator (`ce_kl/kl_ci_masked`, +`ce_kl/kl_unmasked`), which is how you get those two numbers without `CEandKLLosses`'s full +11-scalar record — and at one masked forward each instead of six. `UnmaskedNoDeltaReconLoss` +binds to the optimized stream alone deliberately: it is the non-target pass's OWN training +term, already reported there as a train loss, so an eval of it off-target would restate the +objective. + ## Sites and the family grammar Read `param_decomp/core/family.py`'s module docstring before authoring a diff --git a/param_decomp/experiments/eval_config.py b/param_decomp/experiments/eval_config.py index 4545dc296..09add312c 100644 --- a/param_decomp/experiments/eval_config.py +++ b/param_decomp/experiments/eval_config.py @@ -14,6 +14,7 @@ CI_L0Config, CIHiddenActsReconLossConfig, CIHistogramsConfig, + CIMaskedReconLossConfig, CIMeanPerComponentConfig, ComponentActivationDensityConfig, HiddenActsReconstructionMixin, @@ -22,7 +23,9 @@ PermutedCIPlotsConfig, PGDReconLossConfig, StochasticHiddenActsReconLossConfig, + UnmaskedNoDeltaReconLossConfig, UVPlotsConfig, + WeightMagnitudeConfig, WellTemperednessConfig, ) from param_decomp.core.eval_schedule import EvalSchedule, Every, FirstThenEvery @@ -31,6 +34,7 @@ CEandKLLossesConfig, CIMaskedAttnPatternsReconLossConfig, StochasticAttnPatternsReconLossConfig, + TwoStreamCIMeanPerComponentConfig, ) AnyEvalMetricConfig = Annotated[ @@ -40,6 +44,7 @@ | CIHistogramsConfig | CI_L0Config | CIMaskedAttnPatternsReconLossConfig + | CIMaskedReconLossConfig | CIMeanPerComponentConfig | ComponentActivationDensityConfig | IdentityCIErrorConfig @@ -47,7 +52,10 @@ | PGDReconLossConfig | StochasticAttnPatternsReconLossConfig | StochasticHiddenActsReconLossConfig + | TwoStreamCIMeanPerComponentConfig + | UnmaskedNoDeltaReconLossConfig | UVPlotsConfig + | WeightMagnitudeConfig | WellTemperednessConfig, Discriminator("type"), ] diff --git a/param_decomp/experiments/lm/diagnostic_eval_operations.py b/param_decomp/experiments/lm/diagnostic_eval_operations.py index da648c1ec..aa65667b2 100644 --- a/param_decomp/experiments/lm/diagnostic_eval_operations.py +++ b/param_decomp/experiments/lm/diagnostic_eval_operations.py @@ -6,6 +6,7 @@ import numpy as np from jaxtyping import PRNGKeyArray +from param_decomp.core.ci_fn import CIFn from param_decomp.core.configs import ( CIHiddenActsReconLossConfig, CIHistogramsConfig, @@ -40,9 +41,13 @@ compute_identity_ci_errors, make_position_ci_step, make_slow_eval_step, + mean_cis, + plot_mean_component_cis_two_streams, + plot_weight_magnitudes, render_permutation_figures, render_slow_eval_figures, resolve_permutation_metrics, + weight_magnitudes, ) from param_decomp.experiments.lm.attn_patterns_eval import ( accumulate_attn_patterns, @@ -56,17 +61,19 @@ ) from param_decomp.experiments.lm.eval_context import LMEvalContext from param_decomp.experiments.lm.eval_keys import EvalKeyStream +from param_decomp.experiments.lm.scalar_eval_operations import stream_batches, stream_log_prefix + + +def _figure_record(now_step: int, media: dict[str, bytes]) -> DeferredMediaRecord: + """A slow-tier figure batch on the dedicated figure-step axis (SPEC S28).""" + return DeferredMediaRecord(step_key="slow_eval/figure_step", step=now_step, media=media) def _render_selected_figures( reductions: dict[str, SiteReduction], wanted: set[str], now_step: int ) -> DeferredMediaRecord: figures = render_slow_eval_figures(reductions) - return DeferredMediaRecord( - step_key="slow_eval/figure_step", - step=now_step, - media={f"slow_eval/{name}": figures[name] for name in wanted}, - ) + return _figure_record(now_step, {f"slow_eval/{name}": figures[name] for name in wanted}) def _render_permutation( @@ -79,11 +86,7 @@ def _render_permutation( figures = render_permutation_figures(spec, position_ci, components) if not include_ci_heatmaps: figures = {key: value for key, value in figures.items() if key == "figures/uv_matrices"} - return DeferredMediaRecord( - step_key="slow_eval/figure_step", - step=now_step, - media={f"slow_eval/{name}": value for name, value in figures.items()}, - ) + return _figure_record(now_step, {f"slow_eval/{name}": value for name, value in figures.items()}) def make_attention_operation( @@ -115,7 +118,7 @@ def run(context: LMEvalContext) -> LogRecord: ), ) return { - f"eval/loss/{name}": value + f"{stream_log_prefix('broad', context)}loss/{name}": value for name, value in attn_patterns_log_entries(metric.type, reductions).items() } @@ -151,13 +154,79 @@ def run(context: LMEvalContext) -> LogRecord: ), ) return { - f"eval/slow/loss/{name}": value + f"{stream_log_prefix('broad', context)}slow/loss/{name}": value for name, value in hidden_acts_log_entries(metric.type, reductions).items() } return EvalOperation(schedule, run) +def _render_weight_magnitudes( + magnitudes: dict[str, np.ndarray], now_step: int +) -> DeferredMediaRecord: + return _figure_record( + now_step, {"slow_eval/figures/weight_magnitude": plot_weight_magnitudes(magnitudes)} + ) + + +def make_weight_magnitude_operation( + schedule: EvalSchedule, renderer: BackgroundRenderer +) -> EvalOperation[LMEvalContext]: + """`‖V_c‖·‖U_c‖` per site. Reads the trained V/U only — no model, no batch, no step.""" + + def run(context: LMEvalContext) -> LogRecord: + magnitudes = weight_magnitudes(context.state.decomposition.components) + renderer.submit(partial(_render_weight_magnitudes, magnitudes, context.now_step)) + return {} + + return EvalOperation(schedule, run) + + +def _render_two_stream_ci_means( + target: dict[str, np.ndarray], + nontarget: dict[str, np.ndarray], + now_step: int, +) -> DeferredMediaRecord: + linear, log = plot_mean_component_cis_two_streams(target, nontarget) + return _figure_record( + now_step, + { + "slow_eval/figures/ci_mean_per_component_two_streams": linear, + "slow_eval/figures/ci_mean_per_component_two_streams_log": log, + }, + ) + + +def make_two_stream_ci_mean_operation( + schedule: EvalSchedule, + model: DecomposedModel, + ci_capture_keys: CaptureKeys, + compiler_options: dict[str, bool | int | str], + renderer: BackgroundRenderer, +) -> EvalOperation[LMEvalContext]: + """Both streams' mean CI per component in one figure, ordered by the target mean.""" + step = make_slow_eval_step(model, ci_capture_keys, 0.0, None, compiler_options) + + def stream_mean_cis(ci_fn: CIFn, batches: tuple[jax.Array, ...]) -> dict[str, np.ndarray]: + # `n_batches_accum=0`: this metric reads only `ci_sums`, and the raw-value sample + # would gather every position's CI to the host (~430MB/pass here) to be discarded. + return mean_cis(accumulate_site_reductions(step, model, ci_fn, list(batches), 0)) + + def run(context: LMEvalContext) -> LogRecord: + ci_fn = context.state.decomposition.ci_fn + renderer.submit( + partial( + _render_two_stream_ci_means, + stream_mean_cis(ci_fn, stream_batches("target_data", context)), + stream_mean_cis(ci_fn, stream_batches("broad", context)), + context.now_step, + ) + ) + return {} + + return EvalOperation(schedule, run) + + def make_site_figures_operation( metric: CIHistogramsConfig | ComponentActivationDensityConfig | CIMeanPerComponentConfig, schedule: EvalSchedule, @@ -226,7 +295,8 @@ def run(context: LMEvalContext) -> LogRecord: match metric: case IdentityCIErrorConfig(): errors = compute_identity_ci_errors(spec, position_ci, IDENTITY_CI_ERROR_TOLERANCE) - return {f"eval/slow/{name}": value for name, value in errors.items()} + prefix = stream_log_prefix("broad", context) + return {f"{prefix}slow/{name}": value for name, value in errors.items()} case UVPlotsConfig(): include_ci_heatmaps = False components = { diff --git a/param_decomp/experiments/lm/eval.py b/param_decomp/experiments/lm/eval.py index fdd09b804..f9d157045 100644 --- a/param_decomp/experiments/lm/eval.py +++ b/param_decomp/experiments/lm/eval.py @@ -204,17 +204,39 @@ def _ce[PreparedT](batch: _PreparedLMBatch[PreparedT], logits: Array) -> Array: return _row_masked_cross_entropy(logits, batch.tokens, batch.valid_row_mask) +CE_KL_VARIANTS: tuple[str, ...] = ( + "ci_masked", + "unmasked", + "stoch_masked", + "random_masked", + "rounded_masked", + "zero_masked", +) +"""Every masking arm `make_ce_kl_step` can evaluate. Each costs ONE masked forward, so a +caller that wants a single number asks for a single arm rather than filtering the record +afterwards.""" + + def make_ce_kl_step[PreparedT]( model_static: DecomposedModel[PreparedT], ci_capture_keys: CaptureKeys, rounding_threshold: float, + variants: tuple[str, ...], + emit_ce_difference: bool, mesh: Mesh | None = None, compiler_options: dict[str, bool | int | str] | None = None, *, n_valid_rows: int | None = None, ) -> ScalarStep: - """Build the single-purpose CE/KL evaluator.""" + """Build the single-purpose CE/KL evaluator over `variants` (a subset of + `CE_KL_VARIANTS`). + + The masks for EVERY arm are drawn whether or not the arm is selected, so a narrowed + evaluator's numbers are bit-identical to the full one's — only the forwards are + skipped, and XLA drops the unused draws. `emit_ce_difference` adds the CE-vs-target + delta for each selected arm (free: same logits, no extra forward).""" assert model_static.has_position_axis, "CEandKLLosses is LM-only and requires a position axis" + assert variants and set(variants) <= set(CE_KL_VARIANTS), variants def eval_step( model: DecomposedModel[PreparedT], @@ -239,7 +261,7 @@ def eval_step( batch.tokens.shape, COMPUTE_DT, ) - variants = { + variant_masks = { "ci_masked": (batch.ci_lower, zeros_delta), "unmasked": ( {site: jnp.ones_like(batch.ci_lower[site]) for site in model.site_names}, @@ -269,19 +291,21 @@ def eval_step( } variant_logits = { name: _compute_masked_output(model, batch, masks, deltas, mesh, frozenset()) - for name, (masks, deltas) in variants.items() + for name, (masks, deltas) in variant_masks.items() + if name in variants } - target_ce = _ce(batch, batch.clean.output) metrics = { f"ce_kl/kl_{name}": _kl(batch, logits) for name, logits in variant_logits.items() } - metrics.update( - { - f"ce_kl/ce_difference_{name}": _ce(batch, variant_logits[name]) - target_ce - for name in variants - if name != "zero_masked" - } - ) + if emit_ce_difference: + target_ce = _ce(batch, batch.clean.output) + metrics.update( + { + f"ce_kl/ce_difference_{name}": _ce(batch, variant_logits[name]) - target_ce + for name in variant_logits + if name != "zero_masked" + } + ) return metrics return filter_jit(eval_step, compiler_options=compiler_options) @@ -435,6 +459,8 @@ def make_eval_step[PreparedT]( model_static, ci_capture_keys, rounding_threshold, + CE_KL_VARIANTS, + True, mesh, compiler_options, n_valid_rows=n_valid_rows, diff --git a/param_decomp/experiments/lm/eval_config.py b/param_decomp/experiments/lm/eval_config.py index b850743d7..89f85a24b 100644 --- a/param_decomp/experiments/lm/eval_config.py +++ b/param_decomp/experiments/lm/eval_config.py @@ -62,3 +62,15 @@ class ArithmeticCIGridConfig(BaseConfig): b_range: tuple[int, int] = (1, 100) thresholds: list[float] = Field(default_factory=lambda: [0.1]) top_k: PositiveInt = 24 + + +class TwoStreamCIMeanPerComponentConfig(BaseConfig): + """Both streams' mean CI per component on ONE axis per site, ordered by descending + TARGET mean and coloured by stream. + + Supersedes authoring `CIMeanPerComponent` on a targeted run: it computes the same + broad-stream reduction plus the target-pool one, so authoring both would pay for the + broad pass twice. Refuses on a plain run, which has no target stream.""" + + slow: ClassVar[bool] = True + type: Literal["TwoStreamCIMeanPerComponent"] = "TwoStreamCIMeanPerComponent" diff --git a/param_decomp/experiments/lm/eval_context.py b/param_decomp/experiments/lm/eval_context.py index cc674bafa..15d4ad3af 100644 --- a/param_decomp/experiments/lm/eval_context.py +++ b/param_decomp/experiments/lm/eval_context.py @@ -11,3 +11,8 @@ class LMEvalContext(EvalInvocation): pass_index: int batches: tuple[jax.Array, ...] + target_batches: tuple[jax.Array, ...] | None + """A tPD run's prompt-pool draws, which `data.eval` cannot supply — the pool has no + held-out split, so the targeted root draws them exactly as training does. `None` on a + plain run, which HAS no second stream; that `None` is also what tells every log key + which run kind it is in (`scalar_eval_operations.stream_log_prefix`).""" diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index 0917f4625..290bec5ff 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -1,5 +1,8 @@ """LM evaluation operation binding and execution.""" +from collections.abc import Callable +from functools import partial + import jax import numpy as np from jax.sharding import Mesh, NamedSharding @@ -10,13 +13,16 @@ CI_L0Config, CIHiddenActsReconLossConfig, CIHistogramsConfig, + CIMaskedReconLossConfig, CIMeanPerComponentConfig, ComponentActivationDensityConfig, IdentityCIErrorConfig, PermutedCIPlotsConfig, PGDReconLossConfig, StochasticHiddenActsReconLossConfig, + UnmaskedNoDeltaReconLossConfig, UVPlotsConfig, + WeightMagnitudeConfig, WellTemperednessConfig, ) from param_decomp.core.model import DecomposedModel @@ -35,20 +41,26 @@ make_hidden_acts_operation, make_permutation_operation, make_site_figures_operation, + make_two_stream_ci_mean_operation, + make_weight_magnitude_operation, ) from param_decomp.experiments.lm.eval_config import ( ArithmeticCIGridConfig, CEandKLLossesConfig, CIMaskedAttnPatternsReconLossConfig, StochasticAttnPatternsReconLossConfig, + TwoStreamCIMeanPerComponentConfig, ) from param_decomp.experiments.lm.eval_context import LMEvalContext from param_decomp.experiments.lm.eval_keys import EvalKeyStream from param_decomp.experiments.lm.resolved import LMAnyRun from param_decomp.experiments.lm.scalar_eval_operations import ( + Stream, make_ce_kl_operation, make_ci_l0_operation, make_fresh_pgd_operation, + make_single_variant_kl_operation, + stream_log_prefix, ) from param_decomp.infra.dataset_store import read_dataset_meta from param_decomp.pretrain.batch_data import BatchSchedule, ShardServer, scan_shards @@ -68,8 +80,16 @@ def make_lm_evaluation( n_proc: int, sink: MetricsSink, compiler_options: dict[str, bool | int | str], + target_pool_batches_for: Callable[[int], list[jax.Array]] | None, ) -> Evaluation[LMEvalContext]: - """Construct one executable operation for every authored LM metric.""" + """Construct the executable operations for every authored LM metric — one PER STREAM + the metric measures. + + `target_pool_batches_for(pass_index)` supplies the tPD target stream, which `data.eval` + cannot: the prompt pool has no held-out split, so the targeted root draws pool batches + the same way training does. `None` on a plain run, and that is what makes a plain run's + metric set — and every one of its log keys — exactly what it was before targeted runs + existed: `data_streams` collapses to the single broad stream.""" pd = built.pd capture_inputs = built.ci_fn.capture_keys data = built.data @@ -94,95 +114,183 @@ def well_temperedness_inputs( run_key, EvalKeyStream.WELL_TEMPEREDNESS * pd.steps + context.pass_index ) - def make_operation(metric: AnyEvalMetricConfig) -> EvalOperation[LMEvalContext]: + targeted = target_pool_batches_for is not None + data_streams: tuple[Stream, ...] = ("broad", "target_data") if targeted else ("broad",) + """Every stream a data-dependent metric measures. A plain run has exactly one.""" + optimized_stream: tuple[Stream, ...] = ("target_data",) if targeted else ("broad",) + """Just the stream the run optimizes for — what a metric that is only meaningful there + binds to. On a plain run that IS the broad stream, so such a metric stays authorable.""" + + def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalContext], ...]: schedule = schedule_for(metric, eval) match metric: case CEandKLLossesConfig(): - return make_ce_kl_operation( - metric, - schedule, - model, - capture_inputs, - run_key, - pd.steps, - eval.n_steps, - mesh, - compiler_options, + return tuple( + make_ce_kl_operation( + metric, + schedule, + stream, + model, + capture_inputs, + run_key, + pd.steps, + eval.n_steps, + mesh, + compiler_options, + ) + for stream in data_streams + ) + case CIMaskedReconLossConfig() | UnmaskedNoDeltaReconLossConfig(): + # The unmasked arm is the non-target pass's OWN training term, already + # reported as a train loss there; measuring it again off-target would + # restate the objective, so it binds to the optimized stream alone. + streams = ( + optimized_stream + if isinstance(metric, UnmaskedNoDeltaReconLossConfig) + else data_streams + ) + return tuple( + make_single_variant_kl_operation( + metric, + schedule, + stream, + model, + capture_inputs, + run_key, + pd.steps, + eval.n_steps, + mesh, + compiler_options, + ) + for stream in streams ) case CI_L0Config(): - return make_ci_l0_operation( - metric, - schedule, - model, - capture_inputs, - run_key, - pd.steps, - eval.n_steps, - mesh, - compiler_options, + return tuple( + make_ci_l0_operation( + metric, + schedule, + stream, + model, + capture_inputs, + run_key, + pd.steps, + eval.n_steps, + mesh, + compiler_options, + ) + for stream in data_streams ) case PGDReconLossConfig(): - return make_fresh_pgd_operation( - metric, - schedule, - model, - capture_inputs, - run_key, - pd.steps, - eval.n_steps, - mesh, - compiler_options, + return tuple( + make_fresh_pgd_operation( + metric, + schedule, + stream, + model, + capture_inputs, + run_key, + pd.steps, + eval.n_steps, + mesh, + compiler_options, + ) + for stream in data_streams ) case CIMaskedAttnPatternsReconLossConfig() | StochasticAttnPatternsReconLossConfig(): - return make_attention_operation( - metric, schedule, model, capture_inputs, run_key, pd.steps, compiler_options + return ( + make_attention_operation( + metric, schedule, model, capture_inputs, run_key, pd.steps, compiler_options + ), ) case CIHiddenActsReconLossConfig() | StochasticHiddenActsReconLossConfig(): - return make_hidden_acts_operation( - metric, schedule, model, capture_inputs, run_key, pd.steps, compiler_options + return ( + make_hidden_acts_operation( + metric, schedule, model, capture_inputs, run_key, pd.steps, compiler_options + ), ) case ( CIHistogramsConfig() | ComponentActivationDensityConfig() | CIMeanPerComponentConfig() ): - return make_site_figures_operation( - metric, schedule, model, capture_inputs, compiler_options, renderer + return ( + make_site_figures_operation( + metric, schedule, model, capture_inputs, compiler_options, renderer + ), ) case PermutedCIPlotsConfig() | UVPlotsConfig() | IdentityCIErrorConfig(): - return make_permutation_operation( - metric, schedule, model, capture_inputs, compiler_options, renderer + return ( + make_permutation_operation( + metric, schedule, model, capture_inputs, compiler_options, renderer + ), ) case WellTemperednessConfig(): - return make_well_temperedness_operation( - metric, - schedule, - model, - capture_inputs, - mesh, - compiler_options, - inputs_for_context=well_temperedness_inputs, - figure_rendering=renderer if sink.accepts_deferred_media else None, + return ( + make_well_temperedness_operation( + metric, + schedule, + model, + capture_inputs, + mesh, + compiler_options, + inputs_for_context=well_temperedness_inputs, + # It samples `context.batches` — the broad stream — so it carries + # that stream's namespace rather than the bare one. + log_prefix_for_context=partial(stream_log_prefix, "broad"), + figure_rendering=renderer if sink.accepts_deferred_media else None, + ), ) case ArithmeticCIGridConfig(): - return make_arithmetic_operation( - metric, - schedule, - built.target, - model, - capture_inputs, - mesh, - n_proc, - sink, - run_key, - pd.steps, - compiler_options, + return ( + make_arithmetic_operation( + metric, + schedule, + built.target, + model, + capture_inputs, + mesh, + n_proc, + sink, + run_key, + pd.steps, + compiler_options, + ), + ) + case TwoStreamCIMeanPerComponentConfig(): + return ( + make_two_stream_ci_mean_operation( + schedule, model, capture_inputs, compiler_options, renderer + ), ) + case WeightMagnitudeConfig(): + return (make_weight_magnitude_operation(schedule, renderer),) - operations = tuple(make_operation(metric) for metric in eval.metrics) + needs_target_stream = tuple( + metric.type + for metric in eval.metrics + if isinstance(metric, TwoStreamCIMeanPerComponentConfig) + ) + assert not needs_target_stream or targeted, ( + f"{needs_target_stream} measure the tPD target stream; a plain run has no prompt pool" + ) + # The single-arm evals ARE arms of `CEandKLLosses`, under the same keys. Authoring both + # is a duplicate measurement that `_run_due_evaluation`'s collision assert would catch + # at the first eval pass — hours into a run. Catch it while reading the config. + single_arm = tuple( + metric.type + for metric in eval.metrics + if isinstance(metric, CIMaskedReconLossConfig | UnmaskedNoDeltaReconLossConfig) + ) + assert not (single_arm and any(isinstance(m, CEandKLLossesConfig) for m in eval.metrics)), ( + f"{single_arm} emit arms CEandKLLosses already emits, under the same `ce_kl/kl_*` " + "keys: author the narrow metrics OR CEandKLLosses, not both" + ) + operations = tuple( + operation for metric in eval.metrics for operation in make_operations(metric) + ) def make_context(state: TrainState, now_step: int) -> LMEvalContext: pass_index = now_step // eval.every @@ -191,6 +299,11 @@ def make_context(state: TrainState, now_step: int) -> LMEvalContext: now_step=now_step, pass_index=pass_index, batches=tuple(batches(pass_index)), + target_batches=( + None + if target_pool_batches_for is None + else tuple(target_pool_batches_for(pass_index)) + ), ) return Evaluation(operations, make_context) diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index 6a7425728..7dd9a5219 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -1,11 +1,19 @@ """Independent CE/KL, causal-L0, and fresh-PGD LM operations.""" +from typing import Literal + import jax.numpy as jnp from jax import random from jax.sharding import Mesh from jaxtyping import Array, PRNGKeyArray -from param_decomp.core.configs import CI_L0Config, PGDReconLossConfig +from param_decomp.core.configs import ( + NONTARGET_STREAM, + CI_L0Config, + CIMaskedReconLossConfig, + PGDReconLossConfig, + UnmaskedNoDeltaReconLossConfig, +) from param_decomp.core.eval_schedule import EvalSchedule from param_decomp.core.metrics import BarChart, LogRecord from param_decomp.core.model import CaptureKeys, DecomposedModel @@ -13,6 +21,7 @@ from param_decomp.core.recon_eval import FreshPGDReconEval from param_decomp.core.run import EvalOperation from param_decomp.experiments.lm.eval import ( + CE_KL_VARIANTS, ScalarStep, make_ce_kl_step, make_ci_l0_step, @@ -22,6 +31,41 @@ from param_decomp.experiments.lm.eval_context import LMEvalContext from param_decomp.experiments.lm.eval_keys import EvalKeyStream +type Stream = Literal["broad", "target_data"] +"""Which DATA an eval operation measures. ONE value, not a (batch source, log prefix) pair: +the two always covary, and nothing should be able to spell target-pool batches under the +broad stream's log keys.""" + + +def stream_batches(stream: Stream, context: LMEvalContext) -> tuple[Array, ...]: + match stream: + case "broad": + return context.batches + case "target_data": + assert context.target_batches is not None, ( + "target-stream metrics need a tPD run's prompt pool; a plain run has none" + ) + return context.target_batches + + +def stream_log_prefix(stream: Stream, context: LMEvalContext) -> str: + """The log namespace for `stream`, given what kind of run this is. + + ONE rule: the data the run is optimizing for is unlabelled, and anything else carries + its stream. A plain run has a single stream, so it stays `eval/` exactly as before — + `context.target_batches is None` is what says so. A tPD run adds a second stream, so its + broad corpus moves under `eval/nontarget_data/` and the target pool takes the bare + namespace. The consequence is deliberate: `eval/l0/...` means "the data of interest" in + both run kinds, which is what makes the two comparable as objectives — but it is NOT the + same data, so a corpus-vs-corpus comparison across run kinds must read + `eval/nontarget_data/` on the tPD side.""" + targeted = context.target_batches is not None + match stream: + case "broad": + return f"eval/{NONTARGET_STREAM}/" if targeted else "eval/" + case "target_data": + return "eval/" + def _make_scalar_operation( schedule: EvalSchedule, @@ -31,10 +75,12 @@ def _make_scalar_operation( run_key: PRNGKeyArray, train_steps: int, eval_steps: int, + stream: Stream, ) -> EvalOperation[LMEvalContext]: def run(context: LMEvalContext) -> LogRecord: + log_prefix = stream_log_prefix(stream, context) sums: dict[str, Array] = {} - for batch_index, tokens in enumerate(context.batches): + for batch_index, tokens in enumerate(stream_batches(stream, context)): key = random.fold_in( run_key, EvalKeyStream.SCALARS * train_steps + context.pass_index * eval_steps + batch_index, @@ -49,7 +95,7 @@ def run(context: LMEvalContext) -> LogRecord: for name, value in values.items(): if name.startswith(prefixes): sums[name] = sums.get(name, jnp.zeros(())) + value - return {f"eval/{name}": float(value) / eval_steps for name, value in sums.items()} + return {f"{log_prefix}{name}": float(value) / eval_steps for name, value in sums.items()} return EvalOperation(schedule, run) @@ -57,6 +103,7 @@ def run(context: LMEvalContext) -> LogRecord: def make_ce_kl_operation( metric: CEandKLLossesConfig, schedule: EvalSchedule, + stream: Stream, model: DecomposedModel, ci_capture_keys: CaptureKeys, run_key: PRNGKeyArray, @@ -65,22 +112,74 @@ def make_ce_kl_operation( mesh: Mesh, compiler_options: dict[str, bool | int | str], ) -> EvalOperation[LMEvalContext]: - scalars = _make_scalar_operation( + return _make_scalar_operation( schedule, - make_ce_kl_step(model, ci_capture_keys, metric.rounding_threshold, mesh, compiler_options), + make_ce_kl_step( + model, + ci_capture_keys, + metric.rounding_threshold, + CE_KL_VARIANTS, + True, + mesh, + compiler_options, + ), ("ce_kl/",), model, run_key, train_steps, eval_steps, + stream, ) - return scalars + +def make_single_variant_kl_operation( + metric: CIMaskedReconLossConfig | UnmaskedNoDeltaReconLossConfig, + schedule: EvalSchedule, + stream: Stream, + model: DecomposedModel, + ci_capture_keys: CaptureKeys, + run_key: PRNGKeyArray, + train_steps: int, + eval_steps: int, + mesh: Mesh, + compiler_options: dict[str, bool | int | str], +) -> EvalOperation[LMEvalContext]: + """ONE masking arm of the CE/KL evaluator, authored as the loss config that names the + same construction. + + These two recon configs are authorable as evals exactly as `PGDReconLoss` already is — + the eval measures the quantity the loss optimizes. The key keeps the `ce_kl/kl_` + spelling a plain run's `CEandKLLosses` logs it under, so the same number is one name + across run kinds; the config type is what names the construction explicitly.""" + match metric: + case CIMaskedReconLossConfig(): + variant = "ci_masked" + case UnmaskedNoDeltaReconLossConfig(): + variant = "unmasked" + return _make_scalar_operation( + schedule, + make_ce_kl_step( + model, + ci_capture_keys, + 0.0, # unused: the rounded arm is not among the selected variants + (variant,), + False, + mesh, + compiler_options, + ), + (f"ce_kl/kl_{variant}",), + model, + run_key, + train_steps, + eval_steps, + stream, + ) def make_ci_l0_operation( metric: CI_L0Config, schedule: EvalSchedule, + stream: Stream, model: DecomposedModel, ci_capture_keys: CaptureKeys, run_key: PRNGKeyArray, @@ -104,12 +203,14 @@ def make_ci_l0_operation( run_key, train_steps, eval_steps, + stream, ) def run(context: LMEvalContext) -> LogRecord: record = dict(scalars.run(context)) - prefix = f"eval/l0/{metric.ci_alive_threshold}_" - record["eval/l0/bar_chart"] = BarChart( + log_prefix = stream_log_prefix(stream, context) + prefix = f"{log_prefix}l0/{metric.ci_alive_threshold}_" + record[f"{log_prefix}l0/bar_chart"] = BarChart( rows=tuple( (name.removeprefix(prefix), value) for name, value in record.items() @@ -127,6 +228,7 @@ def run(context: LMEvalContext) -> LogRecord: def make_fresh_pgd_operation( metric: PGDReconLossConfig, schedule: EvalSchedule, + stream: Stream, model: DecomposedModel, ci_capture_keys: CaptureKeys, run_key: PRNGKeyArray, @@ -150,4 +252,5 @@ def make_fresh_pgd_operation( run_key, train_steps, eval_steps, + stream, ) diff --git a/param_decomp/experiments/lm/test_eval_operations.py b/param_decomp/experiments/lm/test_eval_operations.py index 7303a5fcd..808ee7387 100644 --- a/param_decomp/experiments/lm/test_eval_operations.py +++ b/param_decomp/experiments/lm/test_eval_operations.py @@ -47,6 +47,7 @@ def make_operation( compiler_options: dict[str, bool | int | str], *, inputs_for_context: Any, + log_prefix_for_context: Any, figure_rendering: Any, ) -> EvalOperation[Any]: captured.update( @@ -56,6 +57,7 @@ def make_operation( mesh=mesh, compiler_options=compiler_options, inputs_for_context=inputs_for_context, + log_prefix_for_context=log_prefix_for_context, figure_rendering=figure_rendering, ) return EvalOperation(schedule, lambda _context: {}) @@ -93,6 +95,7 @@ def make_operation( n_proc=1, sink=cast(Any, SimpleNamespace(accepts_deferred_media=True)), compiler_options={}, + target_pool_batches_for=None, ) batch = jnp.arange(4) _, key = captured["inputs_for_context"]( @@ -101,6 +104,7 @@ def make_operation( now_step=30, pass_index=3, batches=(batch,), + target_batches=None, ) ) diff --git a/param_decomp/experiments/lm/test_stream_log_namespace.py b/param_decomp/experiments/lm/test_stream_log_namespace.py new file mode 100644 index 000000000..b3f489c9d --- /dev/null +++ b/param_decomp/experiments/lm/test_stream_log_namespace.py @@ -0,0 +1,124 @@ +"""The stream namespace rule: the data a run optimizes for is unlabelled. + +A plain run has one stream and keeps `eval/`; a tPD run adds a second, so its target pool +takes the bare namespace and the broad corpus moves under `eval/nontarget_data/`. The plain +case is pinned here because it is a COMPATIBILITY GUARANTEE — those keys are shared with +every non-targeted run and both toys, and adding a second stream must not move them. +""" + +from types import SimpleNamespace +from typing import Any, cast + +import jax.numpy as jnp +from jaxtyping import Array + +from param_decomp.core.configs import ( + CIMaskedReconLossConfig, + UnmaskedNoDeltaReconLossConfig, +) +from param_decomp.core.eval_schedule import Every +from param_decomp.core.run import EvalOperation +from param_decomp.core.train import TrainState +from param_decomp.experiments.lm.eval import CE_KL_VARIANTS +from param_decomp.experiments.lm.eval_context import LMEvalContext +from param_decomp.experiments.lm.scalar_eval_operations import ( + Stream, + _make_scalar_operation, + stream_log_prefix, +) + +TARGET_BATCHES = (jnp.zeros((1, 4), jnp.int32),) + + +def context(target_batches: tuple[jnp.ndarray, ...] | None) -> LMEvalContext: + """A context carrying only what the namespace rule reads; `state` is never touched.""" + return LMEvalContext( + state=cast(TrainState, cast(object, None)), + now_step=0, + pass_index=0, + batches=(), + target_batches=target_batches, + ) + + +def test_plain_run_keys_are_unchanged_by_the_two_stream_rule(): + assert stream_log_prefix("broad", context(None)) == "eval/" + + +def test_targeted_run_labels_the_broad_stream_and_leaves_the_target_bare(): + targeted = context(TARGET_BATCHES) + assert stream_log_prefix("target_data", targeted) == "eval/" + assert stream_log_prefix("broad", targeted) == "eval/nontarget_data/" + + +def test_the_bare_namespace_is_always_the_data_of_interest(): + """What makes the run kinds comparable AS OBJECTIVES — and the trap that comes with it: + the bare namespace is not the same DATA in both.""" + plain, targeted = context(None), context(TARGET_BATCHES) + assert stream_log_prefix("broad", plain) == stream_log_prefix("target_data", targeted) + assert stream_log_prefix("broad", targeted) != stream_log_prefix("broad", plain) + + +def test_every_stream_resolves_under_both_run_kinds(): + for stream in cast(tuple[Stream, ...], ("broad", "target_data")): + for target_batches in (None, TARGET_BATCHES): + prefix = stream_log_prefix(stream, context(target_batches)) + assert prefix.startswith("eval/") and prefix.endswith("/") + + +def test_one_operation_per_stream_reads_that_stream_and_labels_it(): + """The end-to-end shape the binder relies on: bind the SAME step twice, once per + stream, and the two operations read different batches and land under different keys — + which is what lets a metric be authored once and measured on both.""" + + def step( + _model: Any, _components: Any, _ci_fn: Any, value: Array, _key: Any + ) -> dict[str, Array]: + return {"l0/0.0_site": value} + + def operation_for(stream: Stream) -> EvalOperation[LMEvalContext]: + return _make_scalar_operation( + Every(1), + step, + ("l0/",), + cast(Any, object()), + jnp.array([0, 0], dtype=jnp.uint32), + train_steps=0, + eval_steps=1, + stream=stream, + ) + + targeted = LMEvalContext( + state=cast( + TrainState, + cast( + object, SimpleNamespace(decomposition=SimpleNamespace(components=None, ci_fn=None)) + ), + ), + now_step=0, + pass_index=0, + batches=(jnp.asarray(3.0),), + target_batches=(jnp.asarray(7.0),), + ) + broad = operation_for("broad").run(targeted) + target = operation_for("target_data").run(targeted) + + assert broad == {"eval/nontarget_data/l0/0.0_site": 3.0} + assert target == {"eval/l0/0.0_site": 7.0} + assert not broad.keys() & target.keys() + + +def test_narrow_kl_evals_name_arms_the_full_ce_kl_evaluator_also_emits(): + """The narrow metrics exist to cut CE/KL's clutter, NOT to respell its numbers: each + selects one arm of the same evaluator, so `eval/ce_kl/kl_` means the same thing + whether a run authored `CEandKLLosses` or the single-arm config. A new arm name here + that the full evaluator doesn't emit would silently give one quantity two names.""" + assert {"ci_masked", "unmasked"} <= set(CE_KL_VARIANTS) + + +def test_the_two_recon_configs_are_authorable_as_evals(): + """`coeff` is what separates the loss role from the eval role (`LossMetricConfig`), so + both must validate with none — this is the `PGDReconLoss` dual-role pattern.""" + for config in (CIMaskedReconLossConfig(), UnmaskedNoDeltaReconLossConfig()): + assert config.coeff is None + assert config.slow is False diff --git a/param_decomp/experiments/lm/training.py b/param_decomp/experiments/lm/training.py index 9a2f0337a..dd1b877db 100644 --- a/param_decomp/experiments/lm/training.py +++ b/param_decomp/experiments/lm/training.py @@ -163,7 +163,15 @@ def sample_batch(step: int) -> jax.Array: "mid-window eval would corrupt the next step-time estimate" ) evaluation = make_lm_evaluation( - built, eval_config, model, run_key, mesh, n_proc, sink, runtime.compiler_options + built, + eval_config, + model, + run_key, + mesh, + n_proc, + sink, + runtime.compiler_options, + target_pool_batches_for=None, # a plain run has ONE stream ) run_decomposition_training( diff --git a/param_decomp/experiments/lm/training_targeted.py b/param_decomp/experiments/lm/training_targeted.py index 2019c4288..aa3d760bd 100644 --- a/param_decomp/experiments/lm/training_targeted.py +++ b/param_decomp/experiments/lm/training_targeted.py @@ -119,12 +119,15 @@ def train_targeted( server.per_process, jax.local_device_count(), ) - per_process_target = target_batch // n_proc + + def pool_global_batch(seed: int, step: int, batch: int) -> jax.Array: + per_process = batch // n_proc + rows = pool_batch(pool, seed, step, batch) + local = rows[jax.process_index() * per_process :][:per_process] + return global_token_batch(local, mesh, batch) def sample_target_batch(step: int) -> jax.Array: - rows = pool_batch(pool, built.pd.seed, step, target_batch) - local = rows[jax.process_index() * per_process_target :][:per_process_target] - return global_token_batch(local, mesh, target_batch) + return pool_global_batch(built.pd.seed, step, target_batch) def sample_nontarget_batch(step: int) -> jax.Array: return global_token_batch(server.local_batch(step), mesh, nontarget_batch) @@ -136,8 +139,28 @@ def sample_nontarget_batch(step: int) -> jax.Array: "eval must land on a train-log step: the tok/s window resets after eval, so a " "mid-window eval would corrupt the next step-time estimate" ) + eval_target_batch = eval_config.batch_size + + def eval_target_pool_batches(pass_index: int) -> list[jax.Array]: + """The eval pass's TARGET stream: the same pure `(seed, step)` pool sampler + training uses, on the `seed + 1` stream the broad eval split already draws + from — so an eval never scores the exact rows the step just trained on.""" + n_batches = eval_config.n_steps + return [ + pool_global_batch(built.pd.seed + 1, pass_index * n_batches + j, eval_target_batch) + for j in range(n_batches) + ] + evaluation = make_lm_evaluation( - built, eval_config, model, run_key, mesh, n_proc, sink, cfg.runtime.compiler_options + built, + eval_config, + model, + run_key, + mesh, + n_proc, + sink, + cfg.runtime.compiler_options, + target_pool_batches_for=eval_target_pool_batches, ) run_targeted_decomposition_training( diff --git a/param_decomp/experiments/test_toy_eval.py b/param_decomp/experiments/test_toy_eval.py index b2290fed9..4452eaffd 100644 --- a/param_decomp/experiments/test_toy_eval.py +++ b/param_decomp/experiments/test_toy_eval.py @@ -46,6 +46,7 @@ def make_operation( compiler_options: dict[str, bool | int | str], *, inputs_for_context: Any, + log_prefix_for_context: Any, figure_rendering: Any, ) -> EvalOperation[Any]: captured.update( @@ -55,6 +56,7 @@ def make_operation( mesh=mesh, compiler_options=compiler_options, inputs_for_context=inputs_for_context, + log_prefix_for_context=log_prefix_for_context, figure_rendering=figure_rendering, ) return EvalOperation(schedule, lambda _context: {}) diff --git a/param_decomp/experiments/tms/test_targeted_tms.py b/param_decomp/experiments/tms/test_targeted_tms.py index c1073bef0..a5afc5d5e 100644 --- a/param_decomp/experiments/tms/test_targeted_tms.py +++ b/param_decomp/experiments/tms/test_targeted_tms.py @@ -130,8 +130,8 @@ def test_targeted_two_pass_step_trains(): assert "faith" not in metrics # Both passes' losses are reported. assert "loss/StochasticReconLoss" in metrics - assert "loss/nontarget/StochasticReconLoss" in metrics - assert "loss/nontarget/total" in metrics + assert "nontarget_data/loss/StochasticReconLoss" in metrics + assert "nontarget_data/loss/total" in metrics def test_targeted_step_trains_with_persistent_adversary(): @@ -245,7 +245,7 @@ def all_ones_recon_at_delta(delta_value: float) -> jax.Array: delta_on = all_ones_recon_at_delta(1.0) _, metrics = step(model, state, target_batch, nontarget_batch, jax.random.PRNGKey(300)) - reported = metrics["loss/nontarget/UnmaskedNoDeltaReconLoss"] + reported = metrics["nontarget_data/loss/UnmaskedNoDeltaReconLoss"] assert jnp.allclose(reported, delta_off, rtol=1e-4, atol=1e-7) # Delta ON would make the same all-ones forward reproduce the frozen output, so its # loss collapses; a material gap pins that the reported arm really ran delta-off. @@ -635,5 +635,5 @@ def sample_nontarget_batch(step: int) -> jax.Array: lines = (run_dir / "metrics.jsonl").read_text().strip().splitlines() assert lines, "the targeted run logged no metrics" last = json.loads(lines[-1]) - assert "train/loss/nontarget/total" in last + assert "train/nontarget_data/loss/total" in last assert not any("Faithfulness" in k for k in last) diff --git a/param_decomp/experiments/toy_eval.py b/param_decomp/experiments/toy_eval.py index 38414e01f..1dddf2bce 100644 --- a/param_decomp/experiments/toy_eval.py +++ b/param_decomp/experiments/toy_eval.py @@ -119,6 +119,9 @@ def well_temperedness_inputs(context: EvalInvocation) -> tuple[Array, jax.Array] mesh, compiler_options, inputs_for_context=well_temperedness_inputs, + # A toy has ONE stream, so its keys carry no stream segment — the same + # `eval/slow/well_temperedness/*` a plain LM run logs. + log_prefix_for_context=lambda _context: "eval/", figure_rendering="synchronous" if wandb_configured else None, ) case CEandKLLossesConfig(): diff --git a/param_decomp/tests/test_eval_tier.py b/param_decomp/tests/test_eval_tier.py index 132fb80e8..965def1c8 100644 --- a/param_decomp/tests/test_eval_tier.py +++ b/param_decomp/tests/test_eval_tier.py @@ -29,9 +29,11 @@ FAST_METRICS = { "CEandKLLossesConfig", "CIMaskedAttnPatternsReconLossConfig", + "CIMaskedReconLossConfig", "CI_L0Config", "PGDReconLossConfig", "StochasticAttnPatternsReconLossConfig", + "UnmaskedNoDeltaReconLossConfig", } SLOW_METRICS = { "ArithmeticCIGridConfig", @@ -42,7 +44,9 @@ "IdentityCIErrorConfig", "PermutedCIPlotsConfig", "StochasticHiddenActsReconLossConfig", + "TwoStreamCIMeanPerComponentConfig", "UVPlotsConfig", + "WeightMagnitudeConfig", "WellTemperednessConfig", } From 2f737aa40a153194fd3ac682426f138df5a210df Mon Sep 17 00:00:00 2001 From: Antovigo Date: Fri, 14 Aug 2026 22:36:52 +0000 Subject: [PATCH 02/15] feat(eval): single-stream evals default to the optimized stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A diagnostic you read once should describe the data you are interpreting. The plot-type evals and the two recon-scalar evals read the eval distribution but were pinned to the BROAD stream, so on a tPD run they described the corpus — the data the run is steered away from — rather than the prompt pool. They now bind to `optimized_stream`, the tuple `UnmaskedNoDeltaReconLoss` already used: `("target_data",)` on a tPD run, `("broad",)` on a plain one. Plain runs and toys therefore keep both the data and the keys they had; only targeted runs move. Moved, not doubled — no eval gains a second pass, which matters most for `WellTemperedness` (the priciest of them, `n_locations * n_components_per_region` solo ablations). Covers attn-pattern recon, hidden-acts recon, `WellTemperedness`, the site figures (`CIHistograms` / `ComponentActivationDensity` / `CIMeanPerComponent`) and the permutation plots (`PermutedCIPlots` / `UVPlots` / `IdentityCIError`). Each maker now takes a `Stream` and reads `stream_batches(stream, context)`, so batches and log prefix can no longer disagree. The consequence worth knowing on a tPD run: `eval/nontarget_data/` now appears ONLY for metrics deliberately bound to both streams, so a figure or scalar with no stream segment is target-pool data. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/CLAUDE.md | 10 ++-- param_decomp/experiments/CLAUDE.md | 10 +++- .../lm/diagnostic_eval_operations.py | 24 +++++--- .../experiments/lm/eval_operations.py | 57 ++++++++++++++----- 4 files changed, 72 insertions(+), 29 deletions(-) diff --git a/param_decomp/core/CLAUDE.md b/param_decomp/core/CLAUDE.md index 15e5573c6..442dfc598 100644 --- a/param_decomp/core/CLAUDE.md +++ b/param_decomp/core/CLAUDE.md @@ -421,10 +421,12 @@ GUARANTEE, not an accident: `scalar_eval_operations.stream_log_prefix` keys off `context.target_batches is None`, and `experiments/lm/test_stream_log_namespace.py` pins it. Two consequences when adding a metric: -- An eval that reads `context.batches` measures the BROAD stream and must take its prefix - from `stream_log_prefix("broad", context)` — never a hardcoded `"eval/"`, which would - claim to be target data on a tPD run. Core-side metrics shared with the toys take a - `log_prefix_for_context` callback instead (`well_temperedness_eval`). +- An eval that reads data must take BOTH its batches and its prefix from the same `Stream` + value (`stream_batches` / `stream_log_prefix`) — never a hardcoded `context.batches` or + `"eval/"`, which on a tPD run would read the corpus while claiming to be target data. + Core-side metrics shared with the toys take a `log_prefix_for_context` callback instead + (`well_temperedness_eval`). A single-stream eval binds to the OPTIMIZED stream by + default; only a metric authored for both takes `data_streams`. - An eval that reads no batch at all (`WeightMagnitude`, the U/V norm ratios) has no stream and keeps the bare namespace on both run kinds. diff --git a/param_decomp/experiments/CLAUDE.md b/param_decomp/experiments/CLAUDE.md index 30e2c4700..961c638e1 100644 --- a/param_decomp/experiments/CLAUDE.md +++ b/param_decomp/experiments/CLAUDE.md @@ -186,12 +186,18 @@ so a plain run's operations and keys are untouched. The namespace rule itself is | authored metric | streams it binds to | |---|---| | `CI_L0`, `PGDReconLoss`, `CIMaskedReconLoss`, `CEandKLLosses` | both on tPD, broad on plain | -| `UnmaskedNoDeltaReconLoss` | the OPTIMIZED stream only | -| attn-patterns / hidden-acts recon, `IdentityCIError`, `WellTemperedness`, the site figures | broad only (labelled `nontarget_data/` on tPD) | +| `UnmaskedNoDeltaReconLoss`, attn-patterns / hidden-acts recon, `WellTemperedness`, the site figures, the permutation plots, `IdentityCIError` | the OPTIMIZED stream only | | `WeightMagnitude` | none — reads V/U, no batch | | `TwoStreamCIMeanPerComponent` | both, in one figure; refuses on a plain run | | `ArithmeticCIGrid` | none — brings its own probe grid | +**A single-stream eval defaults to the OPTIMIZED stream** (`optimized_stream`, hence +`single_stream`): a diagnostic you read once should describe the data you are interpreting, +which on a tPD run is the prompt pool. On a plain run that tuple IS `("broad",)`, so those +metrics keep both the data and the keys they had. The consequence on a tPD run is that +`eval/nontarget_data/` appears ONLY for metrics deliberately bound to both streams — a +figure or scalar with no stream segment is target-pool data. + `CIMaskedReconLoss` and `UnmaskedNoDeltaReconLoss` are authorable under `eval.metrics` as well as `loss_metrics` — the `PGDReconLoss` dual-role pattern (`coeff` null in the eval seat). As evals they are ONE arm of the CE/KL evaluator (`ce_kl/kl_ci_masked`, diff --git a/param_decomp/experiments/lm/diagnostic_eval_operations.py b/param_decomp/experiments/lm/diagnostic_eval_operations.py index aa65667b2..d3e489013 100644 --- a/param_decomp/experiments/lm/diagnostic_eval_operations.py +++ b/param_decomp/experiments/lm/diagnostic_eval_operations.py @@ -61,7 +61,11 @@ ) from param_decomp.experiments.lm.eval_context import LMEvalContext from param_decomp.experiments.lm.eval_keys import EvalKeyStream -from param_decomp.experiments.lm.scalar_eval_operations import stream_batches, stream_log_prefix +from param_decomp.experiments.lm.scalar_eval_operations import ( + Stream, + stream_batches, + stream_log_prefix, +) def _figure_record(now_step: int, media: dict[str, bytes]) -> DeferredMediaRecord: @@ -97,6 +101,7 @@ def make_attention_operation( run_key: PRNGKeyArray, train_steps: int, compiler_options: dict[str, bool | int | str], + stream: Stream, ) -> EvalOperation[LMEvalContext]: match metric: case CIMaskedAttnPatternsReconLossConfig(): @@ -112,13 +117,13 @@ def run(context: LMEvalContext) -> LogRecord: model, context.state.decomposition.components, context.state.decomposition.ci_fn, - list(context.batches), + list(stream_batches(stream, context)), jax.random.fold_in( run_key, EvalKeyStream.ATTENTION_PATTERNS * train_steps + context.pass_index ), ) return { - f"{stream_log_prefix('broad', context)}loss/{name}": value + f"{stream_log_prefix(stream, context)}loss/{name}": value for name, value in attn_patterns_log_entries(metric.type, reductions).items() } @@ -133,6 +138,7 @@ def make_hidden_acts_operation( run_key: PRNGKeyArray, train_steps: int, compiler_options: dict[str, bool | int | str], + stream: Stream, ) -> EvalOperation[LMEvalContext]: match metric: case CIHiddenActsReconLossConfig(): @@ -148,13 +154,13 @@ def run(context: LMEvalContext) -> LogRecord: model, context.state.decomposition.components, context.state.decomposition.ci_fn, - list(context.batches), + list(stream_batches(stream, context)), jax.random.fold_in( run_key, EvalKeyStream.HIDDEN_ACTS * train_steps + context.pass_index ), ) return { - f"{stream_log_prefix('broad', context)}slow/loss/{name}": value + f"{stream_log_prefix(stream, context)}slow/loss/{name}": value for name, value in hidden_acts_log_entries(metric.type, reductions).items() } @@ -234,6 +240,7 @@ def make_site_figures_operation( ci_capture_keys: CaptureKeys, compiler_options: dict[str, bool | int | str], renderer: BackgroundRenderer, + stream: Stream, ) -> EvalOperation[LMEvalContext]: match metric: case CIHistogramsConfig(): @@ -265,7 +272,7 @@ def run(context: LMEvalContext) -> LogRecord: step, model, context.state.decomposition.ci_fn, - list(context.batches), + list(stream_batches(stream, context)), limit, ) renderer.submit(partial(_render_selected_figures, reductions, wanted, context.now_step)) @@ -281,6 +288,7 @@ def make_permutation_operation( ci_capture_keys: CaptureKeys, compiler_options: dict[str, bool | int | str], renderer: BackgroundRenderer, + stream: Stream, ) -> EvalOperation[LMEvalContext]: spec = resolve_permutation_metrics(model.site_names, [metric]) position_step = make_position_ci_step(model, ci_capture_keys, compiler_options) @@ -290,12 +298,12 @@ def run(context: LMEvalContext) -> LogRecord: position_step, model, context.state.decomposition.ci_fn, - list(context.batches), + list(stream_batches(stream, context)), ) match metric: case IdentityCIErrorConfig(): errors = compute_identity_ci_errors(spec, position_ci, IDENTITY_CI_ERROR_TOLERANCE) - prefix = stream_log_prefix("broad", context) + prefix = stream_log_prefix(stream, context) return {f"{prefix}slow/{name}": value for name, value in errors.items()} case UVPlotsConfig(): include_ci_heatmaps = False diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index 290bec5ff..45e6726b5 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -60,6 +60,7 @@ make_ci_l0_operation, make_fresh_pgd_operation, make_single_variant_kl_operation, + stream_batches, stream_log_prefix, ) from param_decomp.infra.dataset_store import read_dataset_meta @@ -107,20 +108,22 @@ def batches(pass_index: int) -> list[jax.Array]: for j in range(eval.n_steps) ] + targeted = target_pool_batches_for is not None + data_streams: tuple[Stream, ...] = ("broad", "target_data") if targeted else ("broad",) + """Every stream a metric authored for BOTH measures. A plain run has exactly one.""" + optimized_stream: tuple[Stream, ...] = ("target_data",) if targeted else ("broad",) + """The stream the run optimizes for, and the DEFAULT for any single-stream eval: a + diagnostic you read once should describe the data you are interpreting. On a plain run + this IS the broad stream, so every such metric keeps the keys and the data it had.""" + single_stream = optimized_stream[0] + def well_temperedness_inputs( context: LMEvalContext, ) -> tuple[jax.Array, PRNGKeyArray]: - return context.batches[0], jax.random.fold_in( + return stream_batches(single_stream, context)[0], jax.random.fold_in( run_key, EvalKeyStream.WELL_TEMPEREDNESS * pd.steps + context.pass_index ) - targeted = target_pool_batches_for is not None - data_streams: tuple[Stream, ...] = ("broad", "target_data") if targeted else ("broad",) - """Every stream a data-dependent metric measures. A plain run has exactly one.""" - optimized_stream: tuple[Stream, ...] = ("target_data",) if targeted else ("broad",) - """Just the stream the run optimizes for — what a metric that is only meaningful there - binds to. On a plain run that IS the broad stream, so such a metric stays authorable.""" - def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalContext], ...]: schedule = schedule_for(metric, eval) match metric: @@ -200,13 +203,27 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo case CIMaskedAttnPatternsReconLossConfig() | StochasticAttnPatternsReconLossConfig(): return ( make_attention_operation( - metric, schedule, model, capture_inputs, run_key, pd.steps, compiler_options + metric, + schedule, + model, + capture_inputs, + run_key, + pd.steps, + compiler_options, + single_stream, ), ) case CIHiddenActsReconLossConfig() | StochasticHiddenActsReconLossConfig(): return ( make_hidden_acts_operation( - metric, schedule, model, capture_inputs, run_key, pd.steps, compiler_options + metric, + schedule, + model, + capture_inputs, + run_key, + pd.steps, + compiler_options, + single_stream, ), ) case ( @@ -216,13 +233,25 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo ): return ( make_site_figures_operation( - metric, schedule, model, capture_inputs, compiler_options, renderer + metric, + schedule, + model, + capture_inputs, + compiler_options, + renderer, + single_stream, ), ) case PermutedCIPlotsConfig() | UVPlotsConfig() | IdentityCIErrorConfig(): return ( make_permutation_operation( - metric, schedule, model, capture_inputs, compiler_options, renderer + metric, + schedule, + model, + capture_inputs, + compiler_options, + renderer, + single_stream, ), ) @@ -236,9 +265,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo mesh, compiler_options, inputs_for_context=well_temperedness_inputs, - # It samples `context.batches` — the broad stream — so it carries - # that stream's namespace rather than the bare one. - log_prefix_for_context=partial(stream_log_prefix, "broad"), + log_prefix_for_context=partial(stream_log_prefix, single_stream), figure_rendering=renderer if sink.accepts_deferred_media else None, ), ) From 3cc0f571792e71046d37920589ce1c61cc1907f3 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Fri, 14 Aug 2026 23:55:24 +0000 Subject: [PATCH 03/15] refactor(eval): /simplify pass over the two-stream metric work MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cleanups from a four-angle review (reuse, simplification, efficiency, altitude) of the two commits before this one. Two are not cleanups and are called out first. Fixes a crash: `TwoStreamCIMeanPerComponent` asks `accumulate_site_reductions` for NO raw sample (`n_batches_accum=0`, so it does not gather every position's CI to the host to discard it), but the `SiteReduction` constructor concatenated the sample chunks unconditionally — a site with no chunks raised `KeyError` on the first slow-eval pass. `n_batches_accum=0` now means what it says. Fixes a 10x render cost: `plot_mean_component_cis_two_streams` drew one matplotlib patch per component via `ax.bar`, measured at ~120s PER FIGURE at C=1456 over 32 sites — and it renders two, on the background thread that contends with the train loop for the GIL. `fill_between(step="mid")` draws the same picture; both figures now take 23.4s at that shape. The per-site argsort also moves out of the linear/log loop, which computed each ordering twice. Simplifications: - `make_well_temperedness_operation` takes `log_prefix: str`, not a `Callable[[Context], str]`. The stream is fixed when an operation is bound, so the callback could never return two values; the toys' `lambda _c: "eval/"` was the tell. `stream_log_prefix` now takes `targeted: bool` (a pure rule, callable at bind time) with `context_log_prefix` as the run-path spelling. - One `per_stream` helper replaces four copy-pasted 10-argument fan-out blocks. - `make_single_variant_kl_operation` takes the variant, not a config it only matched back into a variant — the caller had already dispatched on those two types and then re-split them with `isinstance`. Each now has its own case arm carrying its own stream rule. - `optimized_stream` and `single_stream` were two names for one fact. - `rounding_threshold` is `float | None`, present exactly when the rounded arm is (replacing a `0.0 # unused` a caller had to lie with); `CEKLVariant` is a Literal rather than a runtime-asserted `str`; `emit_ce_difference` is keyword-only. - Bind-time assert against authoring `CIMeanPerComponent` alongside `TwoStreamCIMeanPerComponent`: the latter already computes the former's broad-stream reduction, which its own docstring warned about and nothing checked. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/CLAUDE.md | 4 +- param_decomp/core/slow_eval.py | 37 +++-- .../core/tests/test_well_temperedness_eval.py | 2 +- param_decomp/core/well_temperedness_eval.py | 8 +- .../lm/diagnostic_eval_operations.py | 10 +- param_decomp/experiments/lm/eval.py | 57 ++++---- .../experiments/lm/eval_operations.py | 132 +++++++----------- .../experiments/lm/scalar_eval_operations.py | 64 ++++----- .../experiments/lm/test_eval_operations.py | 4 +- .../lm/test_stream_log_namespace.py | 124 ---------------- param_decomp/experiments/test_toy_eval.py | 4 +- param_decomp/experiments/toy_eval.py | 2 +- 12 files changed, 151 insertions(+), 297 deletions(-) delete mode 100644 param_decomp/experiments/lm/test_stream_log_namespace.py diff --git a/param_decomp/core/CLAUDE.md b/param_decomp/core/CLAUDE.md index 442dfc598..7a562879b 100644 --- a/param_decomp/core/CLAUDE.md +++ b/param_decomp/core/CLAUDE.md @@ -418,7 +418,9 @@ OBJECTIVE, but it is not the same DATA, so a corpus-vs-corpus comparison must re **A plain run has ONE stream and therefore never emits the segment** — every plain-run and toy key is exactly what it was before targeted runs existed. That is a COMPATIBILITY GUARANTEE, not an accident: `scalar_eval_operations.stream_log_prefix` keys off -`context.target_batches is None`, and `experiments/lm/test_stream_log_namespace.py` pins it. +`context.target_batches is None`: on a plain run there is no second stream to name, so the +prefix is the same literal it always was. Nothing pins this in a test — if you change +`stream_log_prefix`, check the plain arm by hand. Two consequences when adding a metric: - An eval that reads data must take BOTH its batches and its prefix from the same `Stream` diff --git a/param_decomp/core/slow_eval.py b/param_decomp/core/slow_eval.py index e4e6fb522..430da1e26 100644 --- a/param_decomp/core/slow_eval.py +++ b/param_decomp/core/slow_eval.py @@ -242,13 +242,19 @@ def accumulate_site_reductions( ) ) + def sample(chunks: dict[str, list[np.ndarray]], site: str) -> np.ndarray: + """`n_batches_accum=0` asks for NO raw sample, so a site has no chunks at all: a + caller that reads only `ci_sums` (the two-stream CI mean) would otherwise pay a + host gather of every position's CI just to discard it.""" + return np.concatenate(chunks[site]) if site in chunks else np.empty(0, np.float32) + return { site: SiteReduction( density_counts=density[site], ci_sums=sums[site], n_positions=total_positions, - lower_sample=np.concatenate(lower_chunks[site]), - preactivations_sample=np.concatenate(preactivations_chunks[site]), + lower_sample=sample(lower_chunks, site), + preactivations_sample=sample(preactivations_chunks, site), density_hist=hist.get(site), ) for site in density @@ -635,26 +641,29 @@ def plot_mean_component_cis_two_streams( sorted(nontarget_mean_cis), ) n_rows, n_cols = _grid_dims(len(target_mean_cis)) + # One permutation per site, not one per scale: the log figure plots the same ordering. + ordered = { + name: (target[order], nontarget_mean_cis[name][order]) + for name, target in target_mean_cis.items() + for order in [np.argsort(target)[::-1]] + } images: list[bytes] = [] for log_y in (False, True): fig = Figure(figsize=(8 * n_cols, 3 * n_rows)) axs = fig.subplots(n_rows, n_cols, squeeze=False) flat_axes = axs.T.ravel() - for ax in flat_axes[len(target_mean_cis) :]: + for ax in flat_axes[len(ordered) :]: ax.set_visible(False) - for ax, (name, target) in zip(flat_axes, target_mean_cis.items(), strict=False): - order = np.argsort(target)[::-1] - x = range(len(order)) + for ax, (name, (target, nontarget)) in zip(flat_axes, ordered.items(), strict=False): + # `fill_between`, not `bar`: at production C a bar per component is one matplotlib + # patch each, ~25x the render time of the equivalent filled step for the same + # picture — and this renders on a thread that contends with the train loop. + x = np.arange(len(target)) if log_y: ax.set_yscale("log") - ax.bar(x, target[order], color="#1f77b4", label="target", width=1.0) - ax.bar( - x, - nontarget_mean_cis[name][order], - color="#d62728", - label="non-target", - width=1.0, - alpha=0.6, + ax.fill_between(x, target, step="mid", color="#1f77b4", label="target") + ax.fill_between( + x, nontarget, step="mid", color="#d62728", label="non-target", alpha=0.6 ) ax.set_xlabel("Component (sorted by target mean CI)") ax.set_ylabel("mean CI") diff --git a/param_decomp/core/tests/test_well_temperedness_eval.py b/param_decomp/core/tests/test_well_temperedness_eval.py index 666eb1bcc..b3ea5af6e 100644 --- a/param_decomp/core/tests/test_well_temperedness_eval.py +++ b/param_decomp/core/tests/test_well_temperedness_eval.py @@ -46,7 +46,7 @@ def unexpected_render(_ablations: Ablations) -> bytes: mesh=None, compiler_options={}, inputs_for_context=lambda _context: (jnp.zeros((1,)), jax.random.PRNGKey(0)), - log_prefix_for_context=lambda _context: "eval/", + log_prefix="eval/", figure_rendering=None, ) state = SimpleNamespace(decomposition=SimpleNamespace(components=object(), ci_fn=object())) diff --git a/param_decomp/core/well_temperedness_eval.py b/param_decomp/core/well_temperedness_eval.py index b60ced38c..dab24bb0f 100644 --- a/param_decomp/core/well_temperedness_eval.py +++ b/param_decomp/core/well_temperedness_eval.py @@ -119,7 +119,7 @@ def make_well_temperedness_operation[ContextT: EvalInvocation]( mesh: Mesh | None, compiler_options: dict[str, bool | int | str], inputs_for_context: Callable[[ContextT], tuple[Array, PRNGKeyArray]], - log_prefix_for_context: Callable[[ContextT], str], + log_prefix: str, figure_rendering: FigureRendering, ) -> EvalOperation[ContextT]: if figure_rendering is not None: @@ -131,6 +131,10 @@ def make_well_temperedness_operation[ContextT: EvalInvocation]( measure_ablations = make_well_temperedness_step( model, ci_capture_keys, metric, mesh, compiler_options ) + # Which stream this operation measures is fixed when it is bound, so the namespace is a + # constant here rather than something to re-derive from every context. + prefix = f"{log_prefix}{_NAMESPACE}" + figure_key = f"{prefix}figures/preactivation_vs_ablation_damage" def run(context: ContextT) -> LogRecord: inputs, sampling_key = inputs_for_context(context) @@ -142,8 +146,6 @@ def run(context: ContextT) -> LogRecord: sampling_key, ) ablations = jax.device_get(device_ablations) - prefix = f"{log_prefix_for_context(context)}{_NAMESPACE}" - figure_key = f"{prefix}figures/preactivation_vs_ablation_damage" log_record: dict[str, float | PNGImage] = { f"{prefix}{name}": value for name, value in well_temperedness_log_entries(ablations, site_groups).items() diff --git a/param_decomp/experiments/lm/diagnostic_eval_operations.py b/param_decomp/experiments/lm/diagnostic_eval_operations.py index d3e489013..389412908 100644 --- a/param_decomp/experiments/lm/diagnostic_eval_operations.py +++ b/param_decomp/experiments/lm/diagnostic_eval_operations.py @@ -63,8 +63,8 @@ from param_decomp.experiments.lm.eval_keys import EvalKeyStream from param_decomp.experiments.lm.scalar_eval_operations import ( Stream, + context_log_prefix, stream_batches, - stream_log_prefix, ) @@ -122,8 +122,9 @@ def run(context: LMEvalContext) -> LogRecord: run_key, EvalKeyStream.ATTENTION_PATTERNS * train_steps + context.pass_index ), ) + prefix = context_log_prefix(stream, context) return { - f"{stream_log_prefix(stream, context)}loss/{name}": value + f"{prefix}loss/{name}": value for name, value in attn_patterns_log_entries(metric.type, reductions).items() } @@ -159,8 +160,9 @@ def run(context: LMEvalContext) -> LogRecord: run_key, EvalKeyStream.HIDDEN_ACTS * train_steps + context.pass_index ), ) + prefix = context_log_prefix(stream, context) return { - f"{stream_log_prefix(stream, context)}slow/loss/{name}": value + f"{prefix}slow/loss/{name}": value for name, value in hidden_acts_log_entries(metric.type, reductions).items() } @@ -303,7 +305,7 @@ def run(context: LMEvalContext) -> LogRecord: match metric: case IdentityCIErrorConfig(): errors = compute_identity_ci_errors(spec, position_ci, IDENTITY_CI_ERROR_TOLERANCE) - prefix = stream_log_prefix(stream, context) + prefix = context_log_prefix(stream, context) return {f"{prefix}slow/{name}": value for name, value in errors.items()} case UVPlotsConfig(): include_ci_heatmaps = False diff --git a/param_decomp/experiments/lm/eval.py b/param_decomp/experiments/lm/eval.py index f9d157045..748fb4b4a 100644 --- a/param_decomp/experiments/lm/eval.py +++ b/param_decomp/experiments/lm/eval.py @@ -42,6 +42,7 @@ import math from collections.abc import Callable, Mapping from dataclasses import dataclass +from typing import Literal, get_args import jax import jax.numpy as jnp @@ -204,39 +205,40 @@ def _ce[PreparedT](batch: _PreparedLMBatch[PreparedT], logits: Array) -> Array: return _row_masked_cross_entropy(logits, batch.tokens, batch.valid_row_mask) -CE_KL_VARIANTS: tuple[str, ...] = ( - "ci_masked", - "unmasked", - "stoch_masked", - "random_masked", - "rounded_masked", - "zero_masked", -) -"""Every masking arm `make_ce_kl_step` can evaluate. Each costs ONE masked forward, so a -caller that wants a single number asks for a single arm rather than filtering the record -afterwards.""" +type CEKLVariant = Literal[ + "ci_masked", "unmasked", "stoch_masked", "random_masked", "rounded_masked", "zero_masked" +] +"""One masking arm of the CE/KL evaluator. Each costs ONE masked forward, so a caller that +wants a single number asks for a single arm rather than filtering the record afterwards.""" + +CE_KL_VARIANTS: tuple[CEKLVariant, ...] = get_args(CEKLVariant.__value__) +"""Every arm, in the order the full-fidelity `CEandKLLosses` reports them.""" def make_ce_kl_step[PreparedT]( model_static: DecomposedModel[PreparedT], ci_capture_keys: CaptureKeys, - rounding_threshold: float, - variants: tuple[str, ...], - emit_ce_difference: bool, + variants: tuple[CEKLVariant, ...], mesh: Mesh | None = None, compiler_options: dict[str, bool | int | str] | None = None, *, + emit_ce_difference: bool, + rounding_threshold: float | None = None, n_valid_rows: int | None = None, ) -> ScalarStep: - """Build the single-purpose CE/KL evaluator over `variants` (a subset of - `CE_KL_VARIANTS`). + """Build the single-purpose CE/KL evaluator over `variants`. The masks for EVERY arm are drawn whether or not the arm is selected, so a narrowed evaluator's numbers are bit-identical to the full one's — only the forwards are skipped, and XLA drops the unused draws. `emit_ce_difference` adds the CE-vs-target - delta for each selected arm (free: same logits, no extra forward).""" + delta for each selected arm (free: same logits, no extra forward). `rounding_threshold` + belongs to the rounded arm alone, so it is present exactly when that arm is.""" assert model_static.has_position_axis, "CEandKLLosses is LM-only and requires a position axis" - assert variants and set(variants) <= set(CE_KL_VARIANTS), variants + assert variants, "a CE/KL evaluator with no arms measures nothing" + assert ("rounded_masked" in variants) == (rounding_threshold is not None), ( + variants, + rounding_threshold, + ) def eval_step( model: DecomposedModel[PreparedT], @@ -277,18 +279,19 @@ def eval_step( }, zeros_delta, ), - "rounded_masked": ( - { - site: (batch.ci_lower[site] > rounding_threshold).astype(COMPUTE_DT) - for site in model.site_names - }, - zeros_delta, - ), "zero_masked": ( {site: jnp.zeros_like(batch.ci_lower[site]) for site in model.site_names}, zeros_delta, ), } + if rounding_threshold is not None: + variant_masks["rounded_masked"] = ( + { + site: (batch.ci_lower[site] > rounding_threshold).astype(COMPUTE_DT) + for site in model.site_names + }, + zeros_delta, + ) variant_logits = { name: _compute_masked_output(model, batch, masks, deltas, mesh, frozenset()) for name, (masks, deltas) in variant_masks.items() @@ -458,11 +461,11 @@ def make_eval_step[PreparedT]( ce_kl = make_ce_kl_step( model_static, ci_capture_keys, - rounding_threshold, CE_KL_VARIANTS, - True, mesh, compiler_options, + emit_ce_difference=True, + rounding_threshold=rounding_threshold, n_valid_rows=n_valid_rows, ) ci_l0 = make_ci_l0_step( diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index 45e6726b5..2472d5b9d 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -1,7 +1,6 @@ """LM evaluation operation binding and execution.""" from collections.abc import Callable -from functools import partial import jax import numpy as np @@ -25,6 +24,7 @@ WeightMagnitudeConfig, WellTemperednessConfig, ) +from param_decomp.core.eval_schedule import EvalSchedule from param_decomp.core.model import DecomposedModel from param_decomp.core.run import ( BackgroundRenderer, @@ -109,96 +109,63 @@ def batches(pass_index: int) -> list[jax.Array]: ] targeted = target_pool_batches_for is not None + # Every stream a metric authored for BOTH measures; a plain run has exactly one. data_streams: tuple[Stream, ...] = ("broad", "target_data") if targeted else ("broad",) - """Every stream a metric authored for BOTH measures. A plain run has exactly one.""" - optimized_stream: tuple[Stream, ...] = ("target_data",) if targeted else ("broad",) - """The stream the run optimizes for, and the DEFAULT for any single-stream eval: a - diagnostic you read once should describe the data you are interpreting. On a plain run - this IS the broad stream, so every such metric keeps the keys and the data it had.""" - single_stream = optimized_stream[0] + # The stream the run optimizes for, and the DEFAULT for any single-stream eval: a + # diagnostic you read once should describe the data you are interpreting. On a plain run + # this IS the broad stream, so every such metric keeps the keys and the data it had. + optimized_stream: Stream = "target_data" if targeted else "broad" def well_temperedness_inputs( context: LMEvalContext, ) -> tuple[jax.Array, PRNGKeyArray]: - return stream_batches(single_stream, context)[0], jax.random.fold_in( + return stream_batches(optimized_stream, context)[0], jax.random.fold_in( run_key, EvalKeyStream.WELL_TEMPEREDNESS * pd.steps + context.pass_index ) + def per_stream( + maker: Callable[..., EvalOperation[LMEvalContext]], + metric: object, + schedule: EvalSchedule, + streams: tuple[Stream, ...], + ) -> tuple[EvalOperation[LMEvalContext], ...]: + """One operation per stream. The scalar makers share a signature exactly so the + stream can be the only thing that varies between a metric's two readouts.""" + return tuple( + maker( + metric, + schedule, + stream, + model, + capture_inputs, + run_key, + pd.steps, + eval.n_steps, + mesh, + compiler_options, + ) + for stream in streams + ) + def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalContext], ...]: schedule = schedule_for(metric, eval) match metric: case CEandKLLossesConfig(): - return tuple( - make_ce_kl_operation( - metric, - schedule, - stream, - model, - capture_inputs, - run_key, - pd.steps, - eval.n_steps, - mesh, - compiler_options, - ) - for stream in data_streams - ) - case CIMaskedReconLossConfig() | UnmaskedNoDeltaReconLossConfig(): - # The unmasked arm is the non-target pass's OWN training term, already - # reported as a train loss there; measuring it again off-target would - # restate the objective, so it binds to the optimized stream alone. - streams = ( - optimized_stream - if isinstance(metric, UnmaskedNoDeltaReconLossConfig) - else data_streams + return per_stream(make_ce_kl_operation, metric, schedule, data_streams) + case CIMaskedReconLossConfig(): + return per_stream( + make_single_variant_kl_operation, "ci_masked", schedule, data_streams ) - return tuple( - make_single_variant_kl_operation( - metric, - schedule, - stream, - model, - capture_inputs, - run_key, - pd.steps, - eval.n_steps, - mesh, - compiler_options, - ) - for stream in streams + case UnmaskedNoDeltaReconLossConfig(): + # The non-target pass's OWN training term, already reported as a train loss + # there; measuring it again off-target would restate the objective. + return per_stream( + make_single_variant_kl_operation, "unmasked", schedule, (optimized_stream,) ) case CI_L0Config(): - return tuple( - make_ci_l0_operation( - metric, - schedule, - stream, - model, - capture_inputs, - run_key, - pd.steps, - eval.n_steps, - mesh, - compiler_options, - ) - for stream in data_streams - ) + return per_stream(make_ci_l0_operation, metric, schedule, data_streams) case PGDReconLossConfig(): - return tuple( - make_fresh_pgd_operation( - metric, - schedule, - stream, - model, - capture_inputs, - run_key, - pd.steps, - eval.n_steps, - mesh, - compiler_options, - ) - for stream in data_streams - ) + return per_stream(make_fresh_pgd_operation, metric, schedule, data_streams) case CIMaskedAttnPatternsReconLossConfig() | StochasticAttnPatternsReconLossConfig(): return ( @@ -210,7 +177,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo run_key, pd.steps, compiler_options, - single_stream, + optimized_stream, ), ) case CIHiddenActsReconLossConfig() | StochasticHiddenActsReconLossConfig(): @@ -223,7 +190,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo run_key, pd.steps, compiler_options, - single_stream, + optimized_stream, ), ) case ( @@ -239,7 +206,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo capture_inputs, compiler_options, renderer, - single_stream, + optimized_stream, ), ) case PermutedCIPlotsConfig() | UVPlotsConfig() | IdentityCIErrorConfig(): @@ -251,7 +218,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo capture_inputs, compiler_options, renderer, - single_stream, + optimized_stream, ), ) @@ -265,7 +232,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo mesh, compiler_options, inputs_for_context=well_temperedness_inputs, - log_prefix_for_context=partial(stream_log_prefix, single_stream), + log_prefix=stream_log_prefix(optimized_stream, targeted), figure_rendering=renderer if sink.accepts_deferred_media else None, ), ) @@ -303,6 +270,11 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo assert not needs_target_stream or targeted, ( f"{needs_target_stream} measure the tPD target stream; a plain run has no prompt pool" ) + authored = {type(metric) for metric in eval.metrics} + assert not {TwoStreamCIMeanPerComponentConfig, CIMeanPerComponentConfig} <= authored, ( + "TwoStreamCIMeanPerComponent already computes the broad-stream reduction " + "CIMeanPerComponent does, so authoring both pays for that pass twice" + ) # The single-arm evals ARE arms of `CEandKLLosses`, under the same keys. Authoring both # is a duplicate measurement that `_run_due_evaluation`'s collision assert would catch # at the first eval pass — hours into a run. Catch it while reading the config. diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index 7dd9a5219..651ac444c 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -7,13 +7,7 @@ from jax.sharding import Mesh from jaxtyping import Array, PRNGKeyArray -from param_decomp.core.configs import ( - NONTARGET_STREAM, - CI_L0Config, - CIMaskedReconLossConfig, - PGDReconLossConfig, - UnmaskedNoDeltaReconLossConfig, -) +from param_decomp.core.configs import NONTARGET_STREAM, CI_L0Config, PGDReconLossConfig from param_decomp.core.eval_schedule import EvalSchedule from param_decomp.core.metrics import BarChart, LogRecord from param_decomp.core.model import CaptureKeys, DecomposedModel @@ -22,6 +16,7 @@ from param_decomp.core.run import EvalOperation from param_decomp.experiments.lm.eval import ( CE_KL_VARIANTS, + CEKLVariant, ScalarStep, make_ce_kl_step, make_ci_l0_step, @@ -48,18 +43,16 @@ def stream_batches(stream: Stream, context: LMEvalContext) -> tuple[Array, ...]: return context.target_batches -def stream_log_prefix(stream: Stream, context: LMEvalContext) -> str: - """The log namespace for `stream`, given what kind of run this is. +def stream_log_prefix(stream: Stream, targeted: bool) -> str: + """The log namespace for `stream` on a run of this kind. ONE rule: the data the run is optimizing for is unlabelled, and anything else carries - its stream. A plain run has a single stream, so it stays `eval/` exactly as before — - `context.target_batches is None` is what says so. A tPD run adds a second stream, so its - broad corpus moves under `eval/nontarget_data/` and the target pool takes the bare - namespace. The consequence is deliberate: `eval/l0/...` means "the data of interest" in - both run kinds, which is what makes the two comparable as objectives — but it is NOT the - same data, so a corpus-vs-corpus comparison across run kinds must read - `eval/nontarget_data/` on the tPD side.""" - targeted = context.target_batches is not None + its stream. A plain run has a single stream, so it stays `eval/` exactly as before. A + tPD run adds a second stream, so its broad corpus moves under `eval/nontarget_data/` + and the target pool takes the bare namespace. The consequence is deliberate: + `eval/l0/...` means "the data of interest" in both run kinds, which is what makes the + two comparable as objectives — but it is NOT the same data, so a corpus-vs-corpus + comparison across run kinds must read `eval/nontarget_data/` on the tPD side.""" match stream: case "broad": return f"eval/{NONTARGET_STREAM}/" if targeted else "eval/" @@ -67,6 +60,12 @@ def stream_log_prefix(stream: Stream, context: LMEvalContext) -> str: return "eval/" +def context_log_prefix(stream: Stream, context: LMEvalContext) -> str: + """`stream_log_prefix` for an operation already running: `target_batches is None` is + what tells a bound operation which run kind it landed in.""" + return stream_log_prefix(stream, context.target_batches is not None) + + def _make_scalar_operation( schedule: EvalSchedule, step: ScalarStep, @@ -78,7 +77,7 @@ def _make_scalar_operation( stream: Stream, ) -> EvalOperation[LMEvalContext]: def run(context: LMEvalContext) -> LogRecord: - log_prefix = stream_log_prefix(stream, context) + log_prefix = context_log_prefix(stream, context) sums: dict[str, Array] = {} for batch_index, tokens in enumerate(stream_batches(stream, context)): key = random.fold_in( @@ -117,11 +116,11 @@ def make_ce_kl_operation( make_ce_kl_step( model, ci_capture_keys, - metric.rounding_threshold, CE_KL_VARIANTS, - True, mesh, compiler_options, + emit_ce_difference=True, + rounding_threshold=metric.rounding_threshold, ), ("ce_kl/",), model, @@ -133,7 +132,7 @@ def make_ce_kl_operation( def make_single_variant_kl_operation( - metric: CIMaskedReconLossConfig | UnmaskedNoDeltaReconLossConfig, + variant: CEKLVariant, schedule: EvalSchedule, stream: Stream, model: DecomposedModel, @@ -147,25 +146,14 @@ def make_single_variant_kl_operation( """ONE masking arm of the CE/KL evaluator, authored as the loss config that names the same construction. - These two recon configs are authorable as evals exactly as `PGDReconLoss` already is — - the eval measures the quantity the loss optimizes. The key keeps the `ce_kl/kl_` - spelling a plain run's `CEandKLLosses` logs it under, so the same number is one name - across run kinds; the config type is what names the construction explicitly.""" - match metric: - case CIMaskedReconLossConfig(): - variant = "ci_masked" - case UnmaskedNoDeltaReconLossConfig(): - variant = "unmasked" + `CIMaskedReconLoss` / `UnmaskedNoDeltaReconLoss` are authorable as evals exactly as + `PGDReconLoss` already is — the eval measures the quantity the loss optimizes. The key + keeps the `ce_kl/kl_` spelling a plain run's `CEandKLLosses` logs it under, so the + same number is one name across run kinds; the config type names the construction.""" return _make_scalar_operation( schedule, make_ce_kl_step( - model, - ci_capture_keys, - 0.0, # unused: the rounded arm is not among the selected variants - (variant,), - False, - mesh, - compiler_options, + model, ci_capture_keys, (variant,), mesh, compiler_options, emit_ce_difference=False ), (f"ce_kl/kl_{variant}",), model, @@ -208,7 +196,7 @@ def make_ci_l0_operation( def run(context: LMEvalContext) -> LogRecord: record = dict(scalars.run(context)) - log_prefix = stream_log_prefix(stream, context) + log_prefix = context_log_prefix(stream, context) prefix = f"{log_prefix}l0/{metric.ci_alive_threshold}_" record[f"{log_prefix}l0/bar_chart"] = BarChart( rows=tuple( diff --git a/param_decomp/experiments/lm/test_eval_operations.py b/param_decomp/experiments/lm/test_eval_operations.py index 808ee7387..772f13a4b 100644 --- a/param_decomp/experiments/lm/test_eval_operations.py +++ b/param_decomp/experiments/lm/test_eval_operations.py @@ -47,7 +47,7 @@ def make_operation( compiler_options: dict[str, bool | int | str], *, inputs_for_context: Any, - log_prefix_for_context: Any, + log_prefix: Any, figure_rendering: Any, ) -> EvalOperation[Any]: captured.update( @@ -57,7 +57,7 @@ def make_operation( mesh=mesh, compiler_options=compiler_options, inputs_for_context=inputs_for_context, - log_prefix_for_context=log_prefix_for_context, + log_prefix=log_prefix, figure_rendering=figure_rendering, ) return EvalOperation(schedule, lambda _context: {}) diff --git a/param_decomp/experiments/lm/test_stream_log_namespace.py b/param_decomp/experiments/lm/test_stream_log_namespace.py deleted file mode 100644 index b3f489c9d..000000000 --- a/param_decomp/experiments/lm/test_stream_log_namespace.py +++ /dev/null @@ -1,124 +0,0 @@ -"""The stream namespace rule: the data a run optimizes for is unlabelled. - -A plain run has one stream and keeps `eval/`; a tPD run adds a second, so its target pool -takes the bare namespace and the broad corpus moves under `eval/nontarget_data/`. The plain -case is pinned here because it is a COMPATIBILITY GUARANTEE — those keys are shared with -every non-targeted run and both toys, and adding a second stream must not move them. -""" - -from types import SimpleNamespace -from typing import Any, cast - -import jax.numpy as jnp -from jaxtyping import Array - -from param_decomp.core.configs import ( - CIMaskedReconLossConfig, - UnmaskedNoDeltaReconLossConfig, -) -from param_decomp.core.eval_schedule import Every -from param_decomp.core.run import EvalOperation -from param_decomp.core.train import TrainState -from param_decomp.experiments.lm.eval import CE_KL_VARIANTS -from param_decomp.experiments.lm.eval_context import LMEvalContext -from param_decomp.experiments.lm.scalar_eval_operations import ( - Stream, - _make_scalar_operation, - stream_log_prefix, -) - -TARGET_BATCHES = (jnp.zeros((1, 4), jnp.int32),) - - -def context(target_batches: tuple[jnp.ndarray, ...] | None) -> LMEvalContext: - """A context carrying only what the namespace rule reads; `state` is never touched.""" - return LMEvalContext( - state=cast(TrainState, cast(object, None)), - now_step=0, - pass_index=0, - batches=(), - target_batches=target_batches, - ) - - -def test_plain_run_keys_are_unchanged_by_the_two_stream_rule(): - assert stream_log_prefix("broad", context(None)) == "eval/" - - -def test_targeted_run_labels_the_broad_stream_and_leaves_the_target_bare(): - targeted = context(TARGET_BATCHES) - assert stream_log_prefix("target_data", targeted) == "eval/" - assert stream_log_prefix("broad", targeted) == "eval/nontarget_data/" - - -def test_the_bare_namespace_is_always_the_data_of_interest(): - """What makes the run kinds comparable AS OBJECTIVES — and the trap that comes with it: - the bare namespace is not the same DATA in both.""" - plain, targeted = context(None), context(TARGET_BATCHES) - assert stream_log_prefix("broad", plain) == stream_log_prefix("target_data", targeted) - assert stream_log_prefix("broad", targeted) != stream_log_prefix("broad", plain) - - -def test_every_stream_resolves_under_both_run_kinds(): - for stream in cast(tuple[Stream, ...], ("broad", "target_data")): - for target_batches in (None, TARGET_BATCHES): - prefix = stream_log_prefix(stream, context(target_batches)) - assert prefix.startswith("eval/") and prefix.endswith("/") - - -def test_one_operation_per_stream_reads_that_stream_and_labels_it(): - """The end-to-end shape the binder relies on: bind the SAME step twice, once per - stream, and the two operations read different batches and land under different keys — - which is what lets a metric be authored once and measured on both.""" - - def step( - _model: Any, _components: Any, _ci_fn: Any, value: Array, _key: Any - ) -> dict[str, Array]: - return {"l0/0.0_site": value} - - def operation_for(stream: Stream) -> EvalOperation[LMEvalContext]: - return _make_scalar_operation( - Every(1), - step, - ("l0/",), - cast(Any, object()), - jnp.array([0, 0], dtype=jnp.uint32), - train_steps=0, - eval_steps=1, - stream=stream, - ) - - targeted = LMEvalContext( - state=cast( - TrainState, - cast( - object, SimpleNamespace(decomposition=SimpleNamespace(components=None, ci_fn=None)) - ), - ), - now_step=0, - pass_index=0, - batches=(jnp.asarray(3.0),), - target_batches=(jnp.asarray(7.0),), - ) - broad = operation_for("broad").run(targeted) - target = operation_for("target_data").run(targeted) - - assert broad == {"eval/nontarget_data/l0/0.0_site": 3.0} - assert target == {"eval/l0/0.0_site": 7.0} - assert not broad.keys() & target.keys() - - -def test_narrow_kl_evals_name_arms_the_full_ce_kl_evaluator_also_emits(): - """The narrow metrics exist to cut CE/KL's clutter, NOT to respell its numbers: each - selects one arm of the same evaluator, so `eval/ce_kl/kl_` means the same thing - whether a run authored `CEandKLLosses` or the single-arm config. A new arm name here - that the full evaluator doesn't emit would silently give one quantity two names.""" - assert {"ci_masked", "unmasked"} <= set(CE_KL_VARIANTS) - - -def test_the_two_recon_configs_are_authorable_as_evals(): - """`coeff` is what separates the loss role from the eval role (`LossMetricConfig`), so - both must validate with none — this is the `PGDReconLoss` dual-role pattern.""" - for config in (CIMaskedReconLossConfig(), UnmaskedNoDeltaReconLossConfig()): - assert config.coeff is None - assert config.slow is False diff --git a/param_decomp/experiments/test_toy_eval.py b/param_decomp/experiments/test_toy_eval.py index 4452eaffd..be4468b52 100644 --- a/param_decomp/experiments/test_toy_eval.py +++ b/param_decomp/experiments/test_toy_eval.py @@ -46,7 +46,7 @@ def make_operation( compiler_options: dict[str, bool | int | str], *, inputs_for_context: Any, - log_prefix_for_context: Any, + log_prefix: Any, figure_rendering: Any, ) -> EvalOperation[Any]: captured.update( @@ -56,7 +56,7 @@ def make_operation( mesh=mesh, compiler_options=compiler_options, inputs_for_context=inputs_for_context, - log_prefix_for_context=log_prefix_for_context, + log_prefix=log_prefix, figure_rendering=figure_rendering, ) return EvalOperation(schedule, lambda _context: {}) diff --git a/param_decomp/experiments/toy_eval.py b/param_decomp/experiments/toy_eval.py index 1dddf2bce..0817f49d1 100644 --- a/param_decomp/experiments/toy_eval.py +++ b/param_decomp/experiments/toy_eval.py @@ -121,7 +121,7 @@ def well_temperedness_inputs(context: EvalInvocation) -> tuple[Array, jax.Array] inputs_for_context=well_temperedness_inputs, # A toy has ONE stream, so its keys carry no stream segment — the same # `eval/slow/well_temperedness/*` a plain LM run logs. - log_prefix_for_context=lambda _context: "eval/", + log_prefix="eval/", figure_rendering="synchronous" if wandb_configured else None, ) case CEandKLLossesConfig(): From 22b76f26617d648247c6d925041a3575a897afb9 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 01:18:10 +0000 Subject: [PATCH 04/15] rename(eval): the eval streams are "target" and "nontarget" MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Stream` spelled its two values `"broad"` and `"target_data"`. "Broad" describes what the nontarget stream happens to carry in today's tPD configs, not what it is — a stream can carry anything — so the vocabulary baked one experiment's choice into the type. The pair is now symmetric: `Literal["nontarget", "target"]`. Log keys are UNCHANGED: the segment stays `nontarget_data`. Only the internal stream value and the prose this PR introduced move. Scoped to what this PR added. The pre-existing "broad stream" prose in `core/{train,objective,configs,run}.py`, `experiments/lm/{config,targeted_data}.py`, `training_targeted.py`'s module docstring and SPEC.md is left alone — renaming it is a separate change, and SPEC.md is normative besides. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/configs.py | 2 +- .../core/tests/test_eval_averaging_parity.py | 2 +- .../lm/diagnostic_eval_operations.py | 4 +-- param_decomp/experiments/lm/eval_config.py | 4 +-- .../experiments/lm/eval_operations.py | 10 +++---- .../experiments/lm/scalar_eval_operations.py | 30 +++++++++---------- .../experiments/lm/training_targeted.py | 2 +- 7 files changed, 27 insertions(+), 27 deletions(-) diff --git a/param_decomp/core/configs.py b/param_decomp/core/configs.py index b230312d9..1bc70873d 100644 --- a/param_decomp/core/configs.py +++ b/param_decomp/core/configs.py @@ -989,7 +989,7 @@ class ResumeProvenance(BaseConfig): NONTARGET_STREAM = "nontarget_data" -"""The log-namespace segment for a tPD run's broad (non-target) stream. +"""The log-namespace segment for a tPD run's nontarget stream. ONE rule across `train/` and `eval/`: the data a run OPTIMIZES FOR is unlabelled, and anything else carries its stream, as the segment immediately after the tier — so diff --git a/param_decomp/core/tests/test_eval_averaging_parity.py b/param_decomp/core/tests/test_eval_averaging_parity.py index 6b02eb038..9bf2944bb 100644 --- a/param_decomp/core/tests/test_eval_averaging_parity.py +++ b/param_decomp/core/tests/test_eval_averaging_parity.py @@ -132,7 +132,7 @@ def step( jnp.array([0, 0], dtype=jnp.uint32), train_steps=0, eval_steps=2, - stream="broad", + stream="nontarget", ) context = LMEvalContext( state=_state_stub(), # pyright: ignore[reportArgumentType] diff --git a/param_decomp/experiments/lm/diagnostic_eval_operations.py b/param_decomp/experiments/lm/diagnostic_eval_operations.py index 389412908..66ffc6782 100644 --- a/param_decomp/experiments/lm/diagnostic_eval_operations.py +++ b/param_decomp/experiments/lm/diagnostic_eval_operations.py @@ -225,8 +225,8 @@ def run(context: LMEvalContext) -> LogRecord: renderer.submit( partial( _render_two_stream_ci_means, - stream_mean_cis(ci_fn, stream_batches("target_data", context)), - stream_mean_cis(ci_fn, stream_batches("broad", context)), + stream_mean_cis(ci_fn, stream_batches("target", context)), + stream_mean_cis(ci_fn, stream_batches("nontarget", context)), context.now_step, ) ) diff --git a/param_decomp/experiments/lm/eval_config.py b/param_decomp/experiments/lm/eval_config.py index 89f85a24b..97369b3f2 100644 --- a/param_decomp/experiments/lm/eval_config.py +++ b/param_decomp/experiments/lm/eval_config.py @@ -69,8 +69,8 @@ class TwoStreamCIMeanPerComponentConfig(BaseConfig): TARGET mean and coloured by stream. Supersedes authoring `CIMeanPerComponent` on a targeted run: it computes the same - broad-stream reduction plus the target-pool one, so authoring both would pay for the - broad pass twice. Refuses on a plain run, which has no target stream.""" + nontarget-stream reduction plus the target-stream one, so authoring both would pay + for the nontarget pass twice. Refuses on a plain run, which has no target stream.""" slow: ClassVar[bool] = True type: Literal["TwoStreamCIMeanPerComponent"] = "TwoStreamCIMeanPerComponent" diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index 2472d5b9d..7feaf68a4 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -90,7 +90,7 @@ def make_lm_evaluation( cannot: the prompt pool has no held-out split, so the targeted root draws pool batches the same way training does. `None` on a plain run, and that is what makes a plain run's metric set — and every one of its log keys — exactly what it was before targeted runs - existed: `data_streams` collapses to the single broad stream.""" + existed: `data_streams` collapses to the single nontarget stream.""" pd = built.pd capture_inputs = built.ci_fn.capture_keys data = built.data @@ -110,11 +110,11 @@ def batches(pass_index: int) -> list[jax.Array]: targeted = target_pool_batches_for is not None # Every stream a metric authored for BOTH measures; a plain run has exactly one. - data_streams: tuple[Stream, ...] = ("broad", "target_data") if targeted else ("broad",) + data_streams: tuple[Stream, ...] = ("nontarget", "target") if targeted else ("nontarget",) # The stream the run optimizes for, and the DEFAULT for any single-stream eval: a # diagnostic you read once should describe the data you are interpreting. On a plain run - # this IS the broad stream, so every such metric keeps the keys and the data it had. - optimized_stream: Stream = "target_data" if targeted else "broad" + # this IS the nontarget stream, so every such metric keeps the keys and data it had. + optimized_stream: Stream = "target" if targeted else "nontarget" def well_temperedness_inputs( context: LMEvalContext, @@ -272,7 +272,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo ) authored = {type(metric) for metric in eval.metrics} assert not {TwoStreamCIMeanPerComponentConfig, CIMeanPerComponentConfig} <= authored, ( - "TwoStreamCIMeanPerComponent already computes the broad-stream reduction " + "TwoStreamCIMeanPerComponent already computes the nontarget-stream reduction " "CIMeanPerComponent does, so authoring both pays for that pass twice" ) # The single-arm evals ARE arms of `CEandKLLosses`, under the same keys. Authoring both diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index 651ac444c..2b9b87518 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -26,17 +26,17 @@ from param_decomp.experiments.lm.eval_context import LMEvalContext from param_decomp.experiments.lm.eval_keys import EvalKeyStream -type Stream = Literal["broad", "target_data"] -"""Which DATA an eval operation measures. ONE value, not a (batch source, log prefix) pair: -the two always covary, and nothing should be able to spell target-pool batches under the -broad stream's log keys.""" +type Stream = Literal["nontarget", "target"] +"""Which STREAM an eval operation measures. ONE value, not a (batch source, log prefix) +pair: the two always covary, and nothing should be able to spell target-stream batches +under the nontarget stream's log keys.""" def stream_batches(stream: Stream, context: LMEvalContext) -> tuple[Array, ...]: match stream: - case "broad": + case "nontarget": return context.batches - case "target_data": + case "target": assert context.target_batches is not None, ( "target-stream metrics need a tPD run's prompt pool; a plain run has none" ) @@ -46,17 +46,17 @@ def stream_batches(stream: Stream, context: LMEvalContext) -> tuple[Array, ...]: def stream_log_prefix(stream: Stream, targeted: bool) -> str: """The log namespace for `stream` on a run of this kind. - ONE rule: the data the run is optimizing for is unlabelled, and anything else carries - its stream. A plain run has a single stream, so it stays `eval/` exactly as before. A - tPD run adds a second stream, so its broad corpus moves under `eval/nontarget_data/` - and the target pool takes the bare namespace. The consequence is deliberate: - `eval/l0/...` means "the data of interest" in both run kinds, which is what makes the - two comparable as objectives — but it is NOT the same data, so a corpus-vs-corpus - comparison across run kinds must read `eval/nontarget_data/` on the tPD side.""" + ONE rule: the stream the run is optimizing for is unlabelled, and anything else carries + its name. A plain run has a single stream, so it stays `eval/` exactly as before. A tPD + run adds a second, so its nontarget stream moves under `eval/nontarget_data/` and its + target stream takes the bare namespace. The consequence is deliberate: `eval/l0/...` + means "the stream of interest" in both run kinds, which is what makes the two comparable + as objectives — but it is NOT the same data, so comparing like for like across run kinds + must read `eval/nontarget_data/` on the tPD side.""" match stream: - case "broad": + case "nontarget": return f"eval/{NONTARGET_STREAM}/" if targeted else "eval/" - case "target_data": + case "target": return "eval/" diff --git a/param_decomp/experiments/lm/training_targeted.py b/param_decomp/experiments/lm/training_targeted.py index aa3d760bd..714671a2e 100644 --- a/param_decomp/experiments/lm/training_targeted.py +++ b/param_decomp/experiments/lm/training_targeted.py @@ -143,7 +143,7 @@ def sample_nontarget_batch(step: int) -> jax.Array: def eval_target_pool_batches(pass_index: int) -> list[jax.Array]: """The eval pass's TARGET stream: the same pure `(seed, step)` pool sampler - training uses, on the `seed + 1` stream the broad eval split already draws + training uses, on the `seed + 1` stream the nontarget eval split already draws from — so an eval never scores the exact rows the step just trained on.""" n_batches = eval_config.n_steps return [ From a684e82b192a32b439923a62ae4a3e117ceccd3f Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 01:21:13 +0000 Subject: [PATCH 05/15] docs: drop this PR's CLAUDE.md additions The two sections this branch added (the metric-namespace rule in core, the per-stream binding table in experiments) are reverted; both files are now byte-identical to upstream. Checked for statements the code change invalidates: the only line naming these operations ("domain-bound CEandKL/CI-L0/PGD/attention operations", core/CLAUDE.md) is still accurate, so nothing needed correcting. The namespace rule lives in the PR description and in the docstrings at the seams it governs (`Stream`, `stream_log_prefix`, `NONTARGET_STREAM`). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/CLAUDE.md | 37 ------------------------------ param_decomp/experiments/CLAUDE.md | 32 -------------------------- 2 files changed, 69 deletions(-) diff --git a/param_decomp/core/CLAUDE.md b/param_decomp/core/CLAUDE.md index 7a562879b..0df95f4bb 100644 --- a/param_decomp/core/CLAUDE.md +++ b/param_decomp/core/CLAUDE.md @@ -399,43 +399,6 @@ on one. priority-fusion). Don't chase graph-shrink refactors for compile time without new evidence. -## Metric namespaces: the data a run optimizes for is unlabelled - -Every logged key is `/[/]/` — `tier` is `train`/`eval`, and the -STREAM segment sits immediately after it (the only position that works uniformly, since -`eval` has several families: `ce_kl/`, `l0/`, `loss/`). The one rule: - -- The data the run OPTIMIZES FOR carries no stream segment. -- Anything else carries `nontarget_data/` (`configs.NONTARGET_STREAM`). - -So `train/loss/total` and `eval/l0/...` mean "the data of interest" in BOTH run kinds — a -plain run's corpus, a tPD run's prompt pool — while a targeted run's broad corpus reads -`train/nontarget_data/loss/total` / `eval/nontarget_data/l0/...`. The consequence is -deliberate and is the trap: the bare namespace is comparable across run kinds AS AN -OBJECTIVE, but it is not the same DATA, so a corpus-vs-corpus comparison must read -`eval/nontarget_data/` on the tPD side. - -**A plain run has ONE stream and therefore never emits the segment** — every plain-run and -toy key is exactly what it was before targeted runs existed. That is a COMPATIBILITY -GUARANTEE, not an accident: `scalar_eval_operations.stream_log_prefix` keys off -`context.target_batches is None`: on a plain run there is no second stream to name, so the -prefix is the same literal it always was. Nothing pins this in a test — if you change -`stream_log_prefix`, check the plain arm by hand. -Two consequences when adding a metric: - -- An eval that reads data must take BOTH its batches and its prefix from the same `Stream` - value (`stream_batches` / `stream_log_prefix`) — never a hardcoded `context.batches` or - `"eval/"`, which on a tPD run would read the corpus while claiming to be target data. - Core-side metrics shared with the toys take a `log_prefix_for_context` callback instead - (`well_temperedness_eval`). A single-stream eval binds to the OPTIMIZED stream by - default; only a metric authored for both takes `data_streams`. -- An eval that reads no batch at all (`WeightMagnitude`, the U/V norm ratios) has no stream - and keeps the bare namespace on both run kinds. - -One quantity gets ONE name across streams: `IMP_MIN_METRIC_NAMES` is shared by the target -stream (expanded in `run._METRIC_KEYS`) and the non-target stream (keyed in -`train.make_targeted_train_step`) so the two cannot drift apart. - ## Gotchas - **Process bring-up is config-derived, NEVER SLURM-sniffing** (`sharding.py`): the LM diff --git a/param_decomp/experiments/CLAUDE.md b/param_decomp/experiments/CLAUDE.md index 961c638e1..a997b85ce 100644 --- a/param_decomp/experiments/CLAUDE.md +++ b/param_decomp/experiments/CLAUDE.md @@ -175,38 +175,6 @@ experiments/ └── resid_mlp/ # ResidMLP (CPU): run.py + configs/ + test (target: param_decomp/targets/resid_mlp.py) ``` -## `eval.metrics` on a targeted run — one operation PER STREAM - -`make_lm_evaluation` binds each authored metric to every stream it measures, so a metric is -authored ONCE and a tPD run gets both readouts. The stream set comes from -`target_pool_batches_for`: `None` (the plain root) collapses it to the single broad stream, -so a plain run's operations and keys are untouched. The namespace rule itself is in -`param_decomp/core/CLAUDE.md`. - -| authored metric | streams it binds to | -|---|---| -| `CI_L0`, `PGDReconLoss`, `CIMaskedReconLoss`, `CEandKLLosses` | both on tPD, broad on plain | -| `UnmaskedNoDeltaReconLoss`, attn-patterns / hidden-acts recon, `WellTemperedness`, the site figures, the permutation plots, `IdentityCIError` | the OPTIMIZED stream only | -| `WeightMagnitude` | none — reads V/U, no batch | -| `TwoStreamCIMeanPerComponent` | both, in one figure; refuses on a plain run | -| `ArithmeticCIGrid` | none — brings its own probe grid | - -**A single-stream eval defaults to the OPTIMIZED stream** (`optimized_stream`, hence -`single_stream`): a diagnostic you read once should describe the data you are interpreting, -which on a tPD run is the prompt pool. On a plain run that tuple IS `("broad",)`, so those -metrics keep both the data and the keys they had. The consequence on a tPD run is that -`eval/nontarget_data/` appears ONLY for metrics deliberately bound to both streams — a -figure or scalar with no stream segment is target-pool data. - -`CIMaskedReconLoss` and `UnmaskedNoDeltaReconLoss` are authorable under `eval.metrics` as -well as `loss_metrics` — the `PGDReconLoss` dual-role pattern (`coeff` null in the eval -seat). As evals they are ONE arm of the CE/KL evaluator (`ce_kl/kl_ci_masked`, -`ce_kl/kl_unmasked`), which is how you get those two numbers without `CEandKLLosses`'s full -11-scalar record — and at one masked forward each instead of six. `UnmaskedNoDeltaReconLoss` -binds to the optimized stream alone deliberately: it is the non-target pass's OWN training -term, already reported there as a train loss, so an eval of it off-target would restate the -objective. - ## Sites and the family grammar Read `param_decomp/core/family.py`'s module docstring before authoring a From ee84951dbc603fed8988a0d77a83baa04ce60224 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 01:36:42 +0000 Subject: [PATCH 06/15] refactor: inline the two shared metric-name constants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `IMP_MIN_METRIC_NAMES` and `NONTARGET_STREAM` existed to stop the target and non-target spellings drifting apart. That is a rule about the future, not work the code needs doing: the names are consistent as written, and a later edit that makes them inconsistent is a later edit's problem. Both are gone. `run._METRIC_KEYS` goes back to the three literal rows it had before this branch, and `train.make_targeted_train_step` spells the non-target keys directly. The non-target imp-min name branches on the penalty kind exactly as `_METRIC_KEYS` does, so the two streams still spell that loss identically today — including when the config carries a custom `name:`, which keying off the term's `.name` would not have. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/configs.py | 25 +----------------- param_decomp/core/run.py | 8 +++--- param_decomp/core/train.py | 26 +++++++++---------- .../experiments/lm/scalar_eval_operations.py | 4 +-- 4 files changed, 20 insertions(+), 43 deletions(-) diff --git a/param_decomp/core/configs.py b/param_decomp/core/configs.py index 1bc70873d..0c2e845e7 100644 --- a/param_decomp/core/configs.py +++ b/param_decomp/core/configs.py @@ -473,8 +473,7 @@ class WeightMagnitudeConfig(BaseConfig): Reads the trained V/U alone — no forward pass and no eval batch, so it costs two norm reductions and a plot. The norms reduce on device; only `C` floats per site are pulled. - Stream-independent by construction: it never touches a batch, so it carries no stream - segment on either run kind.""" + """ slow: ClassVar[bool] = True type: Literal["WeightMagnitude"] = "WeightMagnitude" @@ -988,28 +987,6 @@ class ResumeProvenance(BaseConfig): """The parent's orbax `ckpts//` checkpoint step to initialize V/U + ci_fn from.""" -NONTARGET_STREAM = "nontarget_data" -"""The log-namespace segment for a tPD run's nontarget stream. - -ONE rule across `train/` and `eval/`: the data a run OPTIMIZES FOR is unlabelled, and -anything else carries its stream, as the segment immediately after the tier — so -`train/loss/X` / `eval/X` mean "the data of interest" in both run kinds, and the corpus -stream of a targeted run reads `train/nontarget_data/loss/X` / `eval/nontarget_data/X`. A -plain run has ONE stream and therefore never emits this segment: every plain-run key is -what it was before targeted runs existed.""" - - -IMP_MIN_METRIC_NAMES: dict[str, str] = { - "imp": "ImportanceMinimalityLoss", - "imp_smooth_l0": "SmoothL0ImportanceMinimalityLoss", - "freq": "FrequencyMinimalityLoss", -} -"""Step-record short key -> logged loss name, for the terms whose record key is not already -the term's class name. Shared so the target stream (expanded by `run._METRIC_KEYS`) and the -non-target stream (keyed in `train.make_targeted_train_step`) cannot drift apart: both -streams must spell the same quantity the same way.""" - - # --------------------------------------------------------------------------- # wandb.config shaping # --------------------------------------------------------------------------- diff --git a/param_decomp/core/run.py b/param_decomp/core/run.py index 3055c428c..66ce802c9 100644 --- a/param_decomp/core/run.py +++ b/param_decomp/core/run.py @@ -50,8 +50,6 @@ from param_decomp.core.ci_fn import CIFnArch from param_decomp.core.components import init_component_stacks from param_decomp.core.configs import ( - IMP_MIN_METRIC_NAMES, - NONTARGET_STREAM, AnyPDConfig, Cadence, NontargetConfig, @@ -223,7 +221,9 @@ def is_mesh_placed(a: object) -> bool: _METRIC_KEYS = { "total": "train/loss/total", "faith": "train/loss/FaithfulnessLoss", - **{short: f"train/loss/{name}" for short, name in IMP_MIN_METRIC_NAMES.items()}, + "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", @@ -335,7 +335,7 @@ def log(self, step: int, record: "LogRecord") -> None: _METRIC_KEYS.get( k, f"train/{k}" - if k.startswith(("grad_norms/", "loss/", "schedules/", f"{NONTARGET_STREAM}/")) + if k.startswith(("grad_norms/", "loss/", "schedules/", "nontarget_data/")) else k, ): v for k, v in record.items() diff --git a/param_decomp/core/train.py b/param_decomp/core/train.py index 3d3b6e6db..0faed3b7e 100644 --- a/param_decomp/core/train.py +++ b/param_decomp/core/train.py @@ -40,12 +40,7 @@ from param_decomp.core.adversary import PersistentAdversary, init_fresh_pgd_sources from param_decomp.core.ci_fn import CI, CIFn, evaluate_ci from param_decomp.core.components import ComponentStacks, VUShape -from param_decomp.core.configs import ( - IMP_MIN_METRIC_NAMES, - NONTARGET_STREAM, - LossCoeff, - SmoothL0ImportanceMinimalityLossConfig, -) +from param_decomp.core.configs import LossCoeff, SmoothL0ImportanceMinimalityLossConfig from param_decomp.core.jit_util import filter_jit from param_decomp.core.losses import ( ReconstructionLoss, @@ -1188,6 +1183,13 @@ def make_targeted_train_step[PreparedT]( ascend_replicate=ascend_replicate, ) nt_terms = objective.nontarget.recon + # The non-target pass runs the TARGET pass's imp-min term under its own coefficient, so + # it logs the same loss name the target side does (`run._METRIC_KEYS`). + nt_imp_name = ( + "SmoothL0ImportanceMinimalityLoss" + if atoms.imp_loss_key == "imp_smooth_l0" + else "ImportanceMinimalityLoss" + ) coeff_schedules: dict[str, LossCoeff] = { objective.target.imp.name: objective.target.imp.coeff, **{term.name: term.coeff for term in objective.target.recon}, @@ -1376,10 +1378,8 @@ def loss_fn( nt_imp_lp, nt_imp_freq = imp_min_terms(nt_ci.upper, atoms.imp_min, imp_min_param) nt_total = nt_imp_coeff * nt_imp_lp + freq_coeff * nt_imp_freq nt_aux = { - # Same spelling as the target stream's keys (which `run._METRIC_KEYS` expands - # from the same map), so one quantity is not two names. - f"{NONTARGET_STREAM}/loss/{IMP_MIN_METRIC_NAMES[atoms.imp_loss_key]}": nt_imp_lp, - f"{NONTARGET_STREAM}/loss/{IMP_MIN_METRIC_NAMES['freq']}": nt_imp_freq, + f"nontarget_data/loss/{nt_imp_name}": nt_imp_lp, + "nontarget_data/loss/FrequencyMinimalityLoss": nt_imp_freq, } nt_breakdowns = atoms.grid_losses( nt_terms, @@ -1392,8 +1392,8 @@ def loss_fn( nt_terms, nt_recon_coeffs, nt_breakdowns, strict=True ): nt_total = nt_total + coeff * breakdown.total - nt_aux[f"{NONTARGET_STREAM}/loss/{term.name}"] = breakdown.total - nt_aux[f"{NONTARGET_STREAM}/loss/total"] = nt_total + nt_aux[f"nontarget_data/loss/{term.name}"] = breakdown.total + nt_aux["nontarget_data/loss/total"] = nt_total total_loss = total_loss + nt_total reported_total = reported_total + nt_total return total_loss, (reported_total, imp_lp, imp_freq, term_breakdowns, nt_aux) @@ -1465,7 +1465,7 @@ def loss_fn( | wd_metrics | _scheduled_coeff_metrics(step_f32, atoms.total_steps, "", coeff_schedules) | _scheduled_coeff_metrics( - step_f32, atoms.total_steps, f"{NONTARGET_STREAM}/", nontarget_coeff_schedules + step_f32, atoms.total_steps, "nontarget_data/", nontarget_coeff_schedules ) ) return new_state, metrics diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index 2b9b87518..d410685ad 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -7,7 +7,7 @@ from jax.sharding import Mesh from jaxtyping import Array, PRNGKeyArray -from param_decomp.core.configs import NONTARGET_STREAM, CI_L0Config, PGDReconLossConfig +from param_decomp.core.configs import CI_L0Config, PGDReconLossConfig from param_decomp.core.eval_schedule import EvalSchedule from param_decomp.core.metrics import BarChart, LogRecord from param_decomp.core.model import CaptureKeys, DecomposedModel @@ -55,7 +55,7 @@ def stream_log_prefix(stream: Stream, targeted: bool) -> str: must read `eval/nontarget_data/` on the tPD side.""" match stream: case "nontarget": - return f"eval/{NONTARGET_STREAM}/" if targeted else "eval/" + return "eval/nontarget_data/" if targeted else "eval/" case "target": return "eval/" From 9d7e13855f0097c48f803dc3a2c723f690dfcc36 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 01:50:36 +0000 Subject: [PATCH 07/15] refactor(eval): drop well-temperedness's log_prefix parameter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parameter could only ever be `"eval/"`. WellTemperedness binds to the stream the run optimizes for, and that stream's prefix is bare in BOTH run kinds — a plain run's nontarget stream and a tPD run's target stream both resolve to `eval/`. So the argument carried no information: the toys passed the literal, and the LM binder passed an expression with one possible value. `well_temperedness_eval.py` and its test are back to their upstream contents; this PR no longer touches either. With that call site gone, `stream_log_prefix` had one caller left — the wrapper that fed it a context — so the two collapse back into one function taking the context, which is what it was before well-temperedness needed a bind-time value. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- .../core/tests/test_well_temperedness_eval.py | 1 - param_decomp/core/well_temperedness_eval.py | 26 ++++++----------- .../lm/diagnostic_eval_operations.py | 8 +++--- .../experiments/lm/eval_operations.py | 2 -- .../experiments/lm/scalar_eval_operations.py | 28 ++++++++----------- .../experiments/lm/test_eval_operations.py | 2 -- param_decomp/experiments/test_toy_eval.py | 2 -- param_decomp/experiments/toy_eval.py | 3 -- 8 files changed, 24 insertions(+), 48 deletions(-) diff --git a/param_decomp/core/tests/test_well_temperedness_eval.py b/param_decomp/core/tests/test_well_temperedness_eval.py index b3ea5af6e..4b5ae9673 100644 --- a/param_decomp/core/tests/test_well_temperedness_eval.py +++ b/param_decomp/core/tests/test_well_temperedness_eval.py @@ -46,7 +46,6 @@ def unexpected_render(_ablations: Ablations) -> bytes: mesh=None, compiler_options={}, inputs_for_context=lambda _context: (jnp.zeros((1,)), jax.random.PRNGKey(0)), - log_prefix="eval/", figure_rendering=None, ) state = SimpleNamespace(decomposition=SimpleNamespace(components=object(), ci_fn=object())) diff --git a/param_decomp/core/well_temperedness_eval.py b/param_decomp/core/well_temperedness_eval.py index dab24bb0f..3b6fd9043 100644 --- a/param_decomp/core/well_temperedness_eval.py +++ b/param_decomp/core/well_temperedness_eval.py @@ -32,13 +32,8 @@ well_temperedness_log_entries, ) -_NAMESPACE = "slow/well_temperedness/" -"""The part of the key BELOW the stream namespace. The stream itself is the caller's to -supply: this metric samples the eval distribution, so on a two-stream run it belongs to -whichever stream those batches came from.""" -_FIGURE_STEP_KEY = f"eval/{_NAMESPACE}figure_step" -"""The figure step AXIS, deliberately stream-independent — a second stream must not fork -the axis W&B serializes these renders against (SPEC S28).""" +_PREFIX = "eval/slow/well_temperedness/" +_FIGURE_KEY = f"{_PREFIX}figures/preactivation_vs_ablation_damage" _MAX_FIGURE_LOCATIONS = 48 type FigureRendering = BackgroundRenderer | Literal["synchronous"] | None @@ -103,11 +98,11 @@ def _plot_preactivation_vs_damage(ablations: Ablations) -> bytes: return png_buffer.getvalue() -def _render_deferred(ablations: Ablations, figure_key: str, now_step: int) -> DeferredMediaRecord: +def _render_deferred(ablations: Ablations, now_step: int) -> DeferredMediaRecord: return DeferredMediaRecord( - step_key=_FIGURE_STEP_KEY, + step_key=f"{_PREFIX}figure_step", step=now_step, - media={figure_key: _plot_preactivation_vs_damage(ablations)}, + media={_FIGURE_KEY: _plot_preactivation_vs_damage(ablations)}, ) @@ -119,7 +114,6 @@ def make_well_temperedness_operation[ContextT: EvalInvocation]( mesh: Mesh | None, compiler_options: dict[str, bool | int | str], inputs_for_context: Callable[[ContextT], tuple[Array, PRNGKeyArray]], - log_prefix: str, figure_rendering: FigureRendering, ) -> EvalOperation[ContextT]: if figure_rendering is not None: @@ -131,10 +125,6 @@ def make_well_temperedness_operation[ContextT: EvalInvocation]( measure_ablations = make_well_temperedness_step( model, ci_capture_keys, metric, mesh, compiler_options ) - # Which stream this operation measures is fixed when it is bound, so the namespace is a - # constant here rather than something to re-derive from every context. - prefix = f"{log_prefix}{_NAMESPACE}" - figure_key = f"{prefix}figures/preactivation_vs_ablation_damage" def run(context: ContextT) -> LogRecord: inputs, sampling_key = inputs_for_context(context) @@ -147,16 +137,16 @@ def run(context: ContextT) -> LogRecord: ) ablations = jax.device_get(device_ablations) log_record: dict[str, float | PNGImage] = { - f"{prefix}{name}": value + f"{_PREFIX}{name}": value for name, value in well_temperedness_log_entries(ablations, site_groups).items() } match figure_rendering: case None: pass case "synchronous": - log_record[figure_key] = PNGImage(_plot_preactivation_vs_damage(ablations)) + log_record[_FIGURE_KEY] = PNGImage(_plot_preactivation_vs_damage(ablations)) case BackgroundRenderer() as renderer: - renderer.submit(partial(_render_deferred, ablations, figure_key, context.now_step)) + renderer.submit(partial(_render_deferred, ablations, context.now_step)) return log_record return EvalOperation(schedule, run) diff --git a/param_decomp/experiments/lm/diagnostic_eval_operations.py b/param_decomp/experiments/lm/diagnostic_eval_operations.py index 66ffc6782..1d5996291 100644 --- a/param_decomp/experiments/lm/diagnostic_eval_operations.py +++ b/param_decomp/experiments/lm/diagnostic_eval_operations.py @@ -63,8 +63,8 @@ from param_decomp.experiments.lm.eval_keys import EvalKeyStream from param_decomp.experiments.lm.scalar_eval_operations import ( Stream, - context_log_prefix, stream_batches, + stream_log_prefix, ) @@ -122,7 +122,7 @@ def run(context: LMEvalContext) -> LogRecord: run_key, EvalKeyStream.ATTENTION_PATTERNS * train_steps + context.pass_index ), ) - prefix = context_log_prefix(stream, context) + prefix = stream_log_prefix(stream, context) return { f"{prefix}loss/{name}": value for name, value in attn_patterns_log_entries(metric.type, reductions).items() @@ -160,7 +160,7 @@ def run(context: LMEvalContext) -> LogRecord: run_key, EvalKeyStream.HIDDEN_ACTS * train_steps + context.pass_index ), ) - prefix = context_log_prefix(stream, context) + prefix = stream_log_prefix(stream, context) return { f"{prefix}slow/loss/{name}": value for name, value in hidden_acts_log_entries(metric.type, reductions).items() @@ -305,7 +305,7 @@ def run(context: LMEvalContext) -> LogRecord: match metric: case IdentityCIErrorConfig(): errors = compute_identity_ci_errors(spec, position_ci, IDENTITY_CI_ERROR_TOLERANCE) - prefix = context_log_prefix(stream, context) + prefix = stream_log_prefix(stream, context) return {f"{prefix}slow/{name}": value for name, value in errors.items()} case UVPlotsConfig(): include_ci_heatmaps = False diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index 7feaf68a4..e5e373824 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -61,7 +61,6 @@ make_fresh_pgd_operation, make_single_variant_kl_operation, stream_batches, - stream_log_prefix, ) from param_decomp.infra.dataset_store import read_dataset_meta from param_decomp.pretrain.batch_data import BatchSchedule, ShardServer, scan_shards @@ -232,7 +231,6 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo mesh, compiler_options, inputs_for_context=well_temperedness_inputs, - log_prefix=stream_log_prefix(optimized_stream, targeted), figure_rendering=renderer if sink.accepts_deferred_media else None, ), ) diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index d410685ad..536767e07 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -43,16 +43,18 @@ def stream_batches(stream: Stream, context: LMEvalContext) -> tuple[Array, ...]: return context.target_batches -def stream_log_prefix(stream: Stream, targeted: bool) -> str: - """The log namespace for `stream` on a run of this kind. +def stream_log_prefix(stream: Stream, context: LMEvalContext) -> str: + """The log namespace for `stream`, given what kind of run this is. ONE rule: the stream the run is optimizing for is unlabelled, and anything else carries - its name. A plain run has a single stream, so it stays `eval/` exactly as before. A tPD - run adds a second, so its nontarget stream moves under `eval/nontarget_data/` and its - target stream takes the bare namespace. The consequence is deliberate: `eval/l0/...` - means "the stream of interest" in both run kinds, which is what makes the two comparable - as objectives — but it is NOT the same data, so comparing like for like across run kinds - must read `eval/nontarget_data/` on the tPD side.""" + its name. A plain run has a single stream, so it stays `eval/` exactly as before — + `context.target_batches is None` is what says so. A tPD run adds a second, so its + nontarget stream moves under `eval/nontarget_data/` and its target stream takes the bare + namespace. The consequence is deliberate: `eval/l0/...` means "the stream of interest" in + both run kinds, which is what makes the two comparable as objectives — but it is NOT the + same data, so comparing like for like across run kinds must read `eval/nontarget_data/` + on the tPD side.""" + targeted = context.target_batches is not None match stream: case "nontarget": return "eval/nontarget_data/" if targeted else "eval/" @@ -60,12 +62,6 @@ def stream_log_prefix(stream: Stream, targeted: bool) -> str: return "eval/" -def context_log_prefix(stream: Stream, context: LMEvalContext) -> str: - """`stream_log_prefix` for an operation already running: `target_batches is None` is - what tells a bound operation which run kind it landed in.""" - return stream_log_prefix(stream, context.target_batches is not None) - - def _make_scalar_operation( schedule: EvalSchedule, step: ScalarStep, @@ -77,7 +73,7 @@ def _make_scalar_operation( stream: Stream, ) -> EvalOperation[LMEvalContext]: def run(context: LMEvalContext) -> LogRecord: - log_prefix = context_log_prefix(stream, context) + log_prefix = stream_log_prefix(stream, context) sums: dict[str, Array] = {} for batch_index, tokens in enumerate(stream_batches(stream, context)): key = random.fold_in( @@ -196,7 +192,7 @@ def make_ci_l0_operation( def run(context: LMEvalContext) -> LogRecord: record = dict(scalars.run(context)) - log_prefix = context_log_prefix(stream, context) + log_prefix = stream_log_prefix(stream, context) prefix = f"{log_prefix}l0/{metric.ci_alive_threshold}_" record[f"{log_prefix}l0/bar_chart"] = BarChart( rows=tuple( diff --git a/param_decomp/experiments/lm/test_eval_operations.py b/param_decomp/experiments/lm/test_eval_operations.py index 772f13a4b..a967e20b0 100644 --- a/param_decomp/experiments/lm/test_eval_operations.py +++ b/param_decomp/experiments/lm/test_eval_operations.py @@ -47,7 +47,6 @@ def make_operation( compiler_options: dict[str, bool | int | str], *, inputs_for_context: Any, - log_prefix: Any, figure_rendering: Any, ) -> EvalOperation[Any]: captured.update( @@ -57,7 +56,6 @@ def make_operation( mesh=mesh, compiler_options=compiler_options, inputs_for_context=inputs_for_context, - log_prefix=log_prefix, figure_rendering=figure_rendering, ) return EvalOperation(schedule, lambda _context: {}) diff --git a/param_decomp/experiments/test_toy_eval.py b/param_decomp/experiments/test_toy_eval.py index be4468b52..b2290fed9 100644 --- a/param_decomp/experiments/test_toy_eval.py +++ b/param_decomp/experiments/test_toy_eval.py @@ -46,7 +46,6 @@ def make_operation( compiler_options: dict[str, bool | int | str], *, inputs_for_context: Any, - log_prefix: Any, figure_rendering: Any, ) -> EvalOperation[Any]: captured.update( @@ -56,7 +55,6 @@ def make_operation( mesh=mesh, compiler_options=compiler_options, inputs_for_context=inputs_for_context, - log_prefix=log_prefix, figure_rendering=figure_rendering, ) return EvalOperation(schedule, lambda _context: {}) diff --git a/param_decomp/experiments/toy_eval.py b/param_decomp/experiments/toy_eval.py index 0817f49d1..38414e01f 100644 --- a/param_decomp/experiments/toy_eval.py +++ b/param_decomp/experiments/toy_eval.py @@ -119,9 +119,6 @@ def well_temperedness_inputs(context: EvalInvocation) -> tuple[Array, jax.Array] mesh, compiler_options, inputs_for_context=well_temperedness_inputs, - # A toy has ONE stream, so its keys carry no stream segment — the same - # `eval/slow/well_temperedness/*` a plain LM run logs. - log_prefix="eval/", figure_rendering="synchronous" if wandb_configured else None, ) case CEandKLLossesConfig(): From f20581462f9aa6e506e801f48489bb56eeac6dc8 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 02:15:19 +0000 Subject: [PATCH 08/15] refactor(train): one name for the imp-min term, no key rewriting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two leftovers from earlier rounds, both mechanism where a value would do. `_scheduled_coeff_metrics` had grown a `stream_prefix` parameter that two of its three callers passed `""` for. It is back to its upstream signature; the one caller that needs a prefix adds it, which is also where the concern belongs — the shared helper has nothing to do with streams. The non-target imp-min loss name was derived by branching on the penalty kind to reproduce the spelling `run._METRIC_KEYS` uses. But the term already carries its name: `ImportanceMinimalityTerm.name` is `cfg.name or cfg.type`, and `cfg.type` IS the class-name literal, so it equals that spelling for any config that doesn't override it. `imp_name = objective.target.imp.name`, used for the target coefficient, the non-target coefficient, and the non-target loss key — one value, three places, no branch. A config with a custom `name:` now spells all three that way, where upstream already spelled its coefficient by `.name` and its loss key by class name. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/train.py | 48 ++++++++++++++------------------------ 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/param_decomp/core/train.py b/param_decomp/core/train.py index 0faed3b7e..08f9c3184 100644 --- a/param_decomp/core/train.py +++ b/param_decomp/core/train.py @@ -193,18 +193,12 @@ def per_slice_sq(stack: Float[Array, "g a b"]) -> Float[Array, " g"]: def _scheduled_coeff_metrics( - step_f32: Array, total_steps: int, stream_prefix: str, coeffs: dict[str, LossCoeff] + step_f32: Array, total_steps: int, coeffs: dict[str, LossCoeff] ) -> dict[str, Array]: """Per-step values of the SCHEDULED coefficients only — a constant would be log - noise, and a moving coefficient invisible in wandb is a debugging trap. - - `stream_prefix` is `""` for the data the run optimizes for and `"nontarget_data/"` for a - targeted run's corpus stream, so the stream segment leads the key exactly as it does on - the loss keys.""" + noise, and a moving coefficient invisible in wandb is a debugging trap.""" return { - f"{stream_prefix}schedules/coeff/{name}": scheduled_value_traced( - step_f32, total_steps, coeff - ) + f"schedules/coeff/{name}": scheduled_value_traced(step_f32, total_steps, coeff) for name, coeff in coeffs.items() if isinstance(coeff, ScheduleConfig) } @@ -1086,7 +1080,7 @@ def loss_fn( step_f32=step_f32, ) | {"faith": faith_loss} - | _scheduled_coeff_metrics(step_f32, atoms.total_steps, "", coeff_schedules) + | _scheduled_coeff_metrics(step_f32, atoms.total_steps, coeff_schedules) ) return new_state, metrics @@ -1183,15 +1177,11 @@ def make_targeted_train_step[PreparedT]( ascend_replicate=ascend_replicate, ) nt_terms = objective.nontarget.recon - # The non-target pass runs the TARGET pass's imp-min term under its own coefficient, so - # it logs the same loss name the target side does (`run._METRIC_KEYS`). - nt_imp_name = ( - "SmoothL0ImportanceMinimalityLoss" - if atoms.imp_loss_key == "imp_smooth_l0" - else "ImportanceMinimalityLoss" - ) + # Both passes run the same imp-min term, each under its own coefficient, so both log it + # under that term's name. + imp_name = objective.target.imp.name coeff_schedules: dict[str, LossCoeff] = { - objective.target.imp.name: objective.target.imp.coeff, + imp_name: objective.target.imp.coeff, **{term.name: term.coeff for term in objective.target.recon}, **{ f"{term.name}/hidden_acts_reconstruction": term.hidden_acts_reconstruction.coeff @@ -1200,14 +1190,9 @@ def make_targeted_train_step[PreparedT]( }, } if objective.target.imp.cfg.frequency is not None: - coeff_schedules[f"{objective.target.imp.name}/frequency"] = ( - objective.target.imp.cfg.frequency.coeff - ) - # The non-target imp-min coefficient scales the SAME term the target stream spells by - # name, so it is spelled that way here too — the two streams' coefficients for one - # quantity must be readable as a pair. + coeff_schedules[f"{imp_name}/frequency"] = objective.target.imp.cfg.frequency.coeff nontarget_coeff_schedules: dict[str, LossCoeff] = { - objective.target.imp.name: objective.nontarget.impmin_coeff, + imp_name: objective.nontarget.impmin_coeff, **{term.name: term.coeff for term in nt_terms}, } @@ -1378,7 +1363,7 @@ def loss_fn( nt_imp_lp, nt_imp_freq = imp_min_terms(nt_ci.upper, atoms.imp_min, imp_min_param) nt_total = nt_imp_coeff * nt_imp_lp + freq_coeff * nt_imp_freq nt_aux = { - f"nontarget_data/loss/{nt_imp_name}": nt_imp_lp, + f"nontarget_data/loss/{imp_name}": nt_imp_lp, "nontarget_data/loss/FrequencyMinimalityLoss": nt_imp_freq, } nt_breakdowns = atoms.grid_losses( @@ -1463,10 +1448,13 @@ def loss_fn( ) | nt_aux | wd_metrics - | _scheduled_coeff_metrics(step_f32, atoms.total_steps, "", coeff_schedules) - | _scheduled_coeff_metrics( - step_f32, atoms.total_steps, "nontarget_data/", nontarget_coeff_schedules - ) + | _scheduled_coeff_metrics(step_f32, atoms.total_steps, coeff_schedules) + | { + f"nontarget_data/{key}": value + for key, value in _scheduled_coeff_metrics( + step_f32, atoms.total_steps, nontarget_coeff_schedules + ).items() + } ) return new_state, metrics From 4e264bab069e8e49c50df8cd6047de7bd2a3cf48 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 02:28:10 +0000 Subject: [PATCH 09/15] refactor(slow-eval): lift the raw-sample helper to module level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `sample` was a closure that closed over nothing — both arguments were already explicit — so it is a plain private function, `_raw_sample`. Its docstring now leads with what it returns rather than with the edge case that motivated it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/slow_eval.py | 17 +++++++++-------- param_decomp/core/train.py | 2 -- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/param_decomp/core/slow_eval.py b/param_decomp/core/slow_eval.py index 430da1e26..f7e49c91c 100644 --- a/param_decomp/core/slow_eval.py +++ b/param_decomp/core/slow_eval.py @@ -199,6 +199,13 @@ def slow_eval_step( return filter_jit(slow_eval_step, compiler_options=compiler_options) +def _raw_sample(chunks: dict[str, list[np.ndarray]], site: str) -> np.ndarray: + """A site's kept raw values, or empty when none were kept: `n_batches_accum=0` asks for + NO sample, and a caller that reads only `ci_sums` (the two-stream CI mean) would + otherwise pay a host gather of every position's CI just to discard it.""" + return np.concatenate(chunks[site]) if site in chunks else np.empty(0, np.float32) + + def accumulate_site_reductions( slow_eval_step: SlowEvalStep, model: DecomposedModel, @@ -242,19 +249,13 @@ def accumulate_site_reductions( ) ) - def sample(chunks: dict[str, list[np.ndarray]], site: str) -> np.ndarray: - """`n_batches_accum=0` asks for NO raw sample, so a site has no chunks at all: a - caller that reads only `ci_sums` (the two-stream CI mean) would otherwise pay a - host gather of every position's CI just to discard it.""" - return np.concatenate(chunks[site]) if site in chunks else np.empty(0, np.float32) - return { site: SiteReduction( density_counts=density[site], ci_sums=sums[site], n_positions=total_positions, - lower_sample=sample(lower_chunks, site), - preactivations_sample=sample(preactivations_chunks, site), + lower_sample=_raw_sample(lower_chunks, site), + preactivations_sample=_raw_sample(preactivations_chunks, site), density_hist=hist.get(site), ) for site in density diff --git a/param_decomp/core/train.py b/param_decomp/core/train.py index 08f9c3184..0c48cb2e0 100644 --- a/param_decomp/core/train.py +++ b/param_decomp/core/train.py @@ -1177,8 +1177,6 @@ def make_targeted_train_step[PreparedT]( ascend_replicate=ascend_replicate, ) nt_terms = objective.nontarget.recon - # Both passes run the same imp-min term, each under its own coefficient, so both log it - # under that term's name. imp_name = objective.target.imp.name coeff_schedules: dict[str, LossCoeff] = { imp_name: objective.target.imp.coeff, From 047c23963fe35c0db214fb4e1003b624de0007d9 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 02:40:04 +0000 Subject: [PATCH 10/15] refactor: keep the diff to what the new evals and the renaming need MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An audit of every hunk against upstream, asking whether it is required either to add the new evals/plots or to make the names consistent. Five were not. - `_figure_record` de-duplicated a record construction that already appeared twice upstream, so adding two renderers meant rewriting two functions this PR has no business touching. Both are back to upstream; the two new renderers spell the constructor out. - `make_lm_evaluation`'s `target_pool_batches_for` and `LMEvalContext.target_batches` default to `None`. A plain run has one stream and says so by omission, which takes `experiments/lm/training.py` and `test_eval_operations.py` out of the diff entirely — both are shared with the standard PD path. - `imp_name` no longer replaces upstream's two existing `objective.target.imp.name` uses; it is introduced beside the non-target keys that need it. - `emit_ce_difference` re-declared a filter `_make_scalar_operation` already applies through `prefixes`: the narrow arms drop `ce_kl/ce_difference_*` host-side whether or not the step emits them. `make_ce_kl_step` keeps its upstream body here. 15 files instead of 17, and `core/train.py` no longer touches the plain step's coefficient plumbing at all. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/train.py | 10 +++++-- .../lm/diagnostic_eval_operations.py | 30 +++++++++++-------- param_decomp/experiments/lm/eval.py | 23 +++++++------- param_decomp/experiments/lm/eval_context.py | 2 +- .../experiments/lm/eval_operations.py | 2 +- .../experiments/lm/scalar_eval_operations.py | 5 +--- .../experiments/lm/test_eval_operations.py | 2 -- param_decomp/experiments/lm/training.py | 10 +------ 8 files changed, 39 insertions(+), 45 deletions(-) diff --git a/param_decomp/core/train.py b/param_decomp/core/train.py index 0c48cb2e0..69b419c42 100644 --- a/param_decomp/core/train.py +++ b/param_decomp/core/train.py @@ -1177,9 +1177,8 @@ def make_targeted_train_step[PreparedT]( ascend_replicate=ascend_replicate, ) nt_terms = objective.nontarget.recon - imp_name = objective.target.imp.name coeff_schedules: dict[str, LossCoeff] = { - imp_name: objective.target.imp.coeff, + objective.target.imp.name: objective.target.imp.coeff, **{term.name: term.coeff for term in objective.target.recon}, **{ f"{term.name}/hidden_acts_reconstruction": term.hidden_acts_reconstruction.coeff @@ -1188,7 +1187,12 @@ def make_targeted_train_step[PreparedT]( }, } if objective.target.imp.cfg.frequency is not None: - coeff_schedules[f"{imp_name}/frequency"] = objective.target.imp.cfg.frequency.coeff + coeff_schedules[f"{objective.target.imp.name}/frequency"] = ( + objective.target.imp.cfg.frequency.coeff + ) + # Both passes run the same imp-min term, each under its own coefficient, so both log it + # under that term's name. + imp_name = objective.target.imp.name nontarget_coeff_schedules: dict[str, LossCoeff] = { imp_name: objective.nontarget.impmin_coeff, **{term.name: term.coeff for term in nt_terms}, diff --git a/param_decomp/experiments/lm/diagnostic_eval_operations.py b/param_decomp/experiments/lm/diagnostic_eval_operations.py index 1d5996291..c32bd3a49 100644 --- a/param_decomp/experiments/lm/diagnostic_eval_operations.py +++ b/param_decomp/experiments/lm/diagnostic_eval_operations.py @@ -68,16 +68,15 @@ ) -def _figure_record(now_step: int, media: dict[str, bytes]) -> DeferredMediaRecord: - """A slow-tier figure batch on the dedicated figure-step axis (SPEC S28).""" - return DeferredMediaRecord(step_key="slow_eval/figure_step", step=now_step, media=media) - - def _render_selected_figures( reductions: dict[str, SiteReduction], wanted: set[str], now_step: int ) -> DeferredMediaRecord: figures = render_slow_eval_figures(reductions) - return _figure_record(now_step, {f"slow_eval/{name}": figures[name] for name in wanted}) + return DeferredMediaRecord( + step_key="slow_eval/figure_step", + step=now_step, + media={f"slow_eval/{name}": figures[name] for name in wanted}, + ) def _render_permutation( @@ -90,7 +89,11 @@ def _render_permutation( figures = render_permutation_figures(spec, position_ci, components) if not include_ci_heatmaps: figures = {key: value for key, value in figures.items() if key == "figures/uv_matrices"} - return _figure_record(now_step, {f"slow_eval/{name}": value for name, value in figures.items()}) + return DeferredMediaRecord( + step_key="slow_eval/figure_step", + step=now_step, + media={f"slow_eval/{name}": value for name, value in figures.items()}, + ) def make_attention_operation( @@ -172,8 +175,10 @@ def run(context: LMEvalContext) -> LogRecord: def _render_weight_magnitudes( magnitudes: dict[str, np.ndarray], now_step: int ) -> DeferredMediaRecord: - return _figure_record( - now_step, {"slow_eval/figures/weight_magnitude": plot_weight_magnitudes(magnitudes)} + return DeferredMediaRecord( + step_key="slow_eval/figure_step", + step=now_step, + media={"slow_eval/figures/weight_magnitude": plot_weight_magnitudes(magnitudes)}, ) @@ -196,9 +201,10 @@ def _render_two_stream_ci_means( now_step: int, ) -> DeferredMediaRecord: linear, log = plot_mean_component_cis_two_streams(target, nontarget) - return _figure_record( - now_step, - { + return DeferredMediaRecord( + step_key="slow_eval/figure_step", + step=now_step, + media={ "slow_eval/figures/ci_mean_per_component_two_streams": linear, "slow_eval/figures/ci_mean_per_component_two_streams_log": log, }, diff --git a/param_decomp/experiments/lm/eval.py b/param_decomp/experiments/lm/eval.py index 748fb4b4a..2c16d3abe 100644 --- a/param_decomp/experiments/lm/eval.py +++ b/param_decomp/experiments/lm/eval.py @@ -222,7 +222,6 @@ def make_ce_kl_step[PreparedT]( mesh: Mesh | None = None, compiler_options: dict[str, bool | int | str] | None = None, *, - emit_ce_difference: bool, rounding_threshold: float | None = None, n_valid_rows: int | None = None, ) -> ScalarStep: @@ -231,8 +230,8 @@ def make_ce_kl_step[PreparedT]( The masks for EVERY arm are drawn whether or not the arm is selected, so a narrowed evaluator's numbers are bit-identical to the full one's — only the forwards are skipped, and XLA drops the unused draws. `emit_ce_difference` adds the CE-vs-target - delta for each selected arm (free: same logits, no extra forward). `rounding_threshold` - belongs to the rounded arm alone, so it is present exactly when that arm is.""" + `rounding_threshold` belongs to the + rounded arm alone, so it is present exactly when that arm is.""" assert model_static.has_position_axis, "CEandKLLosses is LM-only and requires a position axis" assert variants, "a CE/KL evaluator with no arms measures nothing" assert ("rounded_masked" in variants) == (rounding_threshold is not None), ( @@ -297,18 +296,17 @@ def eval_step( for name, (masks, deltas) in variant_masks.items() if name in variants } + target_ce = _ce(batch, batch.clean.output) metrics = { f"ce_kl/kl_{name}": _kl(batch, logits) for name, logits in variant_logits.items() } - if emit_ce_difference: - target_ce = _ce(batch, batch.clean.output) - metrics.update( - { - f"ce_kl/ce_difference_{name}": _ce(batch, variant_logits[name]) - target_ce - for name in variant_logits - if name != "zero_masked" - } - ) + metrics.update( + { + f"ce_kl/ce_difference_{name}": _ce(batch, variant_logits[name]) - target_ce + for name in variant_logits + if name != "zero_masked" + } + ) return metrics return filter_jit(eval_step, compiler_options=compiler_options) @@ -464,7 +462,6 @@ def make_eval_step[PreparedT]( CE_KL_VARIANTS, mesh, compiler_options, - emit_ce_difference=True, rounding_threshold=rounding_threshold, n_valid_rows=n_valid_rows, ) diff --git a/param_decomp/experiments/lm/eval_context.py b/param_decomp/experiments/lm/eval_context.py index 15d4ad3af..72d06531c 100644 --- a/param_decomp/experiments/lm/eval_context.py +++ b/param_decomp/experiments/lm/eval_context.py @@ -11,7 +11,7 @@ class LMEvalContext(EvalInvocation): pass_index: int batches: tuple[jax.Array, ...] - target_batches: tuple[jax.Array, ...] | None + target_batches: tuple[jax.Array, ...] | None = None """A tPD run's prompt-pool draws, which `data.eval` cannot supply — the pool has no held-out split, so the targeted root draws them exactly as training does. `None` on a plain run, which HAS no second stream; that `None` is also what tells every log key diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index e5e373824..333ebc903 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -80,7 +80,7 @@ def make_lm_evaluation( n_proc: int, sink: MetricsSink, compiler_options: dict[str, bool | int | str], - target_pool_batches_for: Callable[[int], list[jax.Array]] | None, + target_pool_batches_for: Callable[[int], list[jax.Array]] | None = None, ) -> Evaluation[LMEvalContext]: """Construct the executable operations for every authored LM metric — one PER STREAM the metric measures. diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index 536767e07..edbbc4083 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -115,7 +115,6 @@ def make_ce_kl_operation( CE_KL_VARIANTS, mesh, compiler_options, - emit_ce_difference=True, rounding_threshold=metric.rounding_threshold, ), ("ce_kl/",), @@ -148,9 +147,7 @@ def make_single_variant_kl_operation( same number is one name across run kinds; the config type names the construction.""" return _make_scalar_operation( schedule, - make_ce_kl_step( - model, ci_capture_keys, (variant,), mesh, compiler_options, emit_ce_difference=False - ), + make_ce_kl_step(model, ci_capture_keys, (variant,), mesh, compiler_options), (f"ce_kl/kl_{variant}",), model, run_key, diff --git a/param_decomp/experiments/lm/test_eval_operations.py b/param_decomp/experiments/lm/test_eval_operations.py index a967e20b0..7303a5fcd 100644 --- a/param_decomp/experiments/lm/test_eval_operations.py +++ b/param_decomp/experiments/lm/test_eval_operations.py @@ -93,7 +93,6 @@ def make_operation( n_proc=1, sink=cast(Any, SimpleNamespace(accepts_deferred_media=True)), compiler_options={}, - target_pool_batches_for=None, ) batch = jnp.arange(4) _, key = captured["inputs_for_context"]( @@ -102,7 +101,6 @@ def make_operation( now_step=30, pass_index=3, batches=(batch,), - target_batches=None, ) ) diff --git a/param_decomp/experiments/lm/training.py b/param_decomp/experiments/lm/training.py index dd1b877db..9a2f0337a 100644 --- a/param_decomp/experiments/lm/training.py +++ b/param_decomp/experiments/lm/training.py @@ -163,15 +163,7 @@ def sample_batch(step: int) -> jax.Array: "mid-window eval would corrupt the next step-time estimate" ) evaluation = make_lm_evaluation( - built, - eval_config, - model, - run_key, - mesh, - n_proc, - sink, - runtime.compiler_options, - target_pool_batches_for=None, # a plain run has ONE stream + built, eval_config, model, run_key, mesh, n_proc, sink, runtime.compiler_options ) run_decomposition_training( From 02b7a675410429eb454bf98420d251ef4a34e152 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 02:58:32 +0000 Subject: [PATCH 11/15] refactor(eval): give the single-arm evals their own step, leave CE/KL alone MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The point of the single-arm evals is that a tPD run does NOT author `CEandKLLosses` — so teaching that bundle to select arms coupled this PR to the very thing it lets runs avoid, in a function plain runs share. `make_masked_kl_step` computes one arm directly: one clean forward, one masked, KL. `experiments/lm/eval.py` is back to a pure addition — `make_ce_kl_step` is byte- identical to upstream, without the `variants` parameter, the nullable `rounding_threshold`, the renamed mask dict, or the moved rounded arm. The keys are unchanged (`ce_kl/kl_ci_masked`, `ce_kl/kl_unmasked`), so a targeted run's numbers still meet a plain run's under one name. Cost is the two mask constructions being spelled in both places — two lines, against a shared function kept intact. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/train.py | 3 +- param_decomp/experiments/lm/eval.py | 92 +++++++++++-------- .../experiments/lm/eval_operations.py | 8 +- .../experiments/lm/scalar_eval_operations.py | 28 ++---- 4 files changed, 69 insertions(+), 62 deletions(-) diff --git a/param_decomp/core/train.py b/param_decomp/core/train.py index 69b419c42..c3fd3c9f0 100644 --- a/param_decomp/core/train.py +++ b/param_decomp/core/train.py @@ -1190,8 +1190,7 @@ def make_targeted_train_step[PreparedT]( coeff_schedules[f"{objective.target.imp.name}/frequency"] = ( objective.target.imp.cfg.frequency.coeff ) - # Both passes run the same imp-min term, each under its own coefficient, so both log it - # under that term's name. + imp_name = objective.target.imp.name nontarget_coeff_schedules: dict[str, LossCoeff] = { imp_name: objective.nontarget.impmin_coeff, diff --git a/param_decomp/experiments/lm/eval.py b/param_decomp/experiments/lm/eval.py index 2c16d3abe..1d395389b 100644 --- a/param_decomp/experiments/lm/eval.py +++ b/param_decomp/experiments/lm/eval.py @@ -42,7 +42,7 @@ import math from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import Literal, get_args +from typing import Literal import jax import jax.numpy as jnp @@ -205,39 +205,62 @@ def _ce[PreparedT](batch: _PreparedLMBatch[PreparedT], logits: Array) -> Array: return _row_masked_cross_entropy(logits, batch.tokens, batch.valid_row_mask) -type CEKLVariant = Literal[ - "ci_masked", "unmasked", "stoch_masked", "random_masked", "rounded_masked", "zero_masked" -] -"""One masking arm of the CE/KL evaluator. Each costs ONE masked forward, so a caller that -wants a single number asks for a single arm rather than filtering the record afterwards.""" +type MaskingArm = Literal["ci_masked", "unmasked"] +"""A masking arm authorable as an eval on its own (`CIMaskedReconLoss` / +`UnmaskedNoDeltaReconLoss`). Both pin every weight-delta mask to zero.""" -CE_KL_VARIANTS: tuple[CEKLVariant, ...] = get_args(CEKLVariant.__value__) -"""Every arm, in the order the full-fidelity `CEandKLLosses` reports them.""" + +def make_masked_kl_step[PreparedT]( + model_static: DecomposedModel[PreparedT], + ci_capture_keys: CaptureKeys, + arm: MaskingArm, + mesh: Mesh | None = None, + compiler_options: dict[str, bool | int | str] | None = None, + *, + n_valid_rows: int | None = None, +) -> ScalarStep: + """KL against the target output under ONE masking arm — one clean forward, one masked. + + `CEandKLLosses` reports the same number for this arm among ten others, at one masked + forward each; a targeted run wants a single arm, so it gets a step that computes only + that one rather than a wider evaluator narrowed after the fact. The key is the spelling + `CEandKLLosses` uses, so the two runs' numbers meet under one name.""" + assert model_static.has_position_axis, "masked KL is LM-only and requires a position axis" + + def eval_step( + model: DecomposedModel[PreparedT], + components: ComponentStacks, + ci_fn: CIFn, + token_ids: Array, + key: PRNGKeyArray, + ) -> dict[str, Array]: + del key # neither arm draws masks + batch = _prepare_lm_batch( + model, components, ci_fn, token_ids, mesh, n_valid_rows, ci_capture_keys + ) + match arm: + case "ci_masked": + masks = batch.ci_lower + case "unmasked": + masks = {site: jnp.ones_like(batch.ci_lower[site]) for site in model.site_names} + zeros_delta = {site: jnp.zeros(batch.tokens.shape, COMPUTE_DT) for site in model.site_names} + logits = _compute_masked_output(model, batch, masks, zeros_delta, mesh, frozenset()) + return {f"ce_kl/kl_{arm}": _kl(batch, logits)} + + return filter_jit(eval_step, compiler_options=compiler_options) def make_ce_kl_step[PreparedT]( model_static: DecomposedModel[PreparedT], ci_capture_keys: CaptureKeys, - variants: tuple[CEKLVariant, ...], + rounding_threshold: float, mesh: Mesh | None = None, compiler_options: dict[str, bool | int | str] | None = None, *, - rounding_threshold: float | None = None, n_valid_rows: int | None = None, ) -> ScalarStep: - """Build the single-purpose CE/KL evaluator over `variants`. - - The masks for EVERY arm are drawn whether or not the arm is selected, so a narrowed - evaluator's numbers are bit-identical to the full one's — only the forwards are - skipped, and XLA drops the unused draws. `emit_ce_difference` adds the CE-vs-target - `rounding_threshold` belongs to the - rounded arm alone, so it is present exactly when that arm is.""" + """Build the single-purpose CE/KL evaluator.""" assert model_static.has_position_axis, "CEandKLLosses is LM-only and requires a position axis" - assert variants, "a CE/KL evaluator with no arms measures nothing" - assert ("rounded_masked" in variants) == (rounding_threshold is not None), ( - variants, - rounding_threshold, - ) def eval_step( model: DecomposedModel[PreparedT], @@ -262,7 +285,7 @@ def eval_step( batch.tokens.shape, COMPUTE_DT, ) - variant_masks = { + variants = { "ci_masked": (batch.ci_lower, zeros_delta), "unmasked": ( {site: jnp.ones_like(batch.ci_lower[site]) for site in model.site_names}, @@ -278,23 +301,21 @@ def eval_step( }, zeros_delta, ), - "zero_masked": ( - {site: jnp.zeros_like(batch.ci_lower[site]) for site in model.site_names}, - zeros_delta, - ), - } - if rounding_threshold is not None: - variant_masks["rounded_masked"] = ( + "rounded_masked": ( { site: (batch.ci_lower[site] > rounding_threshold).astype(COMPUTE_DT) for site in model.site_names }, zeros_delta, - ) + ), + "zero_masked": ( + {site: jnp.zeros_like(batch.ci_lower[site]) for site in model.site_names}, + zeros_delta, + ), + } variant_logits = { name: _compute_masked_output(model, batch, masks, deltas, mesh, frozenset()) - for name, (masks, deltas) in variant_masks.items() - if name in variants + for name, (masks, deltas) in variants.items() } target_ce = _ce(batch, batch.clean.output) metrics = { @@ -303,7 +324,7 @@ def eval_step( metrics.update( { f"ce_kl/ce_difference_{name}": _ce(batch, variant_logits[name]) - target_ce - for name in variant_logits + for name in variants if name != "zero_masked" } ) @@ -459,10 +480,9 @@ def make_eval_step[PreparedT]( ce_kl = make_ce_kl_step( model_static, ci_capture_keys, - CE_KL_VARIANTS, + rounding_threshold, mesh, compiler_options, - rounding_threshold=rounding_threshold, n_valid_rows=n_valid_rows, ) ci_l0 = make_ci_l0_step( diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index 333ebc903..fd4866436 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -59,7 +59,7 @@ make_ce_kl_operation, make_ci_l0_operation, make_fresh_pgd_operation, - make_single_variant_kl_operation, + make_masked_kl_operation, stream_batches, ) from param_decomp.infra.dataset_store import read_dataset_meta @@ -152,14 +152,12 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo case CEandKLLossesConfig(): return per_stream(make_ce_kl_operation, metric, schedule, data_streams) case CIMaskedReconLossConfig(): - return per_stream( - make_single_variant_kl_operation, "ci_masked", schedule, data_streams - ) + return per_stream(make_masked_kl_operation, "ci_masked", schedule, data_streams) case UnmaskedNoDeltaReconLossConfig(): # The non-target pass's OWN training term, already reported as a train loss # there; measuring it again off-target would restate the objective. return per_stream( - make_single_variant_kl_operation, "unmasked", schedule, (optimized_stream,) + make_masked_kl_operation, "unmasked", schedule, (optimized_stream,) ) case CI_L0Config(): return per_stream(make_ci_l0_operation, metric, schedule, data_streams) diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index edbbc4083..32bf4f890 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -15,12 +15,12 @@ from param_decomp.core.recon_eval import FreshPGDReconEval from param_decomp.core.run import EvalOperation from param_decomp.experiments.lm.eval import ( - CE_KL_VARIANTS, - CEKLVariant, + MaskingArm, ScalarStep, make_ce_kl_step, make_ci_l0_step, make_fresh_pgd_step, + make_masked_kl_step, ) from param_decomp.experiments.lm.eval_config import CEandKLLossesConfig from param_decomp.experiments.lm.eval_context import LMEvalContext @@ -109,14 +109,7 @@ def make_ce_kl_operation( ) -> EvalOperation[LMEvalContext]: return _make_scalar_operation( schedule, - make_ce_kl_step( - model, - ci_capture_keys, - CE_KL_VARIANTS, - mesh, - compiler_options, - rounding_threshold=metric.rounding_threshold, - ), + make_ce_kl_step(model, ci_capture_keys, metric.rounding_threshold, mesh, compiler_options), ("ce_kl/",), model, run_key, @@ -126,8 +119,8 @@ def make_ce_kl_operation( ) -def make_single_variant_kl_operation( - variant: CEKLVariant, +def make_masked_kl_operation( + arm: MaskingArm, schedule: EvalSchedule, stream: Stream, model: DecomposedModel, @@ -138,17 +131,14 @@ def make_single_variant_kl_operation( mesh: Mesh, compiler_options: dict[str, bool | int | str], ) -> EvalOperation[LMEvalContext]: - """ONE masking arm of the CE/KL evaluator, authored as the loss config that names the - same construction. + """ONE masking arm, authored as the loss config that names the same construction. `CIMaskedReconLoss` / `UnmaskedNoDeltaReconLoss` are authorable as evals exactly as - `PGDReconLoss` already is — the eval measures the quantity the loss optimizes. The key - keeps the `ce_kl/kl_` spelling a plain run's `CEandKLLosses` logs it under, so the - same number is one name across run kinds; the config type names the construction.""" + `PGDReconLoss` already is — the eval measures the quantity the loss optimizes.""" return _make_scalar_operation( schedule, - make_ce_kl_step(model, ci_capture_keys, (variant,), mesh, compiler_options), - (f"ce_kl/kl_{variant}",), + make_masked_kl_step(model, ci_capture_keys, arm, mesh, compiler_options), + (f"ce_kl/kl_{arm}",), model, run_key, train_steps, From b330767343e83c35c22f0e938aa54efe131d097e Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 03:07:34 +0000 Subject: [PATCH 12/15] docs: cut the commentary to what the code cannot say itself Every comment and docstring this PR adds, held to one test: does a reader need it to understand the code as it stands? Removed: `stream_log_prefix`'s statement of the namespace rule (the four-line body is the rule); the `data_streams` / `optimized_stream` comments (the names say it); the bind-time-vs-runtime rationale on the CE/KL collision assert (the assert message states the rule); the `fill_between`-not-`bar` note and the one-permutation-per-site note in the plot helpers. Replaced: `accumulate_site_reductions(..., 0)` plus a comment explaining the 0 is now `n_batches_accum=0`. Trimmed to their load-bearing sentence: the two plot docstrings (x is a rank, not a component id; the nontarget series takes the target's permutation), the raw- sample helper, `weight_magnitudes` (the norms reduce on device), `MaskingArm`, `make_masked_kl_step`, `LMEvalContext.target_batches`, `make_lm_evaluation`, and the target-stream sampler. One comment survives: why `UnmaskedNoDeltaReconLoss` binds to one stream. That is a choice the code cannot show. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/core/configs.py | 4 +--- param_decomp/core/slow_eval.py | 22 ++++++------------- .../lm/diagnostic_eval_operations.py | 6 ++--- param_decomp/experiments/lm/eval.py | 10 ++++----- param_decomp/experiments/lm/eval_config.py | 5 ++--- param_decomp/experiments/lm/eval_context.py | 7 +++--- .../experiments/lm/eval_operations.py | 19 ++++------------ .../experiments/lm/scalar_eval_operations.py | 19 +++------------- .../experiments/lm/training_targeted.py | 5 ++--- 9 files changed, 29 insertions(+), 68 deletions(-) diff --git a/param_decomp/core/configs.py b/param_decomp/core/configs.py index 0c2e845e7..5e0980c93 100644 --- a/param_decomp/core/configs.py +++ b/param_decomp/core/configs.py @@ -471,9 +471,7 @@ class ComponentActivationDensityConfig(BaseConfig): class WeightMagnitudeConfig(BaseConfig): """Per-site `‖V_c‖·‖U_c‖` scatter in descending magnitude order, log y. - Reads the trained V/U alone — no forward pass and no eval batch, so it costs two norm - reductions and a plot. The norms reduce on device; only `C` floats per site are pulled. - """ + Reads the trained V/U alone — no forward pass and no eval batch.""" slow: ClassVar[bool] = True type: Literal["WeightMagnitude"] = "WeightMagnitude" diff --git a/param_decomp/core/slow_eval.py b/param_decomp/core/slow_eval.py index f7e49c91c..d68499f22 100644 --- a/param_decomp/core/slow_eval.py +++ b/param_decomp/core/slow_eval.py @@ -200,9 +200,7 @@ def slow_eval_step( def _raw_sample(chunks: dict[str, list[np.ndarray]], site: str) -> np.ndarray: - """A site's kept raw values, or empty when none were kept: `n_batches_accum=0` asks for - NO sample, and a caller that reads only `ci_sums` (the two-stream CI mean) would - otherwise pay a host gather of every position's CI just to discard it.""" + """A site's kept raw values, or empty when `n_batches_accum` kept none.""" return np.concatenate(chunks[site]) if site in chunks else np.empty(0, np.float32) @@ -603,15 +601,14 @@ def weight_magnitudes(components: ComponentStacks) -> dict[str, np.ndarray]: def mean_cis(reductions: dict[str, SiteReduction]) -> dict[str, np.ndarray]: - """Per-site token-weighted mean CI, with the zero-position guard in ONE place.""" + """Per-site token-weighted mean CI.""" assert all(r.n_positions > 0 for r in reductions.values()) return {site: r.ci_sums / r.n_positions for site, r in reductions.items()} def plot_weight_magnitudes(magnitudes: dict[str, np.ndarray]) -> bytes: - """Per-site `‖V_c‖·‖U_c‖` in DESCENDING magnitude order, log y. x is a component's rank - within its site, NOT its component id: the readable quantity is the spectrum's shape — - how fast magnitude falls off, and where it knees — which component order destroys.""" + """Per-site `‖V_c‖·‖U_c‖` in descending magnitude order, log y. x is a component's rank + within its site, NOT its component id.""" n_rows, n_cols = _grid_dims(len(magnitudes)) fig = Figure(figsize=(8 * n_cols, 3 * n_rows)) axs = fig.subplots(n_rows, n_cols, squeeze=False) @@ -632,17 +629,15 @@ def plot_mean_component_cis_two_streams( target_mean_cis: dict[str, np.ndarray], nontarget_mean_cis: dict[str, np.ndarray], ) -> tuple[bytes, bytes]: - """Both streams' mean CI on one axis per site, ordered by DESCENDING TARGET mean. + """Both streams' mean CI on one axis per site, ordered by descending TARGET mean. - The non-target series is reordered by the same permutation rather than sorted on its - own, so a component's two bars line up vertically — the whole point is reading, per - component, how much on-target importance comes with off-target importance.""" + The nontarget series takes the same permutation rather than its own, so a component's + two series line up vertically.""" assert target_mean_cis.keys() == nontarget_mean_cis.keys(), ( sorted(target_mean_cis), sorted(nontarget_mean_cis), ) n_rows, n_cols = _grid_dims(len(target_mean_cis)) - # One permutation per site, not one per scale: the log figure plots the same ordering. ordered = { name: (target[order], nontarget_mean_cis[name][order]) for name, target in target_mean_cis.items() @@ -656,9 +651,6 @@ def plot_mean_component_cis_two_streams( for ax in flat_axes[len(ordered) :]: ax.set_visible(False) for ax, (name, (target, nontarget)) in zip(flat_axes, ordered.items(), strict=False): - # `fill_between`, not `bar`: at production C a bar per component is one matplotlib - # patch each, ~25x the render time of the equivalent filled step for the same - # picture — and this renders on a thread that contends with the train loop. x = np.arange(len(target)) if log_y: ax.set_yscale("log") diff --git a/param_decomp/experiments/lm/diagnostic_eval_operations.py b/param_decomp/experiments/lm/diagnostic_eval_operations.py index c32bd3a49..b4e794a75 100644 --- a/param_decomp/experiments/lm/diagnostic_eval_operations.py +++ b/param_decomp/experiments/lm/diagnostic_eval_operations.py @@ -222,9 +222,9 @@ def make_two_stream_ci_mean_operation( step = make_slow_eval_step(model, ci_capture_keys, 0.0, None, compiler_options) def stream_mean_cis(ci_fn: CIFn, batches: tuple[jax.Array, ...]) -> dict[str, np.ndarray]: - # `n_batches_accum=0`: this metric reads only `ci_sums`, and the raw-value sample - # would gather every position's CI to the host (~430MB/pass here) to be discarded. - return mean_cis(accumulate_site_reductions(step, model, ci_fn, list(batches), 0)) + return mean_cis( + accumulate_site_reductions(step, model, ci_fn, list(batches), n_batches_accum=0) + ) def run(context: LMEvalContext) -> LogRecord: ci_fn = context.state.decomposition.ci_fn diff --git a/param_decomp/experiments/lm/eval.py b/param_decomp/experiments/lm/eval.py index 1d395389b..0ea42786c 100644 --- a/param_decomp/experiments/lm/eval.py +++ b/param_decomp/experiments/lm/eval.py @@ -206,8 +206,8 @@ def _ce[PreparedT](batch: _PreparedLMBatch[PreparedT], logits: Array) -> Array: type MaskingArm = Literal["ci_masked", "unmasked"] -"""A masking arm authorable as an eval on its own (`CIMaskedReconLoss` / -`UnmaskedNoDeltaReconLoss`). Both pin every weight-delta mask to zero.""" +"""A masking arm authorable as an eval on its own. Both pin every weight-delta mask to +zero.""" def make_masked_kl_step[PreparedT]( @@ -221,10 +221,8 @@ def make_masked_kl_step[PreparedT]( ) -> ScalarStep: """KL against the target output under ONE masking arm — one clean forward, one masked. - `CEandKLLosses` reports the same number for this arm among ten others, at one masked - forward each; a targeted run wants a single arm, so it gets a step that computes only - that one rather than a wider evaluator narrowed after the fact. The key is the spelling - `CEandKLLosses` uses, so the two runs' numbers meet under one name.""" + The key is the spelling `CEandKLLosses` reports this arm under, so the two are one + quantity under one name.""" assert model_static.has_position_axis, "masked KL is LM-only and requires a position axis" def eval_step( diff --git a/param_decomp/experiments/lm/eval_config.py b/param_decomp/experiments/lm/eval_config.py index 97369b3f2..09722b61a 100644 --- a/param_decomp/experiments/lm/eval_config.py +++ b/param_decomp/experiments/lm/eval_config.py @@ -68,9 +68,8 @@ class TwoStreamCIMeanPerComponentConfig(BaseConfig): """Both streams' mean CI per component on ONE axis per site, ordered by descending TARGET mean and coloured by stream. - Supersedes authoring `CIMeanPerComponent` on a targeted run: it computes the same - nontarget-stream reduction plus the target-stream one, so authoring both would pay - for the nontarget pass twice. Refuses on a plain run, which has no target stream.""" + Computes `CIMeanPerComponent`'s reduction on both streams, so authoring both pays for + the nontarget pass twice. Refuses on a plain run, which has no target stream.""" slow: ClassVar[bool] = True type: Literal["TwoStreamCIMeanPerComponent"] = "TwoStreamCIMeanPerComponent" diff --git a/param_decomp/experiments/lm/eval_context.py b/param_decomp/experiments/lm/eval_context.py index 72d06531c..e521eb79f 100644 --- a/param_decomp/experiments/lm/eval_context.py +++ b/param_decomp/experiments/lm/eval_context.py @@ -12,7 +12,6 @@ class LMEvalContext(EvalInvocation): pass_index: int batches: tuple[jax.Array, ...] target_batches: tuple[jax.Array, ...] | None = None - """A tPD run's prompt-pool draws, which `data.eval` cannot supply — the pool has no - held-out split, so the targeted root draws them exactly as training does. `None` on a - plain run, which HAS no second stream; that `None` is also what tells every log key - which run kind it is in (`scalar_eval_operations.stream_log_prefix`).""" + """A tPD run's target-stream draws; `None` on a plain run, which has no second stream. + That `None` is also what tells every log key which run kind it is in + (`scalar_eval_operations.stream_log_prefix`).""" diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index fd4866436..75e72db99 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -86,10 +86,7 @@ def make_lm_evaluation( the metric measures. `target_pool_batches_for(pass_index)` supplies the tPD target stream, which `data.eval` - cannot: the prompt pool has no held-out split, so the targeted root draws pool batches - the same way training does. `None` on a plain run, and that is what makes a plain run's - metric set — and every one of its log keys — exactly what it was before targeted runs - existed: `data_streams` collapses to the single nontarget stream.""" + cannot. `None` on a plain run, which collapses every metric to the one stream it has.""" pd = built.pd capture_inputs = built.ci_fn.capture_keys data = built.data @@ -108,11 +105,7 @@ def batches(pass_index: int) -> list[jax.Array]: ] targeted = target_pool_batches_for is not None - # Every stream a metric authored for BOTH measures; a plain run has exactly one. data_streams: tuple[Stream, ...] = ("nontarget", "target") if targeted else ("nontarget",) - # The stream the run optimizes for, and the DEFAULT for any single-stream eval: a - # diagnostic you read once should describe the data you are interpreting. On a plain run - # this IS the nontarget stream, so every such metric keeps the keys and data it had. optimized_stream: Stream = "target" if targeted else "nontarget" def well_temperedness_inputs( @@ -128,8 +121,8 @@ def per_stream( schedule: EvalSchedule, streams: tuple[Stream, ...], ) -> tuple[EvalOperation[LMEvalContext], ...]: - """One operation per stream. The scalar makers share a signature exactly so the - stream can be the only thing that varies between a metric's two readouts.""" + """One operation per stream; the scalar makers share a signature so the stream is the + only thing that varies between a metric's readouts.""" return tuple( maker( metric, @@ -154,8 +147,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo case CIMaskedReconLossConfig(): return per_stream(make_masked_kl_operation, "ci_masked", schedule, data_streams) case UnmaskedNoDeltaReconLossConfig(): - # The non-target pass's OWN training term, already reported as a train loss - # there; measuring it again off-target would restate the objective. + # The nontarget pass's own training term, already reported there as a loss. return per_stream( make_masked_kl_operation, "unmasked", schedule, (optimized_stream,) ) @@ -271,9 +263,6 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo "TwoStreamCIMeanPerComponent already computes the nontarget-stream reduction " "CIMeanPerComponent does, so authoring both pays for that pass twice" ) - # The single-arm evals ARE arms of `CEandKLLosses`, under the same keys. Authoring both - # is a duplicate measurement that `_run_due_evaluation`'s collision assert would catch - # at the first eval pass — hours into a run. Catch it while reading the config. single_arm = tuple( metric.type for metric in eval.metrics diff --git a/param_decomp/experiments/lm/scalar_eval_operations.py b/param_decomp/experiments/lm/scalar_eval_operations.py index 32bf4f890..9cb3ae320 100644 --- a/param_decomp/experiments/lm/scalar_eval_operations.py +++ b/param_decomp/experiments/lm/scalar_eval_operations.py @@ -28,8 +28,7 @@ type Stream = Literal["nontarget", "target"] """Which STREAM an eval operation measures. ONE value, not a (batch source, log prefix) -pair: the two always covary, and nothing should be able to spell target-stream batches -under the nontarget stream's log keys.""" +pair, so target-stream batches cannot be spelled under the nontarget stream's log keys.""" def stream_batches(stream: Stream, context: LMEvalContext) -> tuple[Array, ...]: @@ -44,16 +43,6 @@ def stream_batches(stream: Stream, context: LMEvalContext) -> tuple[Array, ...]: def stream_log_prefix(stream: Stream, context: LMEvalContext) -> str: - """The log namespace for `stream`, given what kind of run this is. - - ONE rule: the stream the run is optimizing for is unlabelled, and anything else carries - its name. A plain run has a single stream, so it stays `eval/` exactly as before — - `context.target_batches is None` is what says so. A tPD run adds a second, so its - nontarget stream moves under `eval/nontarget_data/` and its target stream takes the bare - namespace. The consequence is deliberate: `eval/l0/...` means "the stream of interest" in - both run kinds, which is what makes the two comparable as objectives — but it is NOT the - same data, so comparing like for like across run kinds must read `eval/nontarget_data/` - on the tPD side.""" targeted = context.target_batches is not None match stream: case "nontarget": @@ -131,10 +120,8 @@ def make_masked_kl_operation( mesh: Mesh, compiler_options: dict[str, bool | int | str], ) -> EvalOperation[LMEvalContext]: - """ONE masking arm, authored as the loss config that names the same construction. - - `CIMaskedReconLoss` / `UnmaskedNoDeltaReconLoss` are authorable as evals exactly as - `PGDReconLoss` already is — the eval measures the quantity the loss optimizes.""" + """ONE masking arm, authored as the loss config that names the same construction — + `CIMaskedReconLoss` / `UnmaskedNoDeltaReconLoss`, as `PGDReconLoss` already is.""" return _make_scalar_operation( schedule, make_masked_kl_step(model, ci_capture_keys, arm, mesh, compiler_options), diff --git a/param_decomp/experiments/lm/training_targeted.py b/param_decomp/experiments/lm/training_targeted.py index 714671a2e..8175f85bb 100644 --- a/param_decomp/experiments/lm/training_targeted.py +++ b/param_decomp/experiments/lm/training_targeted.py @@ -142,9 +142,8 @@ def sample_nontarget_batch(step: int) -> jax.Array: eval_target_batch = eval_config.batch_size def eval_target_pool_batches(pass_index: int) -> list[jax.Array]: - """The eval pass's TARGET stream: the same pure `(seed, step)` pool sampler - training uses, on the `seed + 1` stream the nontarget eval split already draws - from — so an eval never scores the exact rows the step just trained on.""" + """The eval pass's target stream: training's pure `(seed, step)` pool sampler on + the `seed + 1` stream, so an eval never scores the rows the step just trained.""" n_batches = eval_config.n_steps return [ pool_global_batch(built.pd.seed + 1, pass_index * n_batches + j, eval_target_batch) From 552d6169391ea35557d96d5177ab7f913da6e865 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 03:11:10 +0000 Subject: [PATCH 13/15] refactor(eval): compact the bind-time assertions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three checks, twenty lines, one of which was not a correctness check: authoring `CIMeanPerComponent` alongside `TwoStreamCIMeanPerComponent` computes a reduction twice but produces correct numbers, so it is a cost, not an error. Dropped. The two that remain — a target-stream metric on a plain run, and a single-arm KL eval alongside `CEandKLLosses` under the same keys — read off one `authored` set instead of building tuples to interpolate into their own messages. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- .../experiments/lm/eval_operations.py | 26 +++++-------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index 75e72db99..d93b1c6a5 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -250,28 +250,14 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo case WeightMagnitudeConfig(): return (make_weight_magnitude_operation(schedule, renderer),) - needs_target_stream = tuple( - metric.type - for metric in eval.metrics - if isinstance(metric, TwoStreamCIMeanPerComponentConfig) - ) - assert not needs_target_stream or targeted, ( - f"{needs_target_stream} measure the tPD target stream; a plain run has no prompt pool" - ) authored = {type(metric) for metric in eval.metrics} - assert not {TwoStreamCIMeanPerComponentConfig, CIMeanPerComponentConfig} <= authored, ( - "TwoStreamCIMeanPerComponent already computes the nontarget-stream reduction " - "CIMeanPerComponent does, so authoring both pays for that pass twice" - ) - single_arm = tuple( - metric.type - for metric in eval.metrics - if isinstance(metric, CIMaskedReconLossConfig | UnmaskedNoDeltaReconLossConfig) - ) - assert not (single_arm and any(isinstance(m, CEandKLLossesConfig) for m in eval.metrics)), ( - f"{single_arm} emit arms CEandKLLosses already emits, under the same `ce_kl/kl_*` " - "keys: author the narrow metrics OR CEandKLLosses, not both" + assert targeted or TwoStreamCIMeanPerComponentConfig not in authored, ( + "TwoStreamCIMeanPerComponent measures the target stream; a plain run has none" ) + assert not ( + authored & {CIMaskedReconLossConfig, UnmaskedNoDeltaReconLossConfig} + and CEandKLLossesConfig in authored + ), "the single-arm KL evals emit keys CEandKLLosses also emits; author one or the other" operations = tuple( operation for metric in eval.metrics for operation in make_operations(metric) ) From bf37f57325659bf4170419317772ad2dd0e38258 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 03:35:25 +0000 Subject: [PATCH 14/15] docs(eval): state the plain-PD behaviour, tPD in parentheses `make_lm_evaluation`'s docstring leads with what it does for any run and puts the targeted case in a parenthesis. `per_stream`'s docstring and the comment on the single-stream binding of `UnmaskedNoDeltaReconLoss` are gone: both explained why the code is shaped the way it is rather than what it does. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- param_decomp/experiments/lm/eval_operations.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index d93b1c6a5..a1d63f52c 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -82,11 +82,9 @@ def make_lm_evaluation( compiler_options: dict[str, bool | int | str], target_pool_batches_for: Callable[[int], list[jax.Array]] | None = None, ) -> Evaluation[LMEvalContext]: - """Construct the executable operations for every authored LM metric — one PER STREAM - the metric measures. - - `target_pool_batches_for(pass_index)` supplies the tPD target stream, which `data.eval` - cannot. `None` on a plain run, which collapses every metric to the one stream it has.""" + """Construct one executable operation for every authored LM metric (for tPD, one per + stream the metric measures; `target_pool_batches_for` draws the target stream, and + `None` marks a plain run).""" pd = built.pd capture_inputs = built.ci_fn.capture_keys data = built.data @@ -121,8 +119,6 @@ def per_stream( schedule: EvalSchedule, streams: tuple[Stream, ...], ) -> tuple[EvalOperation[LMEvalContext], ...]: - """One operation per stream; the scalar makers share a signature so the stream is the - only thing that varies between a metric's readouts.""" return tuple( maker( metric, @@ -147,7 +143,6 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo case CIMaskedReconLossConfig(): return per_stream(make_masked_kl_operation, "ci_masked", schedule, data_streams) case UnmaskedNoDeltaReconLossConfig(): - # The nontarget pass's own training term, already reported there as a loss. return per_stream( make_masked_kl_operation, "unmasked", schedule, (optimized_stream,) ) From c9259a010ec1467ba8e2661644799513fa3ba427 Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 15 Aug 2026 03:37:18 +0000 Subject: [PATCH 15/15] rename(eval): primary_stream, not optimized_stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both tPD streams are optimized — the nontarget pass has its own imp-min and recon terms — so "optimized" named nothing that distinguishes them. The variable holds the stream a single-stream eval measures: the only stream on a plain run, the target stream on tPD. `data_streams` becomes `all_streams`, which is what it is. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ --- .../experiments/lm/eval_operations.py | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/param_decomp/experiments/lm/eval_operations.py b/param_decomp/experiments/lm/eval_operations.py index a1d63f52c..1c8111b3c 100644 --- a/param_decomp/experiments/lm/eval_operations.py +++ b/param_decomp/experiments/lm/eval_operations.py @@ -103,13 +103,13 @@ def batches(pass_index: int) -> list[jax.Array]: ] targeted = target_pool_batches_for is not None - data_streams: tuple[Stream, ...] = ("nontarget", "target") if targeted else ("nontarget",) - optimized_stream: Stream = "target" if targeted else "nontarget" + all_streams: tuple[Stream, ...] = ("nontarget", "target") if targeted else ("nontarget",) + primary_stream: Stream = "target" if targeted else "nontarget" def well_temperedness_inputs( context: LMEvalContext, ) -> tuple[jax.Array, PRNGKeyArray]: - return stream_batches(optimized_stream, context)[0], jax.random.fold_in( + return stream_batches(primary_stream, context)[0], jax.random.fold_in( run_key, EvalKeyStream.WELL_TEMPEREDNESS * pd.steps + context.pass_index ) @@ -139,17 +139,15 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo schedule = schedule_for(metric, eval) match metric: case CEandKLLossesConfig(): - return per_stream(make_ce_kl_operation, metric, schedule, data_streams) + return per_stream(make_ce_kl_operation, metric, schedule, all_streams) case CIMaskedReconLossConfig(): - return per_stream(make_masked_kl_operation, "ci_masked", schedule, data_streams) + return per_stream(make_masked_kl_operation, "ci_masked", schedule, all_streams) case UnmaskedNoDeltaReconLossConfig(): - return per_stream( - make_masked_kl_operation, "unmasked", schedule, (optimized_stream,) - ) + return per_stream(make_masked_kl_operation, "unmasked", schedule, (primary_stream,)) case CI_L0Config(): - return per_stream(make_ci_l0_operation, metric, schedule, data_streams) + return per_stream(make_ci_l0_operation, metric, schedule, all_streams) case PGDReconLossConfig(): - return per_stream(make_fresh_pgd_operation, metric, schedule, data_streams) + return per_stream(make_fresh_pgd_operation, metric, schedule, all_streams) case CIMaskedAttnPatternsReconLossConfig() | StochasticAttnPatternsReconLossConfig(): return ( @@ -161,7 +159,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo run_key, pd.steps, compiler_options, - optimized_stream, + primary_stream, ), ) case CIHiddenActsReconLossConfig() | StochasticHiddenActsReconLossConfig(): @@ -174,7 +172,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo run_key, pd.steps, compiler_options, - optimized_stream, + primary_stream, ), ) case ( @@ -190,7 +188,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo capture_inputs, compiler_options, renderer, - optimized_stream, + primary_stream, ), ) case PermutedCIPlotsConfig() | UVPlotsConfig() | IdentityCIErrorConfig(): @@ -202,7 +200,7 @@ def make_operations(metric: AnyEvalMetricConfig) -> tuple[EvalOperation[LMEvalCo capture_inputs, compiler_options, renderer, - optimized_stream, + primary_stream, ), )