feat(ci-health): CI-fn training-health telemetry (eval/ci_health/*)

[ocg-goodfire]

What

The CI fn is itself a large trained transformer (~10-31B params at 8B scale), but the only in-loop visibility into its training health was per-leaf grad norms. This adds the usual big-model instrumentation as ~220 scalars under eval/ci_health/* on the fast eval cadence, config-gated by a new CIFnHealth entry in eval.metrics (bridge task ci-function-visibility).

Metric families (param_decomp/ci_health.py, new)

• weights/{frob,sigma_max,stable_rank}/<block>/<matrix> — per (block, matrix ∈ wq/wk/wv/wo/w1/w2, + in_proj, out_heads): whole-leaf Frobenius norm (same convention as grad norms → update/weight ratios overlay), spectral norm via batched power iteration over the chunk axis (never full SVD — minutes-slow at 32 chunks × 4096²), stable rank ‖W‖²_F/σ²_max ("weight ranks"). Off the fp32 masters, once per eval pass.
• act/block{i}/* — from a new eval-only instrumented forward (ChunkwiseTransformerCIFn.telemetry): post-RoPE q/k RMS ("QK norm scales"), attention max/RMS logit + softmax entropy on a ≤64-strided query subsample (logit growth / entropy collapse — the telemetry path materializes the score slab flash never does, hence the subsample), in_proj / MLP-hidden / post-block residual RMS. Chunk-reduced: max for logit-max, mean + worst-chunk _min for entropy, mean otherwise.
• logits/* — CI-logit distribution: mean/std + fractions in the lower-squashing zero-gradient (logit > 1) and leak (logit < 0) regions (gradient-starvation signal).
• components/* — mode-collapse scalars off the per-site mean-CI-per-component vector: participation fraction (Σm)²/(C·Σm²) (mean + min over sites) and dead-component fraction (mean + max).

Design notes

• CIBlock.instrumented deliberately duplicates __call__'s forward math (~15 lines) instead of a flagged shared path: the train step must not carry — or DCE-trust away — the explicit score einsum inside the checkpointed scan. Telemetry logits are pinned bit-identical to __call__ in tests; training numerics untouched, no SPEC invariant involved (S28-adjacent, JAX-only keys with no torch analog).
• Power iteration is a fori_loop over a replicated-once fp32 copy per leaf: the first cut (unrolled ×32 over ~50 sharded leaves) demanded ~50k VMM comm-buffer allocations and OOM'd a B200 at dp=8 (job 167430) — the second commit is that fix.
• eval/ci_health/* is dropped from the console line (generalized the grad-norms verbosity filter); the full set still rides to wandb + jsonl.
• Chunkwise-transformer only: the lab _eval conversion refuses the metric for MLP CI fns.

Enabled in

pile_pgd1, pile_ppgd_bsc, llama8b_full32L_HSDP_b32_dp32, llama8b_full32L_HSDP_b128_dp32, llama8b_full32L_seq512_b128_dp128.

Validation

• param_decomp/tests/test_ci_health.py: telemetry↔__call__ logit equality (bitwise), power-iter σ vs full SVD (rtol 1e-2 — fresh-init spectra are near-degenerate, power iter under-estimates ~1%), scalar ranges. Green at default and 4 sim devices.
• Full make test: 460 passed, 7 skipped, 11 xfailed.
• make type clean.
• GPU e2e smoke (pile-4l config, 2 steps, eval.every=1, 1×B200, run p-c1a10e57): 223 eval/ci_health/* keys logged per eval pass, all finite, values textbook at init (attn entropy ≈ log T, stable ranks in the random-matrix ballpark, logit saturation fractions plausible).

🤖 Generated with Claude Code