Skip to content

Explicit placement, broader transformer targets, and evaluation refinements - #1000

Merged
danbraunai-goodfire merged 1 commit into
mainfrom
release/2026-08-20
Aug 24, 2026
Merged

Explicit placement, broader transformer targets, and evaluation refinements#1000
danbraunai-goodfire merged 1 commit into
mainfrom
release/2026-08-20

Conversation

@claude-spd1

@claude-spd1 claude-spd1 commented Aug 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

This release refreshes the core VPD/tPD trainer, its runtime configuration, and the included transformer targets while keeping the package focused on decomposition rather than postprocessing applications.

New and changed loss options

  • NonlinearityLocalityLoss penalizes components whose write vectors spread across too many nonlinear units. Its soft-count threshold is schedulable, and neurons and attention heads can be weighted or excluded independently.
  • EMA frequency minimality is available through frequency.ema_halflife_steps; omitting it retains the existing per-batch estimate.
  • Persistent-PGD adversary objectives now support adversary_objective: e2e | term, selecting output reconstruction alone or the complete reconstruction term including hidden-activation matching.
  • The legacy *LayerwiseLoss and ChunkwiseSubsetReconLoss config types are removed; reconstruction losses use the unified all-site or subset forms.

New and changed runtime/configuration options

  • runtime.replicate, runtime.fsdp, and runtime.tp define the logical device mesh explicitly; process-local device count is supplied at launch.
  • Explicit placement tables can independently place component optimizer state, compute and faithfulness weights, operands, Muon Newton–Schulz work, CI-function weights/activations, frozen target weights, and external/intermediate activations. The owner, zero1, and ddp presets remain available.
  • Checkpointing is now a tagged choice between periodic (with keep_last or keep_all retention) and none for disposable measurement runs.
  • Compiler settings can be tuned-v1, bare, or a complete explicit xla_* mapping, with an explicit persistent compilation-cache directory.
  • Profiling can be disabled, captured in-process with ad_hoc, or configured for Nsight Systems.
  • LM targets now select their attention implementation explicitly.
  • Hugging Face target coverage expands from Llama-3.1-8B and Qwen3-8B-Base to dense Qwen3 0.6B, 1.7B, 4B, 8B, and 14B base/instruct variants.
  • Slow evaluation baselines run at true step 0 on fresh runs rather than after the first optimizer update.

Related Issue

None.

Motivation and Context

Refresh the supported decomposition and training workflows while making consequential placement, topology, checkpointing, and compiler choices explicit and portable across compute environments.

How Has This Been Tested?

  • release-package validation, lock validation, Ruff, and basedpyright
  • seven toy-model configurations completed two training steps, evaluation, and checkpointing from a clean wheel install
  • a tiny Llama configuration completed two training steps, evaluation, and checkpointing from a clean wheel install with current Hugging Face dependencies

Does this PR introduce a breaking change?

Yes. Existing configs must migrate the runtime mesh, placement, checkpointing, compiler-cache, and target-attention fields described above; removed layerwise/chunkwise reconstruction loss aliases are no longer accepted.

@danbraunai-goodfire
danbraunai-goodfire merged commit facf2e7 into main Aug 24, 2026
5 checks passed
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 25, 2026
@Antovigo @claude
perf(targets): re-derive the three-stack frozen layout on the new ana…
ff219fa 
…tomy

Restores stacked_prefix / stacked_tail, dropped from the upstream merge because
goodfire-ai#1000 rewrote glu_transformer.py around the anatomy abstraction at the same time.
`stacked` now covers exactly the decomposed span [split_layer, tail_layer); the two
frozen ends are their own stacks.

Why it is worth a field each: a masked forward that slices its blocks out of one
combined stack materializes a multi-GB copy PER FORWARD (bf16[13,14336,4096] here —
the copies XLA names `wrapped_slice`), and there are ~10 masked forwards per step.
The segment bounds already WERE the stored stacks' boundaries — split_layer ==
first_decomposed and tail_layer == last_decomposed, contiguity asserted — so every
segment maps onto ONE stored stack and `slice_layers` never cuts one up. That
identity is now asserted rather than assumed.

Measured on 2x L40 at the single-block L18 allmerged shape, before this commit:
7.14 s/step and 43.79 GB peak, against the pre-merge 3.581 s / 36.45 GB.

