perf(llama8b): batch weight_deltas per kind (224 site matmul groups → 7)

[claude-spd1]

Description

llama8b.weight_deltas Python-looped over all 224 sites, emitting a per-site fp32 W − (V@U).T matmul group into the loss graph (train-step faith term + faith-warmup step). Per-kind dims are uniform across layers (already asserted by _per_kind_dims) and the frozen weights are layer-stacked, so it now computes one batched fp32 matmul per kind and returns {kind: [n_sites, d_out, d_in]} stacks (rows layer-ascending; per_kind_delta_index maps site → (kind, row)).

SPEC N2/S17 pin the math (fp32 deltas, global mean of squared entries), not the grouping: faithfulness_loss was already grouping-agnostic; its annotation and the DecomposedModel.weight_deltas contract now say so. Toys and llama_simple_mlp keep per-site returns (valid grouping, small site counts).

Second commit: jax.checkpoint at the two faith call sites — the loss reduce's backward needs the deltas (grad = 2Δ/N), and without remat XLA stores every per-kind stack from forward to backward (measured +3.3GiB/rank static temp at dp32). With remat they're recomputed (one matmul per kind) and the change is memory-neutral.

Motivation and Context

2026-07-10 JAX compile audit finding M4 (ledger: lore 2026-07-10--jax-compile-audit-remaining-findings); bridge task batch-weight-deltas-per-kind.

Measurements

Isolated region (faith-warmup step AOT at production 8B shapes/224 sites, CPU backend): optimized-HLO dots 672 → 21, lines 41590 → 27506 (−34%).

GPU A/B at production dp32 (full32L PROFILE config + profile.mem_profile, isolated compile caches; old arm = eb786f7d9):

	old	new (batched)	new (batched + remat)
static temp / rank	47.55 GiB	50.82 GiB	47.09 GiB
jit(step) compile	~223–240 s	—	~212–227 s

jit(step) compile time is flat within noise — consistent with the repo's prior "compile is ~83% priority-fusion, flat across graph structure" finding. The win is the graph itself (651 fewer matmul groups + 224→7 dict entries in this region), not the dp32 step-compile wall clock; the warmup-step graph shrinks the same way. (The A/B runs OOM in the 2-real-step phase of mem_profile in both arms — pre-existing for this config on H100-80GB, arm-independent.)

Step-time impact should be nil (the faith term is a negligible fraction of step FLOPs; remat adds 7 matmuls to the backward) but was not measurable in this config for the same OOM reason.

How Has This Been Tested?

• New test_weight_deltas_match_per_site_reference (3 site sets: all kinds × all layers, partial MLP layer range, heterogeneous per-site C): per-site delta values and V/U faithfulness grads bit-identical to the per-site loop; loss scalar differs only by fp32 summation regrouping (224 → 7 partial sums; feeds no gradient, so trajectories are unchanged).
• test_llama8b.py, stacked_parity, test_generic_model_io.py, test_eval.py, test_llama_simple_mlp.py green; make check clean.

Does this PR introduce a breaking change?

weight_deltas consumers that indexed by site name must go through per_kind_delta_index (in-tree: two test files, updated here). The engine seam (faithfulness_loss) is unchanged.

Crew-Address: task/batch-weight-deltas-per-kind

🤖 Generated with Claude Code