…ot just the batch shell

--signal=B:TERM@300 only signals the batch script; the trainer's SIGTERM-save
handler never fired and preemption (job 50818) hard-killed the ranks after the
grace window — losing everything since the last save_every checkpoint.

jsp-export writes vendored-layout keys regardless of the run yaml's target spec,
so eval must build VendoredLlama (same weights) with model.-prefix patterns
stripped; frozen forced bf16 = what the JAX run actually trains with (its
documented fp32-yaml divergence), and the more faithful eval reference. Unblocks
the C49k offline-eval loop (ckpt-5000 attempt died on the raw-HF dispatch).

…val args (C49k needs --micro-batch-size 4: mb8 OOMs at 157.6+7.8 GiB on a 178 GiB B200)

Per Oli: resilient-launch loops are an anti-pattern this stack shouldn't need. The
one recurring crash class is fixed at the source (command buffers off, ~0% cost);
the loop demonstrably contributed nothing across last night's incidents — it
burned all 5 attempts on each deterministic OOM, could not escape a suspect node
(retries reuse the allocation), and every actual recovery was a diagnosed
resubmission. Plain --requeue + the now-working SIGTERM save = the torch model.

…one eval job per checkpoint save

Replaces the polling sidecar chain (which self-multiplied and hid failures): rank 0
fires sbatch offline_eval_once.sbatch right after each save, fire-and-forget (a
failed submission must not kill a multi-day run — the one place graceful handling
beats fail-fast; it prints loudly instead). Per-run serialization via
--job-name=jsp-oeval-<run> + --dependency=singleton; the run-dir marker dedups
across requeues; optional <run_dir>/offline_eval_args carries run-specific flags
(C49k: --micro-batch-size 4).

…yout, torch postprocess contract

Wrapper gains run_id (p-<8hex>, the torch generate_run_id convention; authored
into the wrapper so resume derives the same identity and the byte-compare pins
it): run dir = out_dir/<run_id>, wandb id = run_id with run_name as the display
name. The pinned torch yaml is additionally written as experiment_config.yaml
(SavedLMRun's contract name), and the push-triggered eval job materializes
model_<step>.pth (safetensors -> torch.save, pruned to newest 2) — together
making a JAX run dir a first-class citizen of harvest/app/postprocess with zero
special-casing. wandb config gains a jax_runtime section (actual device count
etc.) since the upstream yaml's runtime.dp describes the torch run, not this one
(it read 'dp: 32' while running on 64 GPUs). run_id is optional ONLY for the
pre-scheme live C49k wrapper — collapse after it migrates.

…rity fixtures

gen_stacked_fixtures.py runs against the PRE-restructure feature/jax-single-pool-pd
code (stacked DecompVU, LayerRange) and pins, on the tiny L3-5 MLP target: target
weights, per-site V/U, CI-fn/source init leaves, clean/masked/site-input/weight-delta
forwards (S2/S3/S4/N2), and a 2-step training trajectory (metrics + final V/U +
adversary sources, S13-S15). The generated npz was produced from the live base-branch
checkout's venv; the script is excluded from this distribution's basedpyright (it
deliberately imports the old API and only runs on the base branch).

The follow-up site-generality restructure must reproduce these to SPEC D4
reassociation tolerance (clean logits bit-identical).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…e/up/down, per-site C

Generalizes the Llama target from contiguous-layer MLP-only to any per-layer matrix
site, each with its own C. The trainer (train.py) was already site-generic; this
rebuilds the target + plumbing around per-site state:

* llama8b.py — DecompVU becomes {site: (V (d_in,C_s), U (C_s,d_out))} fp32 masters
  (same init scale rule: V~N(0,d_in^-0.5), U~N(0,C^-0.5)); Target is a uniform
  SuffixLayer list from first_decomposed_layer(sites) to the LM head — layers with no
  sites run the plain frozen block, and any site absent from the spec computes the
  exact frozen x@W path (S2: masked_logits(live=()) stays bit-identical to
  clean_logits, S3). q/k/v sites are decomposed BEFORE RoPE/SDPA via FrozenAttn.core
  (head-count reshapes asserted); o applies to the attention output. Clean CI inputs
  (S4): q/k/v <- post-LN1 residual, o <- pre-o_proj attn output, gate/up <- post-LN2
  residual, down <- silu(gate)*up. weight_deltas stays fp32 W - V@U per site (N2).
  Canonical site order = layer-ascending x KIND_ORDER (computation order); site names
  keep the torch-module-path convention layers.{i}.{self_attn,mlp}.{kind}_proj.

* llama8b_sharding.py — per-site placement: V shards C on axis 1 (P(None,'dp')),
  U on axis 0 (P('dp',None)); per-site C % mesh-size asserted; the
  jit-with-out_shardings init pattern and batch_sharded_ci are untouched.

* config.py / torch_config.py / run.py / run_state.py / export.py — TargetConfig
  carries canonical (site, C) pairs; the native {first_layer,last_layer,C} YAML maps
  to the MLP family (unchanged behavior); the torch route accepts any
  (model.)layers.N.{self_attn,mlp}.{kind}_proj target with per-site C (contiguity /
  uniform-C / all-three-kinds asserts dropped; identity-target and other-pattern
  refusals kept). Export reads V/U per site and renames any decomposed matrix's
  frozen weight to .target_weight — attention sites included (the torch vendored
  ComponentLlamaAttention componentizes q/k/v/o, so attention-site export is fully
  supported; round-trip-verified against the real LMComponentModel via the new
  l18_attn fixture case, heterogeneous C, key parity exact).

Equivalence with the stacked implementation (tests/stacked_parity, fixtures pinned
from feature/jax-single-pool-pd): clean/masked/site-input/weight-delta forwards
BIT-identical; 2-step train trajectory rel <= ~2.5e-6 (D4; sole divergence is
clip-global-norm leaf-order reassociation). Cross-framework loss-term equivalence
(tests/equivalence) unchanged and green.

invariance_check gains an abs tolerance floor (1e-7) alongside rel 1e-4: the per-leaf
grad-norm diagnostics include ~1e-4-magnitude norms whose cross-shard reduction
cancellation can graze rel 1e-4 (observed: one ci-fn wv leaf, abs err ~1e-8, step-0
only, non-growing) while every loss term sits at ~1e-6. Verified at 4 and 8 sim
devices (worst loss-term rel 2.5e-6 / 7.1e-6).

Checkpoints are NOT cross-compatible with stacked-era runs (components pytree layout
+ V/U init RNG derivation changed).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…EC §3 site prose

SPEC.md needed no semantic change: §1/§3 already define a site as any selected weight
matrix. Amended the one §3 orientation sentence that described the *implementation* as
contiguous-MLP-only (now: any per-layer q/k/v/o/gate/up/down matrix, per-site C; q/k/v
decomposed before RoPE/SDPA) and added the attention-site examples to the §4.1
site_inputs comment (S4). No invariant changed.

AUDIT.md records: stacked->per-site parity evidence, checkpoint non-compatibility
across the restructure, and the verify_export_torch production-numerics pass becoming
measure-only (the documented GELU/eps divergence is amplified on the tiny attention
fixture; the asserted jax-matched pass is the mapping proof).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…orch PGDReconLoss parity)

