Consistent tPD metric names, and core evals on both streams - #999
Consistent tPD metric names, and core evals on both streams#999Antovigo wants to merge 15 commits into
Conversation
The rule: the data a run OPTIMIZES FOR is unlabelled. Keys become `<tier>/[<stream>/]<family>/<name>`, 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_<arm>` 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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
`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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
`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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
`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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
… alone 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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
`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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JFBxj8yGpCoGZ6KXNP6iaQ
Antovigo
left a comment
There was a problem hiding this comment.
Inline notes on the changes that aren't self-evident from the diff.
— [written by Claude]
| k, f"train/{k}" if k.startswith(("grad_norms/", "loss/", "schedules/")) else k | ||
| k, | ||
| f"train/{k}" | ||
| if k.startswith(("grad_norms/", "loss/", "schedules/", "nontarget_data/")) |
There was a problem hiding this comment.
The renamed non-target keys start with nontarget_data/ rather than loss/, so without this entry they'd fall through to the else arm and be logged with no train/ tier at all.
— [written by Claude]
| ci_alive_threshold: float = 0.0 | ||
|
|
||
|
|
||
| class WeightMagnitudeConfig(BaseConfig): |
There was a problem hiding this comment.
New plot. slow: ClassVar[bool] is required of every eval metric — assert_every_metric_declares_its_tier enforces that each class declares its own tier rather than inheriting one.
— [written by Claude]
| return filter_jit(slow_eval_step, compiler_options=compiler_options) | ||
|
|
||
|
|
||
| def _raw_sample(chunks: dict[str, list[np.ndarray]], site: str) -> np.ndarray: |
There was a problem hiding this comment.
Bug fix, not a refactor. n_batches_accum=0 means 'keep no raw sample', but the constructor used to np.concatenate(chunks[site]) unconditionally — with no chunks that's a KeyError. TwoStreamCIMeanPerComponent passes 0, so it crashed on every invocation before this PR and has never actually run.
— [written by Claude]
| x = np.arange(len(target)) | ||
| if log_y: | ||
| ax.set_yscale("log") | ||
| ax.fill_between(x, target, step="mid", color="#1f77b4", label="target") |
There was a problem hiding this comment.
ax.bar here creates one matplotlib patch per component: measured ~120s per figure at C=1456 over 32 sites, and it renders two figures on the background thread that competes with the train loop for the GIL. fill_between(step='mid') draws the same picture in ~23s for both.
— [written by Claude]
| 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) |
There was a problem hiding this comment.
Sorted descending, so x is a component's rank within its site, not its component id — you can no longer cross-reference a point against other per-component plots.
— [written by Claude]
| from param_decomp.experiments.lm.eval_context import LMEvalContext | ||
| from param_decomp.experiments.lm.eval_keys import EvalKeyStream | ||
|
|
||
| type Stream = Literal["nontarget", "target"] |
There was a problem hiding this comment.
One value, not a (batches, log prefix) pair — stream_batches and stream_log_prefix both take it, so an operation cannot read one stream's batches while writing the other's keys.
— [written by Claude]
| class LMEvalContext(EvalInvocation): | ||
| pass_index: int | ||
| batches: tuple[jax.Array, ...] | ||
| target_batches: tuple[jax.Array, ...] | None = None |
There was a problem hiding this comment.
Defaults to None so the plain LM root and its tests need no change at all — a plain run has one stream and says so by omission. That None is also what every log key reads to tell which run kind it's in.
— [written by Claude]
| ) | ||
|
|
||
| def make_operation(metric: AnyEvalMetricConfig) -> EvalOperation[LMEvalContext]: | ||
| def per_stream( |
There was a problem hiding this comment.
make_operations returns a tuple because one authored metric now yields one operation per stream. all_streams is both on tPD and just the one on a plain run; primary_stream is where single-stream evals bind (target on tPD, the only stream otherwise).
— [written by Claude]
| assert not ( | ||
| authored & {CIMaskedReconLossConfig, UnmaskedNoDeltaReconLossConfig} | ||
| and CEandKLLossesConfig in authored | ||
| ), "the single-arm KL evals emit keys CEandKLLosses also emits; author one or the other" |
There was a problem hiding this comment.
Bind-time rather than at the first eval pass: the operations would otherwise collide on identical ce_kl/kl_* keys and trip _run_due_evaluation's assert hours into a run.
— [written by Claude]
| ) | ||
| eval_target_batch = eval_config.batch_size | ||
|
|
||
| def eval_target_pool_batches(pass_index: int) -> list[jax.Array]: |
There was a problem hiding this comment.
The eval pass needs target-stream batches that data.eval can't supply. Uses training's own pure (seed, step) pool sampler on the seed + 1 stream — the same stream the corpus eval split draws from — so an eval never scores the rows the step just trained on.
— [written by Claude]
Description
The keys are now
<tier>/[<stream>/]<family>/<name>, where:nontarget_datafor the nontarget stream in tPD, omitted otherwiseMotivation
In the current implementation of tPD, all standard evals run on the nontarget data stream. This isn’t explicitly labeled, and there’s no built-in way to run these standard evals on the target data instead. This PRs adds them, so any of the standard evals can be easily enabled or disabled on either of the streams.
There were also inconsistencies in the names of the metrics:
loss/on the train side (train/loss/nontarget/<term>), which only works because train has a single family. Eval has several (ce_kl/,l0/,loss/), so the segment has to come before them — both now put it directly after the tier.record key (
imp_smooth_l0) while the target stream used the class name (SmoothL0ImportanceMinimalityLoss). Both streams now use the class name.IdentityCIError,WellTemperedness, the site figures (CIHistograms,ComponentActivationDensity,CIMeanPerComponent) and the permutation plots (PermutedCIPlots,UVPlots) now run on target data by default.WeightMagnitudeplot, showing the components’ weight magnitudes ||U||*||V||, sorted by descending magnitude.What the metric names look like afterwards
On a plain (non-targeted) run, nothing changes at all — every key is what it is today.
The examples below are all from a targeted run:
train/loss/nontarget/imp_smooth_l0train/nontarget_data/loss/SmoothL0ImportanceMinimalityLosstrain/schedules/coeff/nontarget/impmintrain/nontarget_data/schedules/coeff/SmoothL0ImportanceMinimalityLosstrain/loss/nontarget/totaltrain/nontarget_data/loss/totaltrain/loss/total,train/loss/FaithfulnessLosseval/l0/0.1_<site>eval/nontarget_data/l0/0.1_<site>eval/l0/0.1_<site>eval/loss/PGDReconLosseval/ce_kl/kl_ci_maskedeval/slow/…,slow_eval/figures/…over the nontarget dataCIMaskedReconLossandUnmaskedNoDeltaReconLossbecome authorable undereval.metrics, likePGDReconLossConfig.Authoring a narrow metric and
CEandKLLossesis refused, because they would emit the same keys.TwoStreamCIMeanPerComponentis likeCIMeanPerComponent, except it shows both streams at once in two different colors.How Has This Been Tested?
basedpyrightclean;ruffclean (both also enforced by the pre-commit hook).Full non-slow suite over
core/tests,targets/tests,experiments,tests:1229 passed, 14 skipped, 11 xfailed.
Targeted TMS suite updated to the new train keys and passing.
test_eval_tier.py's authored-metric tier census updated for the four new eval-unionmembers — it caught them, which is what it is for.
The 4-simulated-device pass (
XLA_FLAGS=--xla_force_host_platform_device_count=4) overcore/tests+targets/tests, ascore/CLAUDE.mdrequires: 593 passed, 12 skipped,11 xfailed.
Does this PR introduce a breaking change?
Not for plain runs — a plain run has one stream, so it never emits a stream segment and
every key is what it was.
For targeted runs, yes, in two ways. Every
nontarget/key is renamed, so dashboards andsaved W&B views over tPD runs need updating and old/new runs will not overlay on one panel
for those series. And some evals that ran on nontarget data by default now run on target data.