[kru2710shna]

…tedDeltaNet

During fine-tuning, GatedDeltaNet runs gated_delta_update with use_kernel=not self.training, so training fell back to gated_delta_ops -- a per-token recurrence whose backward scales with sequence length T. The Metal kernel has no VJP, so that loop was the only differentiable path, and step time grew with the number of trained GatedDeltaNet layers (#1956).

Adds gated_delta_chunked: a chunkwise UT (WY) transform mathematically equivalent to gated_delta_ops that cuts the sequential dependency from T to T/chunk_size. The per-chunk inverse uses a Newton-Schulz iteration (exact in ceil(log2(C)) steps for the unipotent I + strictly-lower matrix), which stays on the GPU and is differentiable by plain autograd -- no custom VJP, no CPU round-trip. gated_delta_update routes training to this path for a scalar gate with no padding mask; decode, masks, vector gates and inference are unchanged.

Correctness vs gated_delta_ops (forward and gradients) matches to ~1e-6 across GQA, ragged tails, decay/beta/value extremes and a full-sequence single chunk. Requires normalized keys (the normalize_qk contract).

Measured on M-series: ~20x faster at the Qwen3-Next recurrence shape; a full-model Qwen3.5 training step (6/8 GatedDeltaNet layers) is 7.5x faster at T=64 and 9.3x at T=128, growing with T, and trains at T=1024 where the per-token path runs out of memory.

Adds six tests to tests/test_gated_delta.py.

Fixes #1956

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure:

[lgh-hub]

Hi, thanks for this PR — the chunked training path is exactly what's needed to make qwen3_5-family models trainable on MLX. I ran the test suite and real LoRA training on an Apple-silicon machine (MLX 0.32.4) and ran into three numerical issues, all reproducible and all fixable with a small patch (~90 lines). Sharing in case it's useful.

Root cause 1: TF32 is enabled by default for fp32 matmuls

On Apple GPUs MLX runs fp32 matmuls through TF32 unless MLX_ENABLE_TF32=0 (read at import, no runtime API). A 64×128 fp32 matmul measures ~8.8e-4 relative error on GPU vs ~3.8e-7 on CPU/numpy. gated_delta_ops is elementwise+reduction (true fp32) while gated_delta_chunked is all matmuls, so test_chunked_matches_ops_* fails with a systematic ~1.8e-3 offset — which is what most of the 44 failing tests show. Setting the env var before importing mlx fixes them all (14 tests + 75 subtests pass; tolerances unchanged).

Root cause 2: Newton–Schulz _tri_inv diverges for heavily loaded chunks

For parallel-ish keys with allow_neg_eigval (‖N‖∞ up to ≈106 in the stress tests), the NS partial sums Σ(−N)^m overflow/cancel catastrophically: measured inverse error 1e3–1e13. Replacing _tri_inv with a block-Schur recursive inversion (A=[[A11,0],[A21,A22]] ⇒ A⁻¹=[[X11,0],[−X22·A21·X11,X22]], with an 8×8 NS base case plus 2 Newton polish steps) keeps everything pure-matmul/GPU/autograd-friendly and exact, with intermediates bounded by the true inverse. Killer regimes drop to 1e-3–1e-2 and normal regimes are unchanged or better.

Root cause 3: decay_ratio dead-branch overflow → NaN gradients in real training

This one only shows up in real training, not in the test suite (random test gates don't decay hard enough): when per-token gates get small (real model, some heads: g down to ~0.06 ⇒ Σlog g over a 64-chunk ≈ −170), the upper-triangle entries of decay_ratio = exp(lg_i − lg_j) (i<j, masked away by where in the forward) evaluate to inf. The forward is fine, but in the backward, exp′(x)=exp(x) meets the zero cotangent masked in by where ⇒ 0×inf = NaN, poisoning every in_proj_* gradient. One-line fix: clamp the exponent with mx.minimum(·, 0.0) — for the used lower-triangle (i≥j) the exponent is always ≤0, so this is mathematically a no-op; dead branches become exp(0)=1. (Also added a defensive mx.maximum(g, 1e-8) before log.)

Validation

• pytest tests/test_gated_delta.py: 14 passed + 75 subtests (was 44 failures), tolerances unchanged.
• vjp chunked-vs-ops: fp32 dq/dk/dv/dg/dβ all 3–5e-7. bf16 worst case dk ≈ 3.2e-3 (input quantization floor — reporting as-is rather than hiding it).
• Real bf16 LoRA on a 9B qwen3_5 model, seq 4096, grad checkpointing: gradients finite every step (470 params had NaN before); 200 steps: val 3.56 → 2.84; a full 2-epoch run (6.5k steps) finished with no NaN.
• TF32 and training: with fixes 2 and 3 applied but TF32 left enabled, the same 200-step run matches the TF32-off run within 0.002 val loss at every checkpoint, and throughput was about the same either way (~300–350 tok/s at seq 4096 with grad checkpointing on this machine). So TF32 is what breaks the bit-parity tests, not trainability, and disabling it did not cost measurable speed here. Fixes 2 and 3 are the ones that matter for training.

Minimal patch

diff --git a/mlx_lm/__init__.py b/mlx_lm/__init__.py
@@
 import os
+
+# Apple GPUs run fp32 matmuls through TF32 by default (~1e-3 relative error),
+# which breaks gated_delta_chunked vs gated_delta_ops parity. Must be set
+# before mlx.core is imported; respected if the user overrides it.
+os.environ.setdefault("MLX_ENABLE_TF32", "0")

diff --git a/mlx_lm/models/gated_delta.py b/mlx_lm/models/gated_delta.py
@@ def _tri_inv(A):
-    C = A.shape[-1]
-    n_iter = max(1, math.ceil(math.log2(C))) if C > 1 else 1
-    eye = mx.eye(C, dtype=A.dtype)
-    X = eye + mx.zeros_like(A)
-    for _ in range(n_iter):
-        X = X @ (2.0 * eye - A @ X)
-    return X
+    # Block-Schur recursive inversion of unit lower-triangular A (pure
+    # matmuls, GPU-friendly, autograd-friendly). NS from I diverges when
+    # ‖N‖∞ is large (partial sums overflow/cancel); the recursion keeps
+    # intermediates bounded by the true inverse. Small NS base case + 2
+    # Newton polish steps for accuracy.
+    C = A.shape[-1]
+    def _rec(M):
+        n = M.shape[-1]
+        if n <= 8:
+            eye = mx.eye(n, dtype=M.dtype)
+            X = eye + mx.zeros_like(M)
+            for _ in range(max(1, math.ceil(math.log2(n)))):
+                X = X @ (2.0 * eye - M @ X)
+            return X
+        h = n // 2
+        X11 = _rec(M[..., :h, :h]); X22 = _rec(M[..., h:, h:])
+        X21 = -X22 @ (M[..., h:, :h] @ X11)
+        top = mx.concatenate([X11, mx.zeros_like(X21).swapaxes(-1, -2)], -1)
+        bot = mx.concatenate([X21, X22], -1)
+        return mx.concatenate([top, bot], -2)
+    X = _rec(A)
+    eye = mx.eye(C, dtype=A.dtype)
+    for _ in range(2):
+        X = X @ (2.0 * eye - A @ X)
+    return X

@@ gated_delta_chunked, inside the per-chunk loop
-        gc = g[:, start:end].swapaxes(1, 2).astype(mx.float32)
+        gc = mx.maximum(g[:, start:end].swapaxes(1, 2).astype(mx.float32), 1e-8)
@@
-        decay_ratio = mx.exp(log_gamma[..., :, None] - log_gamma[..., None, :])
+        # Lower triangle (i>=j) always has exponent <= 0, so min(., 0) is a
+        # no-op where the ratio is used; the dead upper branch used to
+        # overflow to inf and produce 0*inf = NaN in the backward.
+        decay_ratio = mx.exp(mx.minimum(
+            log_gamma[..., :, None] - log_gamma[..., None, :], 0.0))

Repro sketch for root cause 3: build a chunk with gates drawn from exp(uniform(-5.4, 0)) (so Σlog g ≈ −170 over 64 tokens), normalized keys, beta in [0.85, 1]; forward is finite but mx.vjp of gated_delta_chunked(...).sum() returns NaN grads for q/k — after the clamp it's finite and matches gated_delta_ops.

Thanks again — happy to iterate if any of this is useful.

[kru2710shna]

Let me take a look and see what can be done in this case !!

[kru2710shna]

Thanks again @lgh-hub — really appreciate the hardware testing and the detailed root-cause write-up.

Root cause 3 (decay NaN) was real and I’ve pushed a fix (f824cc1): the masked upper triangle of decay_ratio overflows to inf under hard gate decay, and its backward poisoned the gradients. I clamp the exponent with minimum(., 0.0), a no-op on the used lower triangle + a maximum(g, 1e-8) guard before log. Added a regression test that reproduces the NaN on the unclamped code (fails before, passes after).

Thank you for catching this — it’s the one that actually breaks real training.

On root causes 1 (TF32) and 2 (_tri_inv divergence): I couldn’t reproduce either on my setup (M-series, MLX 0.32.4). My suite passes with TF32 at its default, and the chunked path matches gated_delta_ops to ~1e-6 forward+backward including allow_neg_eigval; the NS inverse and a block-Schur inverse I tried agree to ~1e-6 everywhere fp32 can represent the result (they only diverge together at ‖N‖ large enough that the true inverse overflows fp32, which the normalize_qk contract keeps us well clear of).

Could you share the exact shapes/seed/config where you saw the 1e13 _tri_inv error and the 44 TF32 failures? If there’s a real device-dependent regime, I’d rather add the block-Schur inverse + TF32 guard with a test that actually reproduces it than ship a change I can’t trigger.

[lgh-hub]

I’m not a programmer, just an AI enthusiast. The issue analysis, patch, and text in my previous comment were actually all completed by an AI coding assistant (Claude). I apologize for not clarifying this earlier, and I realize this violated your AI usage policy.
Since I don't have the capability to fully understand or guarantee the mathematical details in the patch, I won't be submitting a PR.
Regarding the reproduction conditions you asked about, the AI concluded that the inversion issue is hardware-independent. The original algorithm inherently fails under heavy matrix loads, and this happens even with double precision. There is a specific set of inputs without random numbers that can stably reproduce the issue (block size 64, beta=1, g=0.95, with a residual of ~1e8). The TF32 issue only occurs on Apple GPUs, affecting only the benchmark comparisons, not the training itself.
You are welcome to verify this yourself, or simply ignore it. If you find it valuable, please feel free to adopt or adapt it directly—no need to credit me.

[kru2710shna]

@lgh-hub
thank you, you were right. I reproduced your deterministic case (identical keys, g=0.95, block 64) and confirmed it: the old Newton-Schulz inverse reaches ~2.2e9 error there, because NS from the identity diverges when A is near-singular even at modest ‖N‖ — not an ‖N‖-magnitude issue as I’d first assumed. I’ve replaced _tri_inv with a block-Schur recursion (pure matmul, GPU + autograd, no custom VJP), which is exact (~1e-14) on your case, unchanged on well-conditioned inputs, and I pinned your exact construction as a deterministic regression test (5f40854). Both methods still lose accuracy only once the true inverse itself overflows fp32, which normalized keys keep us clear of in practice — so this closes the realistic gap cleanly.

On TF32 (root cause 1): agreed it’s a benchmark-comparison artifact on Apple GPUs, not a training issue — my suite passes with TF32 at its default, so I’ve left the global default alone.

The decay-NaN fix (root cause 3) is in f824cc1. Really appreciate the deterministic repro — it’s what made the inverse fix testable. Credited your report in the commit.