feat(optim): bf16 Adam m/v storage for the V/U + CI-fn optimizers (moments_dtype)

[ocg-goodfire]

What

New OptimizerConfig.moments_dtype: float32 | bfloat16 (default float32 — bit-unchanged, pinned by a new test): bfloat16 stores BOTH Adam moments of the components (V/U) and CI-fn optimizers in bf16 while the fp32 masters stay untouched (SPEC N1). Item #5 of the lower-precision roadmap (bridge task lp-bf16-adam-moments-vu; follows lp-bf16-source-ab / #938 which validated the same technique on the PPGD source Adam).

Implementation is deliberately tiny: optax's mu_dtype covers the first moment natively, and its scale_by_adam update ends with nu = cast_like(nu, state.nu) — so the second moment's storage dtype is whatever it was initialized with. _with_nu_dtype wraps only the init. Update arithmetic still promotes to fp32; bf16 is storage rounding only. m/v inherit the masters' ZeRO-1 ÷N sharding via zeros_like/astype; the #927 fsdp-major reconstruct gathers masters only, so sharding/layout is untouched (confirmed by the memprobe: temp identical, 96.9 GiB both arms).

Validation

Memory (full32L b128/dp32 memprobes, 4 scavenge nodes each):

• resident argument 49.50 → 43.70 GiB = −5.80 GiB/GPU (exactly m+v = ⅓ of the 3×-fp32-master optimizer block)
• runtime peak 146.45 → 140.38 GiB (limit 164.08); runs p-841cc252 (fp32) / p-f9ddec8a (bf16), jobs 170794/170795

Quality (same-seed 10k-step pile-4L PPGD-bsc A/B, runs p-6a9bc1ce fp32 / p-4731b0d9 bf16):

• adversary effectiveness (train PPGD recon): within ±5%, no drift/collapse
• fresh-PGD 20-step eval: bf16 equal or better at every point (−2..−9% pre-spike)
• masked CE/KL (ci/rounded/stoch): ±5%, no systematic degradation; CI-L0 within +3% through step 8k
• faithfulness: converges to identical (+0.0% from step 2600 on)
• both arms share the known end-of-run p-anneal spike (same artifact as the bf16 PPGD source: quality A/B + full32L memory cash-in configs (lp-bf16-source-ab) #938 A/B; bf16 degrades less on eval at the spike)
• ckpt moment probe: v==0 fraction 0.000% in both arms (medians 1e-11..1e-14 ≫ the exponent floor, which bf16 shares with fp32); bf16 arm ckpts store m/v as bfloat16 end-to-end
• live memory corroboration on the 4L arms: −0.90 GB/rank peak

Tests: new tests pin (a) bf16 m/v dtype persists across updates, (b) bf16 trajectory tracks fp32 within storage-rounding scale, (c) fp32 default is bit-identical to plain optax.adamw. Optim-parity/checkpoint/equivalence/stacked-parity suites green at default + 4 sim devices.

The A/B + memprobe config pairs are included (derived from pile_ppgd_bsc.yaml / llama8b_full32L_HSDP_b128_dp32.yaml).

🤖 Generated with Claude Code

https://claude.ai/code/session_01CbD62cDit2QPQBHUT3fzkk