Ported with it: `layers`, the attention-pattern read, `shardings` (all three stacks),
both clean forwards (one scan per stored span, each taking its slice of the absolute
capture-slot tables), the per-site frozen-weight reads (span-relative now), and
`_with_attention_implementation`, which must retarget all three or leave a mixed
model that fails in the prefix, before any decomposed block.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011EqWYtK2A5PugNTcQx2iAE
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 25, 2026
@Antovigo @claude
perf(targets): restore prefix REUSE — run the frozen lead once per st…
8784009 
…ream

The three-stack commit fixed the memory (43.79 -> 28.23 GB peak) but only 8% of the
step time (7.14 -> 6.53 s), which located the rest of the gap precisely: the slice
copies were a MEMORY problem, and the remaining ~1.8x is the frozen lead being
re-run inside every masked forward instead of once per stream per step.

So `ResidualStart` comes back. `prefix_residual` runs embed + the frozen lead and
stop-gradients it; `ForwardSubstrate.prep_stream` substitutes the result for the
batch, and every downstream forward in that stream — clean, taps, each recon-grid
draw, each adversary ascent — resumes from it. The substitution IS the optimization:
`StreamInputs.batch` is the single value all of them receive, so nothing else needed
to change. At an L18-only decomposition that is 18 of every forward's 32 blocks,
plus their remat replays.

`_entry` is the one seam: token ids embed and enter at block 0, a `ResidualStart`
enters at `split_layer`. It also answers the two spellings of the start activation
itself — block `split_layer`'s input and block `split_layer - 1`'s output, which is
what a `resid.<split_layer>` point lowers to — and refuses anything strictly below,
since that prefix was consumed and its internals are gone.

Subtlety worth naming: the served boundary capture must be excluded from the SCAN
layout while staying in the requested set. Leaving it in allocates a buffer for a
block the forward never runs, and the empty buffer overwrites the served value on
read-back — the failure looks like a zeroed capture, not a missing one.

test_prefix_reuse.py comes back with it, migrated to the post-goodfire-ai#1000 API (placement
args, per-stack weight_deltas) and with its pre-dechunk partially-live-chunk case
rewritten as the multi-block span claim, which is what still holds.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011EqWYtK2A5PugNTcQx2iAE
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 25, 2026
@Antovigo @claude
fix(ab_grid): consume the eval context's placed CI fn, not the raw one
41ca10c 
`ci_preactivations` goes through the CI compute lifecycle, and post-goodfire-ai#1000
`materialize_ci_compute_weights` reads `.fn` / `.placement` off a `PlacedCIFn`
to reconstruct a chunkwise fn's compute weights against the run's resolved
placement. The A/B grid operation still reached past the eval context to
`state.decomposition.ci_fn` and handed the bare `ChunkwiseTransformerCIFn`
down, so the first snapshot died with

    AttributeError: 'ChunkwiseTransformerCIFn' object has no attribute 'fn'

`EvalInvocation` exists precisely to carry this pair -- its docstring says
operations consume `placed_ci_fn`, "never a raw (fn, rules) pair". The grid
op was the one violator, and it type-checked only because both of its seams
annotated the parameter `Any`; those were the sole two `ci_fn: Any` in the
tree. Type them as `PlacedCIFn` (alias included) so the checker holds the
contract from here on.

Why the merge shipped it: the grid path is branch-only, so upstream's
placement threading never touched it, it has no tests, and the schedule
carries a `step > 0` guard -- its first firing is `slow_every`, not step 0.
Run 10631 therefore looked healthy for 3900 steps and died at step 4000.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011EqWYtK2A5PugNTcQx2iAE
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 31, 2026
@Antovigo @claude
docs: the upstream-merge campaign report, upstream-fix notes, and swe…
71cf0e2 
…ep plots

report.md sections 7-8.20: the goodfire-ai#1000 merge and its fallout (two ab_grid sharding
bugs, the zero1 layout trap, the fine-tune memory ceiling), the L18 nonlinearity
dose-response at six coefficients, the 4-block runs with their matched control,
and the two fine-tunes.

notes/upstream_fixes.md is the contribution-facing writeup: five items that are
genuinely upstream's (CI-scaled weight-decay crash, hard-coded CI attention
backend, the targeted nonlinearity refusal, zero1 at a degenerate fsdp axis, the
frozen-stack split) plus two of ours kept for the lesson that generalises --
explicit sharding made a bug class single-device tests cannot see, and the repo
already has the multidevice marker that catches it.

