feat(arith-eval): in-loop arithmetic CI + activation grid eval for the L18 MLP#908
Merged
Conversation
An in-loop eval that, on a fixed a x b modular-arithmetic prompt grid, renders per-component causal-importance and activation (x@V) heatmaps at the "=" answer position, plus n_alive / recon / L0 / PGD on the arithmetic probe. - prestage_arithmetic.py: offline generator of the a+b= grid artifact (one prompt per row, fixed length, single-token answers; carries (a, b) + answer_id). - arithmetic_eval.py: one fused CI + x@V grid step, active-set selection computed once (max CI > threshold), n_alive scalars, and faceted top-k heatmap render. - llama8b masked_component_activations seam (x@V; extends the S31 collect path). - config plumbing (ArithmeticCIGridConfig -> ArithmeticEvalConfig) + the ArithmeticGridRenderer rank-0 background renderer, wired via _make_arithmetic_eval into _make_lm_eval_fn on the slow-eval cadence. - llama8b_l18_arith_add.yaml run config; tests for the eval, the llama8b seam, and the lab probe-loading/sharding glue. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The faceted contact sheets had no layout manager, so packed rows let each panel's title (e.g. c13) collide with the x-tick labels of the row above. Use constrained_layout (as slow_eval._plot_ci_matrices does) + a bit more row height. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- plot_component_grids: assert non-empty + return bytes (was bytes|None used as control flow); render_arithmetic_figures skips empty active sets explicitly. - Drop unnecessary quoted forward refs in the lab wiring (imported/defined names). - Remove redundant config defaults from the run yaml; the schema defaults (thresholds=[0.1], top_k=24) are the single source of truth. - Trim arithmetic_eval module docstring to one orienting line; move the detail to param_decomp/CLAUDE.md, and list arithmetic_eval.py / prestage_arithmetic.py in the CLAUDE.md file maps. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…req 2e-6) Dan's jax-l18-b64-200k-eps1e6-seq512 reproduced in the current schema + the ArithmeticCIGrid eval; bitdeer data path + launch_env HF cache. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Make Dan's L18 arith run fit 4 nodes in current code: - TRAIN jit_step: recon scans full 32-layer stack w/ per-layer [32,C,d_in/fsdp] slot -> ~63 GiB C-driven floor. Halve C 49152->24576 -> ~32 GiB. - EVAL jit_eval_step: on the current 4-proc x 8-GPU topology, PGDReconLoss (mask_scope shared_across_batch) and the hidden-acts metrics all-gather the batch -> [32 layers, full batch, seq, mlp] (OOM). Drop PGDReconLoss + Stochastic/CIHiddenActsReconLoss. CEandKLLosses (per-position, required by the in-loop eval) and the CI metrics stay per-device sharded (batch 1). eval.batch_size 32 (must be a multiple of dp). Both remat on; Dan's train batch/seq/data/coeffs kept. Restoring C=49152 + the gathering eval metrics needs the recon scan to skip undecomposed layers' C slot and the eval to shard (or 1-GPU-per-proc). 60-step smoke mirrors. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
compute_hidden_acts_metrics ran unconditionally in the slow tier, unlike the position-CI / UVPlots metrics beside it which are gated on the config naming them. So hidden-acts ran even when eval.metrics omitted CIHiddenActsReconLoss / StochasticHiddenActsReconLoss -- and on the 4-proc x 8-GPU topology it OOMs (full- batch gather) and hits a shape bug (the step derives leading from residual.shape[:-1] = batch-only, so the delta mask is (B,) and mis-broadcasts against the (B,seq,d) activations). Gate it like want_position_ci so the metrics are opt-in. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ions docstring) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Remove the 7 # fmt: skip I added (accept ruff's layout); the pre-existing llama8b:527 skip is left untouched. - Remove comments that restated the adjacent code or an assert/docstring: the bf16-readout precision note, two to_grid shape comments (src + test), the empty-selection guard note, the 'run the fast eval on the probe' note, and the seam-check note that duplicated its assert message. Kept the load-bearing ones: the HLO-baking step-factory note (per CLAUDE.md), the sink/scan invariant assert, the collective-gather tier block, the addition-only prestage constraint, and the test-identity explanations. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…Eval docstring json/pyarrow were imported inside _load_arithmetic_probe; pyarrow is already loaded via param_decomp.data (parquet shards), so the local imports bought nothing — move both to module top. Drop the stale 'recon/L0/PGD' metric names from the _ArithmeticEval docstring (PGD is no longer in the run's eval.metrics; the fast-eval set is config-driven). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Resolved conflicts: - targets/llama8b.py: keep Oli's segmented [frozen prefix -> live -> frozen suffix] masked forward (2f9af23, +29%); re-apply collect_activations (x@V) as a parallel sink in the live block. collect_activations appended as a defaulted param so the 3 existing callers stay byte-identical; gather uses the new layer-first_live offset; masked_component_activations passes it by keyword; ys typed as the (out, acts) tuple; refreshed the stale lax.cond docstring. - lm/run.py: keep the hidden-acts opt-in gate AND thread their compiler_options (co) into compute_hidden_acts_metrics; merge built_run imports (LAUNCH_CONFIG_FILENAME + EvalConfig). - run.py / config.py: merge imports; take their _init_or_restore_state call (adds profile.no_checkpoint + compiler_options). make check clean; full fast suite 463 passed / 11 xfailed (incl. equivalence goldens -> numerics preserved through the masked-forward refactor). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…test The 3 llama8b_l18_arith_*.yaml were our experiment-specific run configs (they bake in the C-halving / eval-trimming / 4-node fitting) — noise for a feature PR, and nothing references them. Remove them; the feature (arithmetic_eval + seam + wiring + prestage_arithmetic + schema + docs) enables via one opt-in metric added to any run config. The exact run config stays reproducible from the run dir's pinned launch_config.yaml + this branch's history. Their yamls were the only exerciser of the ArithmeticCIGrid -> ArithmeticEvalConfig build path; replace that with a unit test in test_config.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…8 config) Add the arith CI-grid eval to the canonical L18-MLP config so it runs by default, artifact_dir on the config's own data mount (prestage the probe there once per cluster, per the inline note). Make the config-path test config-independent (drop any existing ArithmeticCIGrid before injecting a known one). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Feature: _ensure_arithmetic_probe auto-generates the default (add, 1..100) probe on rank 0 at first run if artifact_dir is empty (barrier so other ranks wait), so ArithmeticCIGrid is turnkey — no manual prestage. Arith is Llama-only, so prestage's default tokenizer is the target's; a custom grid stays a manual prestage. C49k note + CLAUDE.md updated; unit tests cover generate-when-missing / idempotent / non-main. Also drop the hidden-acts opt-in gate: that was a bitdeer-topology OOM workaround (base runs compute_hidden_acts_metrics unconditionally and Oli's C49k works), not part of the arith feature. Restores base behavior; the unconditional-run + the leading-shape issue we hit on 4-proc x 8-GPU are a separate, pre-existing fix. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…narrativizing) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ature/jax-arith Clean auto-merge — their per-token-CI-density-heatmap additions (configs.py / built_run.py / lm/config.py / lm/run.py slow-eval wiring) and our arith-eval additions landed in non-overlapping spots. make check clean; 57 tests pass across test_config / test_slow_eval / arith suite. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ature/jax-arith Clean auto-merge — #919 fixes the hidden-acts eval to derive its waist leading from the CI output instead of the token input (the bug we flagged); it touches hidden_acts_eval.py / slow_eval.py / its test, none of which this branch modifies, so no conflicts and no changes to our code. make check clean; 30 tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The ArithmeticCIGrid metric carried an absolute cluster path (artifact_dir) to a parquet probe that was fully derivable from (operation, a_range, b_range) + the target tokenizer — and was already auto-generated in-job. Configs now carry the spec; the probe builds in-memory at startup on every rank (deterministic, so no rank-0 write, no barrier, no row-major order guard) via the new experiments/lm/arithmetic_probe.py. Deletes prestage_arithmetic.py, the parquet/meta.json format, and the cluster-specific path in the C49k config. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Collaborator
Author
|
Pushing a small fix. Ongoing tests surfaced an issue with the eval that depended on kernel selection, which showed up using some reasonable configs but not others. Fixing. |
Collaborator
Author
Actually will wait until Oli's review is in. Don't want to break his work mid-flight. |
…nderer Review fixes on the arithmetic CI-grid eval: - make_eval_step grows a required kw-only n_valid_rows: when set, CE/KL/ L0 and the PGD objective become valid-row-masked means, so the probe's sharding-pad tail rows carry zero weight and the scalars stop varying with device count. None (the corpus tier) keeps the reductions bit-identical. The arith eval builds its own instance with n_valid_rows=n_prompts. - compute_arithmetic_selection replaces the full-grid host allgather (~2 GB/site fp32 at C=49k, x2 grids x every rank): the step now also returns the replicated per-component max CI over real rows only, the selection happens host-side off that (C,) vector, and only the <=top_k shown columns are gathered (index padded to top_k for one jit shape). select_active uses one stable descending ordering per site, so a higher threshold's active set is a prefix of a lower's and every figure indexes the shared columns by prefix. - SlowEvalRenderer + ArithmeticGridRenderer fold into one BackgroundRenderer taking a render closure (submit(partial(...))). - ArithmeticGrid pins contiguous integer axes; the imshow extent is now the outer cell edges (centers +-0.5) so ticks land on operand values. SPEC S28a wording updated for the renderer rename (no semantic change). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Post-merge fix: #948 pointed weight loading at the shared cluster HF cache (llama8b.hf_snapshot_dir, now public), but the arith probe's tokenizer load still resolved against the default HF cache — empty on a cluster node, so the eval would fail at startup while the weights load fine. The tokenizer now loads from the same resolved snapshot dir. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds an in-loop arithmetic-grid eval (
ArithmeticCIGrid) for LM decompositions. On a fixeda × bmodular-arithmetic prompt grid ("<a>+<b>="), it renders — per decomposed component, at the=answer position — ana × bheatmap of causal importance and of the pre-mask activationx@V, plusn_alive/n_droppedand recon / L0 / PGD on the arithmetic probe.prestage_arithmetic.py— offline generator of the grid artifact (one prompt/row, fixed length, single-token answers; carries(a, b)+answer_id). Loaded directly by the eval, not via the streamingLMDataConfigloader.arithmetic_eval.py— one fused step returning both CI andx@Vat the answer position (batch axis kept as the grid); the active set (max CI > threshold) is selected once and shared by the scalars and the top-top_kfigures.llama8b.masked_component_activations— new seam exposingx@Vper live site, extending_run_masked_forward'scollectplumbing (sibling ofmasked_site_outputs, SPEC S31). LM-only; the eval narrows to it via theComponentActivationModelprotocol.ArithmeticCIGridConfig→ArithmeticEvalConfig), the rank-0 backgroundArithmeticGridRenderer, wired through_make_arithmetic_evalinto_make_lm_eval_fnon the slow-eval cadence;configs/llama8b_l18_arith_add.yaml.Related Issue
None.
Motivation and Context
We want to see how well a decomposition reproduces the target model's modular-arithmetic mechanism — the L18 MLP addition neurons from Feucht et al. ("Arithmetic in the Wild"), which fire on periodic functions of the sum
a+b(base-10 Fourier features, periods 2/5/10) and so show diagonal period-stripes on thea × bgrid. Thedown-site activation grids are the target to recover. This ports the relevant slice of the arithmetic analysis from the torchexperiment/8B_targetedwork into the JAX line as a first-class eval, parameterised by site/threshold so it generalises beyond addition.How Has This Been Tested?
tests/test_arithmetic_eval.py(fused step vs hand-rolled CI +x@V; row-major reshape; active-set selection; n_alive / n_dropped; render),tests/test_llama8b.py(the seam:x@V @ U == masked_site_outputsunder all-ones masks),experiments/lm/test_arithmetic_wiring.py(probe-load order guard / missing-artifact / sharding). All green;make type(basedpyright) + ruff clean.x@Vheatmaps render to wandb;n_alive ≈ 93%at init, confirming thetop_kcap and that n_alive starts high and falls as training sparsifies.JAX_PLATFORMS=cpu(cudnn flash rejects the tiny-model's fp32 attention on GPU).Does this PR introduce a breaking change?
No — additive. New optional eval metric;
masked_component_activationsis added only onLlamaDecomposedModel;_run_masked_forwardgains an internalcollect_activationsarg (all call sites updated). Training / recon semantics are unchanged (the eval is figure-tier, off the recon path).