AdversaryConfig = SourceAdamConfig | FreshPGDConfig selects the recon adversary
variant in make_train_step: PPGD keeps its persistent sources + Adam moments +
fused (n_warmup+1)-th ascent (SPEC S13/S14, unchanged); FreshPGD samples per-site
sources every step (unique_per_datapoint (B,T,C+1) or shared_across_batch
(1,1,C+1)), runs n_steps of step_size*sign(grad) with clamp [0,1], and carries NO
cross-step state (TrainState.sources stays empty; the post-backward source
cotangent is unused). Metrics key 'pgd' -> wandb train/loss/PGDReconLoss.
ExperimentConfig.ppgd renamed to .adversary (native yaml key 'ppgd:' unchanged);
the torch-config route accepts PGDReconLoss as the adversary slot. Matches Dan's
p-a005ed60 config (n_steps=1, step_size=1, init random, unique_per_datapoint).

…vs replicated vs bridged) + remaining gaps in dependency order

…LM implementation

The t-9d2b8f02 pile model (4L, d768, 6 heads, GELU MLP, plain rotate-half
RoPE, tied head) as a decomposition target. llama_simple_mlp.py mirrors
llama8b.py's seams: per-site maybe-masked matmuls for all six site kinds
(h.{i}.attn.{q,k,v,o}_proj before-RoPE/after, h.{i}.mlp.c_fc before the
GELU, down_proj after), frozen residual-start suffix + embedding/blocks
prefix, fp32 weight deltas (N2), live=() bit-equal to clean (S2/S3).

Registration: TargetConfig | LlamaSimpleMLPTargetConfig union in config.py
(native target block dispatches on pretrain_run_path); torch_config accepts
kind: pretrained LlamaSimpleMLP specs with h.* wildcard decomposition
patterns expanded over the checkpoint's n_layer; run.py builds the target
by kind and threads prefix_residual_fn; export/offline-eval guarded
llama8b-only.

Weights load from the torch pretrain cache via a one-off .pt -> safetensors
conversion (tools/convert_llama_simple_mlp_checkpoint.py, torch venv; tied
lm_head stored once as wte.weight).

Torch-fixture equivalence (tests/simple_mlp_equivalence/, fp32): tiny
random model (GQA repeat=2) max abs logits diff ~2e-7; real t-9d2b8f02
weights ~5e-5 — pins RoPE construction, tanh-GELU, GQA grouping, rms eps.
Tiny-config site tests mirror test_llama8b.py (masked identity, live=()
bitwise, mixed attn+mlp train step with heterogeneous C: S13/S15/S9/N1).
Suite green at 1 and 8 sim devices; basedpyright + ruff clean.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