notes/plots/nlpenalty/ is 18 figures, one metric per stream per family, plus the
values table. Every comparison is matched code / seed / recipe; the confounded
pre-merge 4L baseline appears in none of them.

Headline findings, all with the run ids that produced them:
- the prior cuts nonlinearities per component 2179 -> 18-63 depending on dose
- sparsity and on-distribution reconstruction are nearly unaffected
- the cost is off-distribution and essentially DOSE-INDEPENDENT
- locality is maintained, not locked in: it relaxes once the prior is removed

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011EqWYtK2A5PugNTcQx2iAE
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 31, 2026
@Antovigo @claude
Delta-pin the non-target CE/KL eval family (T4)
709ee52 
Port of 5c30295 from feature/dual_obj_jax onto the upstream-merge branch,
which forked at 14df37a and so never carried it. Original message:

  A targeted run's non-target-stream reconstruction evals ran with every
  weight-delta mask at 0, measuring a component-only quantity no training
  forward defends (T4 pins the delta fully on off-target). Only the
  fresh-PGD probe had been delta-pinned (T4's 2026-08-20 amendment), which
  is why every other nontarget_data recon key sat persistently high while
  the probe did not.

  Extend the rule to the whole eval family, bound off the shared
  nontarget_delta_pinned predicate (targeted && nontarget stream):

  - CEandKLLosses: ci_masked/random_masked/rounded_masked/zero_masked
    compose delta masks at 1.0; stoch_masked pins instead of drawing
    U[0,1] (the T4 training composition); unmasked keeps the enumerated
    delta-off exception (UnmaskedNoDeltaReconLoss).
  - The standalone masked-KL arms (both CI roles): ci_masked pins;
    unmasked stays delta-off.
  - Plain runs and the target stream are unchanged.

Two conflicts, both from goodfire-ai#1000's eval rewrite, resolved toward the merged code:

- eval.py builds every mask source with the sharding-aware `*_like` helpers,
  never `jnp.zeros(shape)` / bare `random.uniform` — a bare draw lowers
  REPLICATED under the Explicit mesh and held the full eval batch per rank
  (112 GiB per kind at the 32L shape). The pinned constant is `ones_like` on
  the token array, and the pinned `stochastic_deltas` branch skips the draw.
- CE/KL construction now routes through `scalar_step_for`, the one
  config->kernel binding the AOT fit check shares, so `delta_pinned` is
  threaded through the dispatcher rather than bypassing it. The fit check
  keeps the default False, which is also the memory-conservative arm:
  pinning REPLACES the stochastic draw with a constant.

The two new tests built a bare model; post-goodfire-ai#1000 the step builders take a
PlacedModel, so they are wrapped in this file's existing idiom. Verified
non-vacuous: forcing variant_delta back to zeros fails the CE/KL test.

SPEC T4 amended 2026-08-28 (pending Oli sign-off).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01P1KiqKVGfRMMvC5Gw85cKA
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 31, 2026
@Antovigo @claude
Merge feature/dual_obj_jax-upstream-merge into feature/dual_obj_jax
d514a2f 
Brings the upstream PR goodfire-ai#1000 merge (facf2e7, "Explicit placement, broader
transformer targets, and evaluation refinements") and the whole campaign that
followed it onto the working branch, so dual_obj_jax carries BOTH the
dual-objective (tPD) reconstruction work and the nonlinearity-occupancy
penalty on one line of development.

What arrives with it:
- upstream goodfire-ai#1000: explicit placement (PlacedModel / PlacedCIFn), the retired
  recon chunking (SPEC S2 amended), the authored (replicate, fsdp, tp) mesh.
- the tPD port onto those primitives: four-pass targeted step, dual CI roles
  (S37), delta-pinned non-target masks.
- NonlinearityLocalityLoss admitted to the targeted loss union (SPEC S36),
  scored once per step in weight space, with tests pinning that it is
  pass-count invariant.
- the three-stack frozen layout (stacked_prefix / stacked / stacked_tail) and
  prefix reuse.
- GPU-only fixes the CPU suite cannot reach: the CI-scaled decay re-type
  (63db6fa), the authored CI-fn SDPA backend (a9111a4), and the two
  A/B-grid sharding fixes (41ca10c, fb07ebf) plus the multidevice test
  that catches their class on CPU (4cbd48a).

Conflicts, all in the delta-pinning fix both branches carry (5c30295 here,
ported as 709ee52 there) — resolved to the ported side in every case, since
it is the same T4 rule expressed against the post-goodfire-ai#1000 code:
- eval.py: keep the sharding-aware `*_like` mask/delta construction. The
  pre-merge `jnp.zeros(shape)` / bare `random.uniform` forms lower REPLICATED
  under the Explicit mesh (112 GiB per mask kind at the 32L shape).
- scalar_eval_operations.py: `delta_pinned` threads through `scalar_step_for`,
  the one config->kernel binding the AOT fit check shares, rather than
  bypassing it.
- test_eval.py auto-merged BOTH copies of the two delta-pinning tests, and the
  duplicate pair was the pre-merge one built on a bare model — which, being
  second, would have shadowed the working copies and failed. Duplicates removed.

The resulting tree is byte-identical to feature/dual_obj_jax-upstream-merge at
709ee52. 530 tests pass there; the eval suite re-verified here against this
worktree's venv.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01P1KiqKVGfRMMvC5Gw85cKA
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 31, 2026
@Antovigo @claude
Carry the addsub SOTA recipe back into the tree as a reference config
00b5ccc 
The upstream merge 4616e8c (PR goodfire-ai#1000) resolved CONFIGS.md and the LM configs dir
wholly in upstream's favour, which silently took addsub-L18-dual-obj.yaml and its
registry row with it; it reached this branch on d514a2f. A set-difference audit
(local at 14df37a, absent upstream, missing now, minus upstream's deliberate
removals) shows the seat was the ONLY purely-local file that merge dropped.

The recipe returns as a reference config rather than a canonical seat: the merge left
the LM registry at its test-enforced cap of 10, and the recipe did not need to displace
one. It is addsub-L18-16-ntmerged-bsc (p-5b7fa697), the 2026-08-21 campaign winner —
the merged nontarget adversary on both the output and hidden streams, target stream
left at 14-4x — carried onto the post-goodfire-ai#1000 schema. Comment-stripped it is identical
to that run's immutable pin apart from run_name and the six migrations goodfire-ai#1000 forced:
xla attention on the CI fn and the target, the folded cadence.checkpointing block, the
newly required compilation_cache_dir/compiler_options, the explicit replicate mesh axis
for dp/gpus_per_node, and sharding ddp for zero1.

Living outside param_decomp/experiments/lm/configs/ means test_repo_configs_parse.py
does not gate it and no schema change will migrate it — stated in the file's own
header. Validated by hand against tip instead: parse, placement gate, model_dump and
both file round-trips, and the names-never-locations rule.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0191UYkh6wHq8ZM49wMgdVuW
Antovigo added a commit to Antovigo/spd that referenced this pull request Aug 31, 2026
@Antovigo @claude
Drop CIHistograms from the SOTA reference config — goodfire-ai#1000 m…
a615770 
…ade it fatal here

Upstream goodfire-ai#1000 added an assertion in `slow_eval.py::accumulate_site_reductions`: the
CIHistograms value histograms bin against each batch's own min/max, so counts from more
than one eval batch cannot be summed. The config runs `eval.n_steps: 10`, so the metric
killed the run in its FIRST slow eval pass — caught by a 25-step smoke (job 10924),
not by any offline check.

`n_batches_accum: 1` does NOT save it, which is what made this easy to miss: that field
only constrains what the op may accumulate, it does not cap the batch list the assertion
inspects. p-5b7fa697 carried the metric happily for 20k steps pre-merge, where the
assertion did not exist, so the pin is not wrong — this is a migration item that the
recipe's carry-over missed, and the same call the migrated nlcontrol config had already
made. Keeping it would require `eval.n_steps: 1`, rewriting every other metric's
statistics for the sake of two histograms.

Neither `build_targeted_experiment_config` nor the parse gate can reach this: it lives in
the eval pass itself. Re-validated the config offline as before, then confirmed on GPU —
smoke 10927 completed rc=0 with a full AB-grid snapshot (index.html, manifest.js, a
104 MB step_20.js, saved_components/total 1952) and checkpoints at steps 20 and 25.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0191UYkh6wHq8ZM49wMgdVuW
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@danbraunai-goodfire danbraunai-goodfire danbraunai-goodfire approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@claude-spd1 @danbraunai-goodfire