fix(eval): make the slow-eval tier work on multi-host GPU

[danbraunai-goodfire]

Description

The slow/plot eval tier crashed on multi-host GPU (>1 process) at the first slow eval (step 1000). Two independent bugs, both fixed here:

1. cuDNN-fp32 attention (param_decomp/ci_fn.py). The CI-transformer attention always requested the cuDNN flash impl via attn_implementation(), which only accepts fp16/bf16/fp8. Training and the fast eval run the CI fn under bf16 autocast, but the slow tier deliberately reads it out in fp32 (slow_eval.py: "a fp32-CI-fn readout"). On GPU that raised:

NotImplementedError: Q must be fp16/bf16/fp8_e4m3fn/fp8_e5m2, got float32

Fix: pick the impl by dtype — cuDNN for fp16/bf16 (training + fast eval, unchanged), the XLA impl for fp32 (slow eval). The XLA composite materializes (B,H,T,T), but that only matters for the long-sequence target forward, not this seq-512 CI transformer.

2. np.asarray on a process-sharded array (param_decomp/slow_eval.py). accumulate_site_reductions did np.asarray(flat_lower[site]) / flat_logits[site], which keep the dp-sharded batch axis. On >1 process those span non-addressable devices and np.asarray raises:

RuntimeError: Fetching value for `jax.Array` that spans non-addressable (non process local) devices ...

(The density/sum reductions are all-reduced → addressable, so they were fine; only the raw per-position histogram sample is sharded.) Fix: gather across processes with jax.experimental.multihost_utils.process_allgather(..., tiled=True).

Motivation and Context

Together these crash every multi-GPU LM decomposition run at its first slow eval — including the in-flight baseline p-0ff8e5d3, which reached step 1000 then died from the gRPC coordination cascade these exceptions trigger on the other ranks. CPU tests never caught either: attn_implementation() returns "xla" off-GPU, and single-process np.asarray is addressable.

How Has This Been Tested?

• make type / basedpyright clean.
• param_decomp/tests/test_slow_eval.py: 26 passed (MPLBACKEND=agg; the data-path + gather tests pass, incl. accumulate_site_reductions).
• On btdr (2 nodes / 16 GPU): pre-fix the run aborts at the step-1000 slow eval (first with the Q must be fp16… error, then the non-addressable-fetch RuntimeError); with both fixes the slow tier completes and training continues. [validation in progress]

Does this PR introduce a breaking change?

No. Training and fast-eval paths (bf16, addressable reductions) are unchanged; only the fp32 attention path switches to the XLA impl and the sharded histogram sample is gathered across processes.

🤖 Generated with Claude Code

[danbraunai-goodfire]

Validated on btdr (2 nodes / 16 GPU): with both fixes, a run clears the step-1000 slow eval (previously the crash point) and continues training — [eval @ 1000] and [eval @ 2000] both fire, slow-tier metrics (eval/slow/loss/CIHiddenActsReconLoss/*, etc.) log, and there are no Q must be fp16 / non-addressable-fetch errors (run p-a067a8d0).

[danbraunai-goodfire]

superseded by #905