EvalConfig.l0_groups carries the fnmatch site patterns; make_eval_step resolves
members at build time (unmatched group refuses) and emits the torch group keys
(l0/<thr>_<group> = SUM of member-site L0s). Unblocks Dan's pile config whose
eval block groups per-layer + total.

… feature/jax-site-generality

Scope names now spell the stored source shape in tensor order: single_source->c,
broadcast_across_batch->sc, repeat_across_batch->nsc, per_batch_per_position->bsc;
PGD MaskScope -> {c, bc, bsc}. Classes renamed to match (CScope/SCScope/NSCScope/
BSCScope). Code + tests + reference yamls swept together; 514 tests green.

Stored experiment_config.yaml files predate the shape-spelled scope names;
BeforeValidator aliases exactly the literals that occur in stored data
(broadcast_across_batch->sc, per_batch_per_position->bsc; mask scopes
shared_across_batch->c, unique_per_datapoint->bsc). Delete once stored runs
are migrated.

# Conflicts:
#	param_decomp/metrics/base.py
#	param_decomp/metrics/importance_minimality.py
#	param_decomp/optimize.py
#	param_decomp/tests/metrics/test_importance_minimality_loss.py
#	param_decomp_lab/experiments/lm/run.py
#	param_decomp_lab/run_sink.py
#	param_decomp_lab/tests/test_resumption.py

…t references

…ckage

# Conflicts:
#	param_decomp/metrics/base.py
#	param_decomp/metrics/persistent_pgd_recon.py
#	param_decomp/metrics/persistent_pgd_state.py
#	param_decomp/metrics/pgd_utils.py
#	param_decomp/tests/metrics/test_recon_losses.py
#	param_decomp/tests/test_optimize.py
#	param_decomp/tests/test_spd_losses.py
#	param_decomp_lab/experiments/lm/data.py
#	param_decomp_lab/tests/test_flat_vs_two_pool_loss_terms.py
#	param_decomp_lab/tests/test_three_pool_grad_check_distributed.py
#	param_decomp_lab/tests/test_two_pool_grad_check_distributed.py
#	param_decomp_lab/three_pool/optimize.py
#	param_decomp_lab/three_pool/two_pool_optimize.py

… feature/fsdp-lm-trainer

# Conflicts:
#	param_decomp/optimize.py
#	param_decomp_lab/experiments/lm/_gpt2xl_ab/gpt2xl_2pool_b16_8gpu.yaml
#	param_decomp_lab/experiments/lm/_llama8b/llama8b_l18_b1280_seq2048_2pool_30node.yaml
#	param_decomp_lab/experiments/lm/_llama8b/llama8b_l18_b512_2pool_lr_low.yaml
#	param_decomp_lab/three_pool/step_chunkwise.py
#	param_decomp_lab/three_pool/step_pool_a.py
#	pyproject.toml

… flow through the package aliases)

…heckout launches; the frozen clone keeps its spooled copies)

Recomposes the torch submit stack (generate_run_id, create_git_snapshot,
SlurmConfig/generate_script, submit_slurm_job) into a JAX composition root:
mint p- id -> snapshot ref -> immutable shared-FS workspace (clone + both
venvs, built at submit time since 8 srun tasks/node would race in-job
clones) -> stamp run_id into the workspace wrapper yaml -> sbatch.
SlurmConfig grows signal/ntasks_per_node; generate_script takes a setup
override. nano pyproject: cuda extra + in-tree param-decomp-config dep
(drops the provisional git-URL pin).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ide row truncation

Config byte-derived from Dan's stored runs/p-af354eb1/experiment_config.yaml
(pgd1step-ss1.0) with documented edits: staged parquet data, save_every
5000/keep 2. The pile artifact's rows are 513 wide (512+label); the shard
server now truncates to the leading seq_len exactly like the torch loader.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

