From eaf1748ea5a4fea47b39d4c68dcafd585306fa77 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Tue, 23 Sep 2025 15:50:38 -0700
Subject: [PATCH] enable training for qwen3 next (#496)

---
 mlx_lm/models/gated_delta.py |  4 +++-
 mlx_lm/models/qwen3_next.py  | 16 ++++++++++++----
 mlx_lm/quant/dwq.py          |  4 ++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/mlx_lm/models/gated_delta.py b/mlx_lm/models/gated_delta.py
index c9364e6..4f16fd6 100644
--- a/mlx_lm/models/gated_delta.py
+++ b/mlx_lm/models/gated_delta.py
@@ -92,6 +92,7 @@ def _make_gated_delta_kernel():
 _gated_delta_kernel = _make_gated_delta_kernel()
 
 
+@mx.compile
 def _gated_delta_step_ops(
     q: mx.array,
     k: mx.array,
@@ -203,6 +204,7 @@ def gated_delta_update(
     A_log: mx.array,
     dt_bias: mx.array,
     state: Optional[mx.array] = None,
+    use_kernel: bool = True,
 ) -> Tuple[mx.array, mx.array]:
 
     beta = mx.sigmoid(b)
@@ -213,7 +215,7 @@ def gated_delta_update(
         if state is None:
             state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
 
-    if mx.default_device() != mx.gpu or not mx.metal.is_available():
+    if not use_kernel or mx.default_device() != mx.gpu or not mx.metal.is_available():
         return gated_delta_ops(q, k, v, g, beta, state)
     else:
         return gated_delta_kernel(q, k, v, g, beta, state)
diff --git a/mlx_lm/models/qwen3_next.py b/mlx_lm/models/qwen3_next.py
index ad7b27f..64f8fe6 100644
--- a/mlx_lm/models/qwen3_next.py
+++ b/mlx_lm/models/qwen3_next.py
@@ -251,14 +251,22 @@ class Qwen3NextGatedDeltaNet(nn.Module):
             )
         ]
 
-        if cache is not None:
-            state = cache[1]
-
+        state = cache[1] if cache else None
         inv_scale = k.shape[-1] ** -0.5
         q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
         k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
 
-        out, state = gated_delta_update(q, k, v, a, b, self.A_log, self.dt_bias, state)
+        out, state = gated_delta_update(
+            q,
+            k,
+            v,
+            a,
+            b,
+            self.A_log,
+            self.dt_bias,
+            state,
+            use_kernel=not self.training,
+        )
 
         if cache is not None:
             cache[1] = state
diff --git a/mlx_lm/quant/dwq.py b/mlx_lm/quant/dwq.py
index 8066b72..67f7627 100644
--- a/mlx_lm/quant/dwq.py
+++ b/mlx_lm/quant/dwq.py
@@ -324,6 +324,10 @@ def main():
             bits=args.bits,
         )
 
+    if mx.metal.is_available():
+        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        mx.set_wired_limit(max_rec_size)
+
     opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
     dwq_quantize(
         model,