fix(mlx): wrong gated-delta grads for batch rows past the first

mx .at[:, t].add scatters read the update tensor with a wrong batch
stride on the Metal stream, silently corrupting d_q/d_k/d_v/d_g/d_beta
for every batch element except index 0. Collect per-step grads in lists
and stack instead; verified against plain autodiff at B in {2,3,4}.

The underlying Metal scatter indexing bug is fixed upstream in
ml-explore/mlx#3483 but is not in any release yet (latest: 0.31.2).
The list+stack form could revert to the scatter once a fixed mlx is
the minimum supported version, though it is also the cheaper op here.

Author: Lyxot

unsloth_zoo/gated_delta_vjp.py

@@ -122,11 +122,16 @@ def _chunked_vjp(primals, cotangents, outputs):
         # BPTT: `d_state` is the cotangent w.r.t. the RETURNED state at the
         # current step (= input to step t+1). Starts at d_state_out, then
         # propagates through the recurrence + mask.
-        d_q = mx.zeros_like(q_c)
-        d_k = mx.zeros_like(k_c)
-        d_v = mx.zeros_like(v_c)
-        d_g = mx.zeros_like(g_c)
-        d_beta = mx.zeros_like(beta_c)
+        # Per-step grads are collected in lists and stacked afterwards: each
+        # t is produced exactly once, and mx `.at[:, t].add` scatters read
+        # the update with a wrong batch stride (wrong grads for every batch
+        # row past the first; verified against plain autodiff on mlx 0.31).
+        # Fixed upstream in ml-explore/mlx#3483, not yet in any release.
+        d_q_steps = []
+        d_k_steps = []
+        d_v_steps = []
+        d_g_steps = []
+        d_beta_steps = []
         d_state = d_state_out
 
         for t in range(chunk_T - 1, -1, -1):
@@ -171,7 +176,7 @@ def _chunked_vjp(primals, cotangents, outputs):
                 + dy_t[..., None].astype(mx.float32) * q_t[..., None, :].astype(mx.float32)
             )
             d_q_t = (dy_t[..., None].astype(mx.float32) * state_new).sum(axis=-2)
-            d_q = d_q.at[:, t].add(d_q_t.astype(d_q.dtype))
+            d_q_steps.append(d_q_t.astype(q_c.dtype))
 
             # state_new = state_decayed + k[..., None, :] * delta[..., None]
             d_kd = d_state_new
@@ -203,14 +208,19 @@ def _chunked_vjp(primals, cotangents, outputs):
                 d_g_t = d_decay
 
             d_k_t = d_k_t_from_update + d_k_t_from_kv
-            d_k = d_k.at[:, t].add(d_k_t.astype(d_k.dtype))
-            d_v = d_v.at[:, t].add(d_v_t.astype(d_v.dtype))
-            d_g = d_g.at[:, t].add(d_g_t.astype(d_g.dtype))
-            d_beta = d_beta.at[:, t].add(d_beta_t.astype(d_beta.dtype))
+            d_k_steps.append(d_k_t.astype(k_c.dtype))
+            d_v_steps.append(d_v_t.astype(v_c.dtype))
+            d_g_steps.append(d_g_t.astype(g_c.dtype))
+            d_beta_steps.append(d_beta_t.astype(beta_c.dtype))
 
             # d_state_prev = recurrence-derived gradient + mask passthrough.
             d_state = d_state_prev_via_recurrence + d_state_prev_passthrough
 
+        d_q = mx.stack(d_q_steps[::-1], axis=1)
+        d_k = mx.stack(d_k_steps[::-1], axis=1)
+        d_v = mx.stack(d_v_steps[::-1], axis=1)
+        d_g = mx.stack(d_g_steps[::-1], axis=1)
+        d_beta = mx.stack(d_beta_steps[::-1], axis=1)
         d_mask = mx.zeros_like(mask_chunk) if mask_chunk is not None else None
         return d_q, d_k, d_v, d_g, d_beta, d_state, d_mask