llama8b.sbatch / llama8b_slurm.sbatch ran jsp-train from the LIVE checkout
($SLURM_SUBMIT_DIR) — the reason the C49k run is pinned to a frozen clone.
pd-jax-lm launches from an immutable snapshot workspace instead. The XLA env
and srun shape they carried live on in jax_launch.py. offline_eval_loop.sbatch
was the pre-push-trigger eval design, long dead.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Completes the torch slow_eval/step pattern (infra/wandb.py + optimize.py:471)
for the eval/* namespace: pd-offline-eval writes retroactively into the live
run and can't pass step=, so it rides an eval/step axis — but the in-loop
writer never stamped that field, leaving the two writers on incompatible
axes (offline rows rendered at head+1, e.g. 'step 10201' on C49k).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…r's job

The offline sidecar recomputed the fast eval scalars in torch and logged them
under the same eval/* keys the JAX in-loop eval writes, putting two writers
(on incompatible step axes) on one key family — the 'eval at step 10201'
confusion. Now: offline instantiates only slow metrics and logs only
slow_eval/* on its slow_eval/step axis; eval/* has exactly one writer (the
trainer, default axis), matching torch single-pool semantics. Drops the
eval/step axis entirely (reverts 5774791's eval/* half).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Each train log line leads with wall-clock elapsed<eta (e.g. 2:31<6:50) from
the recent step rate, and drops the per-param grad norms that dwarfed the
metrics — the full breakdown still rides to wandb + jsonl. No progress bar
(no in-place redraw in a SLURM .out anyway), so no tqdm dependency.


Claude-Session: https://claude.ai/code/session_01XBgqY1rYzgQbbLdgfQZsbD

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…ot fp32 (#905)

* fix(slow-eval): run the CI-fn readout in training precision (bf16), not fp32

The slow/plot eval tier deliberately read the CI fn out in fp32, but the CI
transformer's attention routes to cuDNN flash, which rejects fp32 — so the slow
tier crashed at the first slow eval on GPU:
    NotImplementedError: Q must be fp16/bf16/fp8_e4m3fn/fp8_e5m2, got float32
(hit by every run, including bf16-only ones — unrelated to fp8 work).

Rather than route fp32 to the XLA attention impl (#885), run the CI-fn readout in
training precision (bf16) — matching train.py / eval.py and the hidden-acts +
attn-pattern slow tiers, which already do. This is the more faithful readout (the
deployed model runs bf16) and keeps cuDNN flash (faster, no (B,H,T,T) materialize).
Reductions/returns stay fp32 so the host accumulation is byte-unchanged.

Also incorporates #885's precision-independent multihost fix: the per-position
histogram sample keeps the dp-sharded batch axis, so `np.asarray` on >1 process
spans non-addressable devices — gather it with process_allgather(tiled=True).

Supersedes #885's ci_fn.py change (no fp32 attention path is needed once the slow
tier is bf16); incorporates its slow_eval.py gather fix.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019uMvZPo7hyAFgGbhDdLrEL

* style(slow-eval): trim the bf16-readout and gather comments to the load-bearing why

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019uMvZPo7hyAFgGbhDdLrEL

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…64 OOM) (#898)

* fix(load_run): replicate HF prefix without cross-host allgather (dp>=64 OOM)

jax.device_put(host_array, P()) runs multihost_utils.assert_equal ->
process_allgather(tiled=True), tiling the ~1GB embedding to ~process_count GB
(OOM at dp=64). Build the replicated global array from each host's local copy
via make_array_from_callback instead.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012PUbea772bseCPww4uv3JE

* docs(sharding): trim place_via_shardings docstring

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011Y8zwFyb74dftPrAFHnej2

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…throughput (#903)

Squash of the full32L HSDP perf line. Buffer donation (donate=all-except-first) + V/U reconstruction hoist take full Llama-8B VPD to b128/4-seq-per-GPU on 32 GPUs at ~3.6x the b32 throughput, numerics bit-identical. Includes env-gated memory/timing instrumentation, the b128 production config, and doc cleanup (PERF_NOTES distilled to Lore; completed migration docs removed). Full per-commit history on branch perf/gather-coalesce-unroll-k; canon in lore 2026-06-29--full32l-hsdp-donation-canonical-scaling-memory.

…review (#900)

* fix(jax): correctness + multi-host robustness fixes from feature/jax review

A review of the JAX trainer against the torch oracle (torch-oracle tag) and
SPEC.md found the training math faithful; these address the issues it surfaced.

- fix(eval): identity_ci_error counted off-diagonal errors only inside the
  min(shape) block, silently ignoring active components in the trailing columns
  of overcomplete decompositions (C > rows — the normal toy/LM case). Count over
  the full matrix minus the block diagonal, matching torch
  IdentityCIPattern.distance_from. Fixed in all three copies (slow_eval.py +
  tms/resid_mlp model.py).

- fix(run): SIGTERM was read from a per-process flag at collective decision
  points (faith-warmup exit, eval entry, orbax save), so a per-task SIGTERM with
  no cross-rank simultaneity could diverge ranks onto different branches and hang
  a collective. Reconcile the flag across ranks (OR-reduce) once per step into a
  local used at every gate; drop the now-redundant per-rank checks inside the LM
  eval_fn (the pass is admitted only after consensus and runs to completion).

- fix(config): wire faithfulness_warmup_weight_decay (was hardcoded 0.0,
  silently ignoring the field) and drop the redundant ==0.0 canonical assert;
  add a fail-fast assert that the persistent-PGD source LR is fn_type=="constant"
  (source_lr ignores decay, so a decaying schedule was silently flattened).

- fix(eval): log the LR the step actually applied (now_step-1), not next step's;
  permute the lower/upper CI heatmaps each by their own-derived permutation
  (torch parity) instead of sharing the upper-derived one.

Verified: make type clean; core tests pass at 1 and 4 simulated devices
(incl. the device-count-invariance suite); slow_eval tests pass under the Agg
backend (validating the heatmap change); toy/schedule/config/launch tests pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UKuMWZoyShCYVd3Q937KfA

* refactor(jax): trim redundant comments, drop dead sigterm_received

The SIGTERM consensus rationale was stated four times; keep it once in the
_sigterm_consensus docstring and drop the two call-site comments that restated
it. Tighten the docstring and the LR-logging comment. Remove sigterm_received()
— now unused after the LM eval_fn stopped reading the per-rank flag.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UKuMWZoyShCYVd3Q937KfA

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(stacked-parity): portable fp32 tolerance, not bit-exact across CPUs

test_clean_output_bit_identical used jnp.array_equal and the sibling forward
pins used rtol=1e-5/atol=1e-6. Those encode bit-exactness against the
fixture-generating host, but ubuntu-latest is a heterogeneous runner pool:
float32 matmul reduction order differs by ~1 ULP across CPU microarchitectures,
so the exact check and the near-zero elements under the tight atol flaked
intermittently (failed 3 of the last 4 CI runs; passed only when a run happened
to land on a matching CPU). The forwards are CI-fn-independent and the math is
unchanged — this is pure reassociation noise, not a regression.

Fold clean_output into the tolerance-based pins (rename -> test_clean_output_matches)
and raise the shared forward tolerance to rtol=1e-4/atol=1e-5. Observed
cross-arch noise is <=3.1e-6 abs (measured on arm64 vs the fixture host), so the
new bound keeps ~3-4x headroom while staying essentially exact for fp32.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01CixstgLX5XaRm8CwAEReXk

* test(llama8b): rebuild vu/ci_fn per state in fresh-pgd test (step donates)

make_train_step now donates the state (donate="all-except-first", #903), so
the shared module-level vu/ci_fn were deleted after the first run_step and the
second crashed on a deleted buffer. Build them fresh inside make_state with the
same deterministic keys — bit-identical init, independent donatable buffers.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011Y8zwFyb74dftPrAFHnej2

* docs(stacked-parity): don't attribute the ~1e-4 magnitude to SPEC D4

D4 is the reassociation concept; the specific tolerance is ours. Reword so the
docstring cites D4 for the idea without misquoting it on the number.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011Y8zwFyb74dftPrAFHnej2

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…h_env (#906)

The SLURM rank env (XLA flags, NCCL/host-memory knobs, PD_* profiling toggles)
was a hardcoded `_RANK_ENV` block in launch.py, so a run's config.yaml didn't
capture what env/XLA flags it ran with, and A/B-ing a flag meant editing a shared
launcher file.

Add `RuntimeConfig.launch_env: LaunchEnv` (param_decomp.configs) as the single
source of truth — the formerly-hardcoded values become its defaults. `LaunchEnv`
carries typed knobs (xla_flags, xla_python_client_{mem_fraction,allocator},
xla_pjrt_gpu_host_memory_limit_gb, nccl_debug, malloc_arena_max), a free-form
`env` escape hatch (merged last), and a typed `ProfileConfig` that renders the
PD_* profiling toggles — so a profile run is a config, not an env hack.
`launch.py::_render_rank_env` renders `LaunchEnv.as_env()` into the exported bash
block; LD_LIBRARY_PATH stays launcher-computed (machine-specific). The old
`pd-lm --allocator` flag is now `launch_env.xla_python_client_allocator`.

Defaults render byte-identically to the old block (test pins this). Inline
(dp is None) path inherits the caller's env, unchanged.


Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

remat=True keeps nothing_saveable (production default, no behavior change); remat=False now uses dots_saveable so retained-activation forwards stop recomputing the matmuls. Gather is JIT-and-freed either way.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

Passes gpu_max_activity_api_events via ProfileOptions.advanced_configuration so a full train step fits (the default ~1M-event cap truncates the step mid-forward).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

Trimmed from full32L to layer 18 only; per-matrix C doubled; single recon chunk (sites_per_chunk 7, coeff 0.5).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

… in-block

The stochastic recon forwards each held a pre-built [n_layer,B,S,C] mask stack;
collapse them to ONE shared CI stack recomputed inside the checkpointed block
(faithful by checkpoint determinism). The CI envelope and the compute weights are
each forward-evaluated ONCE per step (single eqx.filter_vjp / jax.vjp) rather than
once detached for the adversary ascent + once inside the main backward.

Memory (full32L b128 dp32, MEMPROF 142183 -> MEMRECOMP 142698):
  temp arena 113.79 -> 98.88 GiB (-15), runtime peak 163.30 -> 148.38 GiB (-15),
  headroom under the 164.08 GiB limit 0.78 -> 15.70 GiB.

Interface: stack_ci + masked_output_stochastic on the DecomposedModel protocol;
scan targets recompute in-block, toys delegate to run_stochastic_masked_output.
SPEC unchanged (pure recompute restructuring; the stochastic draw need not match
the prior goldens -- only fwd==bwd faithfulness matters).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

The "lazy DDP" experiment replicated the bf16 compute weights fully (optimizer
masters still ÷N) to make the per-layer FSDP all-gather vanish — trading a full
model-scale V/U resident for no gather. It does not scale (holds a full model
copy; we are comms-bound precisely because the model is large), so retire it:
the ProfileConfig.replicate_weights toggle + as_env mapping, the train.py grad
reduce-scatter branch, and the llama8b _reconstruct_compute_weights branch. The
default ÷fsdp compute layout is now unconditional.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

…gs → compiler_options

Replace the debuggy PD_* env-var surface with proper config the trainer reads directly:

- Profiling toggles (mem_profile / time_steps / trace / async_test / leaf_bench /
  no_checkpoint / profile_max_events) are read straight off `ProfileConfig` in run.py —
  no more config -> as_env() -> os.environ round-trip in the same process. Drop
  ProfileConfig.as_env().
- The compute-experiment toggles become RuntimeConfig fields threaded as args, not env:
  `scan_unroll` (native lax.scan(unroll=k)), `gather_fp8`, `ascend_replicate`.
- XLA compiler flags move off `launch_env.xla_flags` (XLA_FLAGS env) onto
  `RuntimeConfig.compiler_options`, passed NATIVELY to every jit via a typed
  `jit_util.filter_jit` helper. Unlike XLA_FLAGS env, compiler_options are in the
  compile-cache key, so a flag A/B actually recompiles (kills the stale-cache confound).
  Default = the tuned MaxText set, so a config with no overrides comes out tuned, not naked.
  `LaunchEnv.as_env()` now renders only the genuine pre-init env (client mem/allocator,
  NCCL, glibc).
- Remove the ci_broadcast experiment (vmap-all-chunks CI fn) and the PersistentPGD
  `start_frac` field (+ its config fixtures).

make check clean; 87 tests pass.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

Resolve per-layer liveness at trace (live_set is static) and run the masked
forward as [frozen prefix] -> [live block] -> [frozen suffix] sub-scans, instead
of one scan over all layers with a per-site lax.cond. The cond was a packing +
scheduling barrier: removing it lets XLA pack the per-layer V/U all-gathers far
more aggressively and prefetch across the live block. Only the live block carries
V/U; frozen segments reuse clean_output's body and gather no V/U.

full32L 4-chunk (sites_per_chunk=56) step-2 wall: 8.088 -> 5.698s (-29.6%);
defining all-gather ops 467 -> 210. Numerically bit-identical (SPEC S2/S3 -- a
frozen segment is the old frozen branch op-for-op; equivalence goldens unchanged).

Assumes layer-aligned (sites_per_chunk % n_decomposed_kinds == 0) + contiguous
chunks, asserted loudly. Drops partial-layer (per_site / layerwise) recon support;
the one partial-layer unit test became a whole-layer ablation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

`pd-lm` runs via `fire.Fire`, which parses `--tags a,b,c` into a tuple (but keeps
a value containing a hyphen as a string). `tags.split(",")` then crashed with
"'tuple' object has no attribute 'split'". Normalize tags whether Fire hands back
a str, a tuple, or None.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

The engine already applies grad clip to whichever optimizer carries a
grad_clip_norm (`_adamw_with_clip` builds `chain(clip_by_global_norm_with_eps,
adamw)` for both), using the torch-parity clip. Only the canonical-config assert
gated it off for the CI-fn optimizer. Drop that assert so CI-fn grad clip is an
optional knob; the components clip stays required (part of the method).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

The ce_kl/ce_unrecovered_<variant> metrics are no longer used in wandb.
Remove them from the in-loop eval pass, the eval-key test, the averaging-
parity test, and the LM headline print (now ce_difference_ci_masked).


Claude-Session: https://claude.ai/code/session_01Ayu51JNPUaYbeWvhMubBL9

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

`--requeue` re-runs the same job id into the same `slurm-<jobid>.out`; SLURM's
default output mode is truncate, so an auto-requeue wiped the prior attempt's log
— destroying the crash evidence that caused the requeue. Append instead.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

Enforce restore_step's documented contract (restore "onto reference's shardings")
explicitly via device_put, guarding against a multi-host restore landing a leaf on
a layout differing from what the jitted step was compiled for — which triggers a
costly entry-reshard on the first post-resume step (suspected cause of the 16-node
resume OOM). No-op when the restore already matched. Efficacy on the OOM is pending
a live-resume repro.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

…-min logging (#907)

Make the smooth-L0 (Geman–McClure) penalty annealed γ:1→0.01 the default
importance-minimality loss for the pile 4-layer decomposition, the γ-endpoint
the redescending-gamma sweep found threshold-free (tail ≤0.01 collapses to ~1%,
L0@0 matches the honest count). Preserve the prior L_p 2→0.4 baseline verbatim as
`jose.yaml`. Frequency term carried over unchanged for parity.

Log the imp-min loss and its annealed schedule param under penalty-kind-specific
wandb keys so L_p and smooth-L0 runs are distinguishable:
  loss   train/loss/ImportanceMinimalityLoss   | .../SmoothL0ImportanceMinimalityLoss
  param  train/schedules/p_imp                 | train/schedules/gamma_imp
and add the SmoothL0ImpMin short name for the wandb.config surface. The key is
chosen once at step-factory build time off the imp-min config kind.


Claude-Session: https://claude.ai/code/session_01Vhn77KtbeFej14pmwLbJoQ

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: oli <oli@goodfire.ai>

* fix(runs): pin run config as launch_config.yaml, not config.yaml

The trainer pins each run's self-contained config to `run_dir/config.yaml`
(`_pin_config_copy`), but `config.yaml` is also wandb's reserved run-config
filename. `MetricsSink.for_run` then `wandb.save`s that path, which symlinks
wandb's `files/config.yaml` onto the pinned copy; wandb writes its own run
config through the symlink and clobbers the pinned launch config with a
`_wandb:`/`<config>:` value dump.

In-loop readers survive (they read before wandb.init), but every off-loop
consumer of a finished run's config — `load_run_dir_config` (open_jax_run /
run_metadata / harvest / autointerp), fine-tune structural-compat, slow-eval
metric re-validation — reads the clobbered dump and fails.

Rename the pinned file to `launch_config.yaml` (a non-reserved basename) via a
single shared `LAUNCH_CONFIG_FILENAME` constant in core, routed through every
writer (LM + toy trainers) and reader. wandb no longer collides, so the pinned
config stays clean and off-loop consumers work directly.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DZBC35BZEkEw2ViomVHE89

* docs(built_run): shorten LAUNCH_CONFIG_FILENAME docstring

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DZBC35BZEkEw2ViomVHE89

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat(slow-eval): opt-in per-token CI density heatmap

Add a config-gated per-token CI density heatmap as an add-on to the
CIHistograms slow-eval metric, sharing its masked-free CI forward. When
`CIHistograms.density_heatmap_n_bins` is set, the jitted step also emits a
per-component on-device bincount (C, n_bins+1): column 0 = underflow
(CI < 1e-9, incl. exact-0 inactive), columns 1..n_bins = log-spaced [1e-9, 1]
bands. Accumulated over EVERY eval batch (uncapped — small reduction, no
raw-value host transfer, unlike the n_batches_accum-capped raw sample).

`plot_ci_density_heatmap` renders it active-conditional on a log y-axis
(drops the underflow column), components sorted desc by mean CI into <=256
rank-blocks, per-column-max (over visible bands) magma color, twin log-axis
mean-CI overlay. Emitted as `figures/ci_density_heatmap`. Verified faithful
against the offline reference on p-13c22bca (33.6M tokens).

SPEC S29 amended (2026-07-01) to carry `density_heatmap_n_bins`.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(slow-eval): bump CI density heatmap x-resolution to 600 blocks

Component-axis rank-blocks 256 -> 600 for finer x granularity in the
density heatmap figure. Pure render param; the on-device histogram and
its n_bins are untouched.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GqFFjKFdzUaa81cptwWrx4

---------

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

… inputs (S31) (#919)

Both hidden-acts eval steps computed leading = residual.shape[:-1] on what is
actually the model INPUTS (an LM's [batch, seq] token ids), yielding (batch,)
instead of (batch, seq). Every leading-shaped tensor (stochastic delta masks,
zero deltas, the clean forward's all-false routes) lost the sequence axis and
crashed site_out's route/delta broadcast the first time a run reached a slow
eval with the metrics enabled — the deterministic CI step too, via
route[..., None] in the clean forward, not just the stochastic delta path.

leading now comes off the CI output [*leading, C] (_waist_leading), matching
how lm.py::stochastic_site_masks already draws delta masks (ci.shape[:-1]) and
staying target-generic. The misnamed residual arg is renamed to inputs: Any per
the DecomposedModel protocol — the misnomer is what invited the bug; the
sibling attn_patterns_eval.py names the same thing tokens.

Regression test with batch != seq (fails pre-fix on both steps) pinning
per-site sums/counts and the token-weighted accumulation.


Claude-Session: https://claude.ai/code/session_016wiNqVkCMzQ8UQtkeEZ8eK

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

The train loop subsamples metrics at the log step rather than reducing
over the window since the last log. That's fine for most metrics, but
loses the shape of grad-norm spikes between logs.

For grad_norms/summary/{components,ci_fns,total} — the clip-relevant
aggregates — accumulate the per-step values across the window and log
min/max/median at each log boundary, replacing the single subsampled
point value. Per-leaf grad norms and all other metrics are unchanged.

The per-step values are held as device handles, not floats: appending a
handle is async, so the loop stays unsynced between logs and the whole
window reduces in a single host transfer at the log step.


Claude-Session: https://claude.ai/code/session_01GjuHowooGKrpNaGCXmiztg

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…922)

The PPGD term builds its plan with route-all routing and a single
persistent source bundle, so every routing draw produced an identical
forward: n_samples > 1 only multiplied recon compute, averaging n
identical losses. The field was 1 in every config; inline it
(plan construction unchanged, SPEC S11/S24 untouched).

Stored run configs carrying n_samples: 1 strip on load via the existing
removed-field shim on the class (same pattern as
use_sigmoid_parameterization); any other value is rejected. Checked-in
run yamls and the compile-probe generator drop the key.


Claude-Session: https://claude.ai/code/session_01KiMmzVhKUWbxbEz4R3Lkqm

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

* fix(checkpoint): coerce restored state onto reference FORMAT not just sharding

The prior device_put-onto-sharding was a confirmed no-op (StandardRestore already
honors the sharding spec — verified via an instrumented resume: restored CI-fn was
correctly ÷N). The real resume-OOM is a ÷1-scale ENTRY RELAYOUT on the first
resumed step: orbax-restored arrays carry a default memory layout that differs from
what the jitted step was compiled for, so the same executable that runs fresh at
~150GB materialized a ~103GB buffer on resume. The fresh-init reference is built by
the same XLA layout assignment as the step, so its `.format` (layout + sharding) IS
the step's expected input layout; coercing the restored tree onto it removes the
entry relayout. Unblocks durable auto-resume (transient crashes can now recover).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR

* docs: de-parameterise config-specific comments

Comments/docstrings stated one-run measurements (job IDs, ~6.6x peak,
~730k trace events), one-config numbers (d=4096, 513=512+label,
16GB->2GB at fsdp=8), or one-topology thresholds (dp>=64, "all 8 local
GPUs") as if they were code properties. Rewritten to state the general
mechanism; specifics dropped. Bridge task: config-comments-cleanup.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Dy5JU3zsCFxDk6ANLdtN9S

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

…the generic input edge (#923)

The residual-start removal (c64b73a / 3883362, 2026-06-24) deleted the
prefix/suffix machinery but left prose describing it in the agent-facing docs and
a few comments — which is why agents kept resurfacing the dead concept (CLAUDE.md
is injected into every session). Sweep:

- CLAUDE.md + param_decomp/CLAUDE.md: full token-input forward, no prefix build
  step, no `prefix_residual_fn` / `first_layer` / "suffix is multi-GB"; fix the
  DecomposedModel bullet (site_inputs -> read_activations, drop the stale
  `(frozen, vu)` framing).
- param_decomp/README.md: llama8b row described "residual-start suffix + Prefix
  harvest" and a deleted builder; now LlamaDecomposedModel full-model token-input,
  build_decomposed_lm / load_decomposed_lm_from_hf.
- gen_probe_config.py / lm/run.py / lm/load_run.py: comment rewords.
- test_generic_model_io.py: the dict-input test had quietly stopped testing the
  engine (it exercised only its own local prefix helper after the seam deletion).
  The synthetic model now consumes the dict batch DIRECTLY (feat_proj folded in
  as a frozen field), so the generic INPUT edge is exercised through the real
  make_train_step again. test_no_bake_invariant reuses the model and follows.
- Toys: delete tms_input_residual (zero callers — its "uniform harvest path"
  rationale died with the seam); resid_mlp_input_residual keeps its one live
  caller but drops the `prefix` param name + prefix framing.

Left intentionally: equivalence/stacked-parity fixture prose ("suffix weights")
describes the residual-fed goldens still xfailed pending embed-internal regen;
llama8b.py's frozen-prefix/suffix sub-scan comment is the live chunkwise-scan
meaning; string-prefix uses untouched.

Validated: test_generic_model_io 2 passed; test_no_bake_invariant passed;
tms+resid_mlp 42 passed / 2 skipped; basedpyright + ruff clean.


Claude-Session: https://claude.ai/code/session_01LzktY4CrJA2c2fGzgZRSiW

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

…tent (#925)

The shape-spelled scope literals (fresh-PGD mask_scope c/bc/bsc, PPGD scope
sc/bsc) are the canonical convention (SPEC §6 SCOPE), but the repo kept
authoring the legacy verbose names, so the stored-run aliases in configs.py
could never become dead. Migrate every repo-authored occurrence:

- param_decomp/configs/*.yaml: mask_scope shared_across_batch -> c,
  unique_per_datapoint -> bsc; scope type per_batch_per_position -> bsc,
  broadcast_across_batch -> sc. Header comments that name this file's own
  settings updated; upstream-torch provenance mentions kept verbatim.
- configs/compile_probe/gen_probe_config.py: emitted scope type -> bsc.
- SPEC.md: defaults table PPGD scope -> `sc`; S16 -> `bsc` (matches §6's
  canonical SCOPE row; no semantic change).
- nano_param_decomp/run.py: comments/docstring -> `bsc` with shape gloss.

The _LEGACY_*_ALIASES BeforeValidators in configs.py stay: 30 stored run
configs under PARAM_DECOMP_OUT_DIR/runs still carry the legacy names
(verified they parse to the same scopes through the aliases).

Validation: all 30 checked-in YAMLs parse via load_config to the intended
scopes; make type clean; test_config + test_recon_log_keys green.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>