This repository was archived by the owner on Jul 17, 2026. It is now read-only.
tests: pin full-fp32 matmuls in the torch↔JAX goldens (GPU TF32 trap)#970
Open
claude-spd1 wants to merge 1 commit into
Open
Conversation
XLA's f32 matmul default on GPU tensor cores is TF32 (~1e-3 rel). The
equivalence/parity goldens assert fp32 tolerances (RTOL 2e-4 / 1e-4, abs
1e-5), so an unpinned GPU run burns most of the tolerance budget on
precision noise: measured on H100, the ppgd term sits at rel 1.4e-4 of
its 2e-4 RTOL under TF32 vs 8e-8 pinned — room for a real port bug to
masquerade as TF32 noise (compile-audit finding M6). CPU CI never sees
this ("highest" is already the CPU default), which is why it stayed
latent.
- autouse `jax.default_matmul_precision("highest")` fixture in
tests/equivalence, tests/stacked_parity, and
tests/simple_mlp_equivalence (all three assert fp32 goldens); the same
pin in jax_equivalence.main(), the CLI path. Test-scope only: train
numerics stay bf16 by design, and the bf16 imp-min seam test is
deliberately not pinned.
- vendored_jax: `attn_implementation` dispatches on dtype — cuDNN flash
attention rejects fp32, so the fp32 goldens previously could not run
on GPU at all (NotImplementedError on H100). bf16/fp16 keep cuDNN;
fp32 takes the XLA composite. Dispatch-only; the production bf16 path
is unchanged.
Verified on H100: all three suites 17 passed / 1 xfailed (identical to
CPU), per-term golden rel err ≤ 8e-8 pinned vs up to 1.4e-4 unpinned.
Crew-Address: task/pin-golden-matmul-precision
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Stacked on #936 (base is its branch; will retarget to
feature/jaxwhen it merges). Two changes:jax.default_matmul_precision("highest")in every fp32-golden test harness —tests/equivalence,tests/stacked_parity, andtests/simple_mlp_equivalence(same fp32-golden class, same trap), as an autouse fixture per module, plus the same pin injax_equivalence.py main()(the CLI path). Test-scope only: train numerics stay bf16 by design, and the bf16 imp-min seam test is deliberately not pinned.vendored_jax.llama.attn_implementationnow dispatches on dtype. cuDNN flash attention rejects fp32, so the fp32 goldens previously could not run on GPU at all — every numeric test crashed withNotImplementedErroron H100. bf16/fp16 keep cuDNN (production path unchanged); fp32 takes the XLA composite. Without this, the precision pin is unreachable on GPU.Related Issue
Compile-audit finding M6 (board task
pin-golden-matmul-precision, promoted from the 2026-07-10 JAX compile audit).Motivation and Context
XLA's f32 matmul default on GPU tensor cores is TF32 (~1e-3 rel). The goldens are the only remaining torch↔JAX parity oracle (torch is retired at tag
torch-oracle) and assert fp32 tolerances. CPU-only CI never sees the difference — "highest" is already the CPU default — but a GPU dev-box run does. #936 removes the_PENDING_REGENxfail that currently masks this, so the pin should land with/right after it.Measured on H100 against the committed goldens (per-term rel err, RTOL 2e-4):
Unpinned, TF32 noise alone eats 70% of ppgd's tolerance budget — that's the "real port bug written off as TF32 noise" masquerade the audit flagged. Pinned, headroom is four orders of magnitude.
How Has This Been Tested?
basedpyrightandruffclean.Does this PR introduce a breaking change?
No. Production numerics untouched: the bf16 attention path still uses cuDNN; the dtype dispatch only enables the previously-crashing fp32 GPU path.
Crew-Address: task/pin-golden-matmul-precision
🤖 Generated with Claude Code