tests: pin full-fp32 matmuls in the torch↔JAX goldens (GPU TF32 trap)

param_decomp/tests/equivalence/jax_equivalence.py

@@ -210,7 +210,10 @@ def per_site(prefix: str) -> dict[str, jnp.ndarray]:
 def main() -> None:
     f = dict(np.load(HERE / "fixtures.npz"))
     ref = json.loads((HERE / "torch_reference.json").read_text())
-    jaxv = compute_jax_terms(f)
+    # Same full-fp32 matmul pin as the pytest fixture: the GPU f32 default is TF32
+    # (~1e-3 rel), which swamps RTOL. No-op on CPU.
+    with jax.default_matmul_precision("highest"):
+        jaxv = compute_jax_terms(f)
     print(f"{'term':6} {'jax':>16} {'torch':>16} {'rel_err':>12}  ok")
     all_ok = True
     for term in ("faith", "imp", "stoch", "ppgd"):

param_decomp/tests/equivalence/test_equivalence.py

@@ -35,6 +35,7 @@
 import json
 from pathlib import Path
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 import pytest
@@ -50,6 +51,16 @@
 ATOL = 1e-5
 
 
+@pytest.fixture(autouse=True)
+def full_fp32_matmuls():
+    """XLA's f32 matmul default on GPU tensor cores is TF32 (~1e-3 rel), which swamps
+    RTOL — unpinned GPU runs fail spuriously, or worse, teach us to read real torch↔JAX
+    divergence as precision noise. Pin full fp32 (a no-op on CPU CI, where "highest" is
+    already the default). Test-scope only: train numerics are bf16 by design."""
+    with jax.default_matmul_precision("highest"):
+        yield
+
+
 def _load_fixtures() -> dict[str, np.ndarray]:
     return dict(np.load(HERE / "fixtures.npz"))
 

param_decomp/tests/simple_mlp_equivalence/test_torch_equivalence.py

@@ -15,6 +15,7 @@
 import json
 from pathlib import Path
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 import pytest
@@ -31,6 +32,15 @@
 REAL_CACHE_DIR = Path("/mnt/data/artifacts/mechanisms/param-decomp/pretrain_cache/spd-t-9d2b8f02")
 
 
+@pytest.fixture(autouse=True)
+def full_fp32_matmuls():
+    """XLA's f32 matmul default on GPU tensor cores is TF32 (~1e-3 rel), which swamps
+    the absolute tolerances below — the goldens are full-fp32 torch logits. No-op on
+    CPU CI, where "highest" is already the default. Test-scope only."""
+    with jax.default_matmul_precision("highest"):
+        yield
+
+
 def _max_abs_diff(a: Array, b: np.ndarray) -> float:
     return float(jnp.max(jnp.abs(a - jnp.asarray(b))))
 

param_decomp/tests/stacked_parity/test_stacked_parity.py

@@ -22,6 +22,7 @@
 from pathlib import Path
 
 import equinox as eqx
+import jax
 import jax.numpy as jnp
 import numpy as np
 import optax
@@ -62,6 +63,16 @@
 RTOL = 1e-4
 ATOL = 1e-5
 
+
+@pytest.fixture(autouse=True)
+def full_fp32_matmuls():
+    """XLA's f32 matmul default on GPU tensor cores is TF32 (~1e-3 rel), which swamps
+    RTOL — the committed goldens were pinned with full-fp32 matmuls. No-op on CPU CI,
+    where "highest" is already the default. Test-scope only."""
+    with jax.default_matmul_precision("highest"):
+        yield
+
+
 STABLE_FIXTURE_METRIC_KEYS = (
     "total", "faith", "imp", "stoch", "ppgd", "p_imp", "src_lr",
     "grad_norms/summary/components", "grad_norms/summary/ci_fns", "grad_norms/summary/total",

vendored_jax/llama.py

@@ -112,16 +112,20 @@ def repeat_kv(x: Float[Array, "b kvh t hd"], n_rep: int) -> Float[Array, "b h t
     return x.reshape(b, kvh * n_rep, t, hd)
 
 
-def attn_implementation() -> Literal["cudnn", "xla"]:
-    """cuDNN flash attention on GPU; the XLA composite elsewhere (CPU tests)."""
-    return "cudnn" if jax.default_backend() == "gpu" else "xla"
+def attn_implementation(dtype: jnp.dtype) -> Literal["cudnn", "xla"]:
+    """cuDNN flash attention on GPU for the half dtypes it supports; the XLA composite
+    elsewhere. cuDNN rejects fp32, so the fp32 paths (CPU tests, the GPU-run
+    equivalence goldens) take the XLA composite."""
+    if jax.default_backend() == "gpu" and dtype in (jnp.bfloat16, jnp.float16):
+        return "cudnn"
+    return "xla"
 
 
 def causal_sdpa(q: Array, k: Array, v: Array) -> Array:
     # q,k,v: (B, H, T, hd); jax.nn.dot_product_attention takes (B, T, H, D).
     qt, kt, vt = (a.transpose(0, 2, 1, 3) for a in (q, k, v))
     out = jax.nn.dot_product_attention(
-        qt, kt, vt, is_causal=True, implementation=attn_implementation()
+        qt, kt, vt, is_causal=True, implementation=attn_implementation(q.dtype)
     )
     return out.transpose(0, 2, 1, 3)