AI-Hypercomputer / AI-Hypercomputer/maxtext
Qwen3-Next GDN hardcodes `Precision.HIGHEST`, ignoring `config.matmul_precision` — is this intentional?
- Vorherrschende Sprache
- Python
- Sterne
- 2.4k
- Forks
- 607
- Ø Merge
- 2 T. 19 Std.
- Gemergte PRs (30 T.)
- 158
Beschreibung
### Summary
The Qwen3-Next Gated DeltaNet path in `src/maxtext/models/qwen3.py` hardcodes `jax.lax.Precision.HIGHEST` in
9 places and never reads `config.matmul_precision`. Every other attention layer in the repo honors that config
(43 usages in `src/maxtext/layers/`). On TPU this means the GDN matmuls always run **6 bf16 MXU passes**
instead of the **1** implied by the default config.
I'd like to ask whether this is deliberate before sending a patch, because the git history suggests it **was**
a conscious choice at one point — and there is existing evidence in this repo that bf16 can destabilize the
delta-rule recurrence. Details below, including the evidence *against* changing it.
Verified at `e190ed60` (current `main` at time of writing).
### The inconsistency
`src/maxtext/configs/base.yml:128`:
```yaml
matmul_precision: "default"
```
Honored across the codebase — 43 usages in `src/maxtext/layers/`, e.g. `layers/attention_mla.py:348`:
```python
logits = jnp.einsum("bthd, bsd -> btsh", q, k, precision=self.config.matmul_precision)
```
Not honored by GDN — `src/maxtext/models/qwen3.py`, 9 sites across all 3 entry points:
```
106 prec = jax.lax.Precision.HIGHEST # naive_jax_chunk_gated_delta_rule (@58)
149 prec = jax.lax.Precision.HIGHEST # └─ scan_body (@146)
256 S = jnp.matmul(k_beta, k_c.swapaxes(-1,-2), precision=jax.lax.Precision.HIGHEST)
# jax_chunk_gated_delta_rule (@183)
275 u_chunks = jnp.matmul(A, v_beta.astype(jnp.float32), precision=jax.lax.Precision.HIGHEST)
279 w_chunks = jnp.matmul(A, k_beta_g, precision=jax.lax.Precision.HIGHEST)
303 prec = jax.lax.Precision.HIGHEST # └─ scan_body (@301), inter-chunk scan
407 v_prime = jnp.matmul(k_cumdecay, state, precision=jax.lax.Precision.HIGHEST)
# jax_ar_gated_delta_rule (@363), decode path
413 attn_inter = jnp.matmul(q_g, state, precision=jax.lax.Precision.HIGHEST)
420 jnp.matmul(k.astype(jnp.float32)[...,None], v_new[...,None,:], precision=jax.lax.Precision.HIGHEST)
```
Notably, the call sites already forward other config values into these same functions — `matmul_precision` is
simply absent from the list (`qwen3.py:878`):
```python
return jax_chunk_gated_delta_rule(
query=q, key=k, value=v, g=g_val, beta=beta_val,
chunk_size=cfg.gdn_chunk_size,
initial_state=init_h,
use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
compute_dtype=cfg.dtype,
)
```
(All 3 call sites: 839 AR/decode, 878 sharded, 892 non-sharded.)
### What the history suggests (why I think this may be intentional)
This looks less like an oversight than like a rationale that got separated from its code:
- `0e03921` (2025-10-20) introduced the reference impl with an explicit comment:
```python
# NOTE: Precision set to HIGHEST for numerical accuracy.
prec = jax.lax.Precision.HIGHEST
```
- `060fecf` (2026-02-13, the optimized GDN impl) carried `Precision.HIGHEST` into the new chunked fast path
**and removed that NOTE comment.** At `main` today, no rationale for HIGHEST survives anywhere in the file.
So the precision choice appears deliberate in origin, but it was never reconciled with `matmul_precision`, and
the reasoning is no longer discoverable in the source. That's really what I'm asking about.
### Measurements
**All numbers are from one configuration and should not be read as general.**
TPU v7x (Ironwood), JAX 0.10.2 / libtpu 0.0.42.1, isolated GDN kernel, fwd+bwd,
`B=4, S=4096, H=32, K=V=128, gdn_chunk_size=64, use_qk_norm_in_gdn=True`:
| precision | fwd+bwd | vs HIGHEST |
|---|---|---|
| HIGHEST (current) | 24.12 ms | 1.00× |
| HIGH | 18.82 ms | 1.28× |
| DEFAULT | **13.92 ms** | **1.73×** |
Accuracy vs an fp32 reference (8 seeds × 4 decay strengths, worst case), at `default`:
output cos **0.999984**, final-state cos **0.999990**, min gradient cos **0.999952** (over q, k, v, g, beta).
Long horizon: recurrent state chained 30× (**122,880 tokens**, weak decay = worst case for compounding) —
state cos **0.999992, flat across all 30 chains**, no drift.
Full-model 500-step seed-locked A/B (Qwen3-Next-80B-A3B MoE, **synthetic data**, HIGHEST vs DEFAULT):
mean |Δloss| **0.0012** over steps ≤100 and **0.0000** over steps ≥400; identical final loss; no NaN.
The gap *shrinks* with horizon rather than compounding.
### Evidence against changing this
I want to flag the counter-evidence myself, since it's material:
- **@bzantium's #4348** reports that bf16 "breaks `(I+S)A = I` enough to destabilize the recurrence" at
`gdn_chunk_size=128`, and works around it with manual bf16x3 operand splitting.
- My results are at **chunk 64** with `use_qk_norm_in_gdn=True`, where I measured **max|S| ≈ 0.2** — the
`(I+S)` inverse is very well conditioned there, which likely explains why 1 bf16 pass suffices *for me*.
That conditioning argument cuts both ways: it implies nothing about other chunk sizes or without qk-norm.
- **Both results can be true.** My data does not show bf16 is safe in general.
- My convergence A/B used **synthetic data**, which approaches a memorization floor early; it stresses
horizon/compounding but is not real-data optimization dynamics.
- On my own pipeline the end-to-end payoff is small (GDN ≈16% of step, which is comms-bound). The kernel win
should matter far more on GDN-dominated dense configs — where #4348 measured +31%.
Taken with the removed NOTE comment above, it's entirely possible HIGHEST is load-bearing and should stay.
### Question
Is the hardcoded `HIGHEST` deliberate (delta-rule numerical stability), or an oversight that predates
`matmul_precision` being plumbed through the other layers?
Depending on your answer I'm happy to send either:
1. **Honor the config** — plumb `matmul_precision` through the GDN entry points and forward
`cfg.matmul_precision` at the 3 call sites, following the existing `gdn_chunk_size` / `use_qk_norm_in_gdn` /
`dtype` convention. Fixes the inconsistency, but **does change default numerics** (HIGHEST → `"default"`).
2. **Opt-in knob** — add `gdn_matmul_precision` defaulting to `"highest"`. Zero behavior change; the speedup is
available to anyone who validates it for their config.
I'd lean toward (1) as the more correct fix, but given the stability concern and the history above I don't want
to change default numerics for everyone on the strength of a single-config result. Happy to follow your
preference.
A related sub-question: should `naive_jax_chunk_gated_delta_rule` (@58) be plumbed too, or stay pinned at
HIGHEST deliberately as the numerical reference? `tests/unit/qwen3_next_vs_reference_test.py:643,1129` already
passes `matmul_precision=highest` explicitly — which the GDN hardcode currently makes a no-op for this path.
This touches the same file as #4348, so I'm happy to coordinate ordering / rebase behind it.
Environment: JAX 0.10.2, libtpu 0.0.42.1, TPU v7x.
Beitragsleitfaden
Bewertung
Dieses Issue wurde noch nicht bewertet.