LoRA trainer has no gradient clipping or non-finite guard: one bad batch loses the whole run

[axiom-of-choice]

Summary

The LoRA trainer has no gradient clipping and no guard against non-finite gradients. A single batch whose gradient overflows puts the LoRA weights into NaN permanently, and every subsequent iteration reports nan with no warning and no way to recover. The only way to notice is to watch the loss.

mx.optimizers.clip_grad_norm(grads, max_norm) already exists in MLX and returns (clipped_grads, total_norm). mlx_lm/tuner/trainer.py never calls it. A grep for clip, isnan, isfinite or norm in that file returns nothing, and there is no --grad-clip / --max-grad-norm flag on mlx_lm.lora.

What I hit

Fine-tuning Qwen/Qwen3-1.7B (LoRA rank 8, 16 layers, bf16), three runs died the same way:

Run	lr	batch	scale	optimizer	schedule	Result
1	1e-5	2	20.0	adam	none	NaN at iter 380, after 370 stable iters (loss 0.95→1.17)
2	5e-6	4	20.0	adam	none	NaN at iter 10, resuming from run 1's iter-200 checkpoint
3	5e-6	4	8.0	adamw	cosine + warmup 100	NaN at iter 10, at lr 4.5e-7

Run 3 is the one I think matters: NaN at a learning rate of 4.5e-7, still inside warmup, where lora_b is barely off its zero init and the model is essentially the base model. This is not divergence from too high an lr, and lowering the lr did not fix it.

Trainable parameters: 0.289% (4.981M/1720.575M)
Iter 1: Val loss 1.472, Val took 49.392s
Iter 10: Train loss nan, Learning Rate 4.500e-07, It/sec 0.071, ...

I bisected the layers separately before filing, so this isn't a bad-data report:

• Data: 8786 rows, token ids within [0, 151668] for vocab 151936, no empty rows, shortest 420 tokens.
• Forward: 20 steps, loss 1.09-1.87, |logit|max 58-66, no NaN or inf.
• Backward (grad-checkpoint on): 25 steps, all gradients finite, |grad|max 0.39-1.48.
• mx.compile + optimizer update: 15 steps with run 3's exact config, no NaN.

Individual batches and the step function are both fine. The NaN only shows up in the real training sequence, on a specific batch.

Correction to an earlier version of this issue: I originally wrote that the batch order is not reproducible, because train() calls iterate_batches without passing seed. That was wrong, and I'd rather fix it than leave it standing. run() in lora.py calls np.random.seed(args.seed) before training, and nothing in the CLI training path draws from numpy in between, so batch order for a normal mlx_lm.lora run is reproducible. There is a real but much smaller bug there, iterate_batches gates on if seed: so an explicit seed=0 is ignored, which I've filed separately as #1660. It is not a factor in this issue.

Two things I measured that shape the fix

1. Clipping alone does not fix this. clip_grad_norm propagates non-finite values rather than removing them:

>>> g = {"a": mx.array([float("nan"), 1.0]), "b": mx.array([2.0, 3.0])}
>>> out, norm = optim.clip_grad_norm(g, 1.0)
>>> float(norm), tree_flatten(out)
(nan, [('a', array([nan, nan])), ('b', array([nan, nan]))])

The finite entries get contaminated too, because everything is scaled by max_norm / norm and norm is NaN. The inf case is worse: the normalizer becomes 0, and inf * 0 is NaN, so infinite gradients also poison every parameter.

So clipping covers large-but-finite gradients, which is a different problem from the one above. Fixing the NaN case needs both halves: clipping, plus a guard that drops the update when the norm is not finite.

2. The guard has to stay in the graph. step() is wrapped in @partial(mx.compile, inputs=state, outputs=state). Branching on float(norm) in Python would force a device sync every iteration. This composition compiles and works:

grad, norm = clip_grad_norm(grad, max_norm)
keep = mx.isfinite(norm)
grad = tree_map(lambda g: mx.where(keep, g, 0), grad)
optimizer.update(model, grad)

I verified clip_grad_norm works inside mx.compile, including across a step with a 1000x gradient spike, so the data-dependent value is not a tracing problem.

One caveat worth documenting rather than fixing: zeroing the gradient is not a complete no-op. A momentum optimizer still moves the weights from its existing state, and AdamW applies weight decay independently of the gradient (w -> w*(1 - lr*wd)), measured at 1.07e-07 relative at lr=5e-6. With Adam, or AdamW at weight_decay=0.0, the weights are exactly unchanged. Either way the bad batch contributes no learning signal, which is the point.

Prior NaN reports

• Val and Training Loss NaN on M4 Max starting on mlx-lm 0.22.2 #49 "Val and Training Loss NaN on M4 Max" — closed, resolved by upgrading mlx to 0.24.1. That was an mlx bug.
• NaN values doing fine tuning with gpt-oss:20b #361 "NaN values doing fine tuning with gpt-oss:20b" — closed, 8 comments, several people reporting the same symptom, went toward data format and MXFP4 quantization.

Neither was caused by missing clipping, and I'm not claiming otherwise. I mention them only as evidence that people do run into NaN during fine-tuning, and that today there's nothing in the trainer that limits the damage.

Proposed change

Deliberately minimal, and off by default so no existing run changes behavior:

1. grad_clip: Optional[float] = None on TrainingArgs.
2. --grad-clip on mlx_lm.lora, plus its CONFIG_DEFAULTS entry.
3. Apply it in the compiled step() between value_and_grad and optimizer.update, on the already-accumulated and averaged gradient rather than per micro-batch, so it composes correctly with grad_accumulation_steps and average_gradients.
4. The non-finite guard via mx.where, as above.

I have this working with tests, and would be happy to open a PR. Two questions before I do:

• Is the non-finite guard something you want in the same change? It's less standard than clipping, and I'd rather ask than assume. Clipping on its own would not have saved any of the three runs above.
• Do you prefer one --grad-clip flag controlling both, or the guard separated behind its own flag?

Environment: mlx-lm 0.31.3, mlx 0.32.0, M2 Pro 32 GB, macOS 26.5.2. Verified against main at e5baded.

[axiom-of-choice]

Follow-up with real-data evidence, from a 3000-iter run of the same Qwen3-1.7B config that previously NaN'd at iter 10, now with both guards active (--grad-clip 1.0):

Iter 10: Train loss 1.3797, LR 4.500e-07, grad_norm 4.232,        skipped 0
Iter 20: Train loss 1.4634, LR 9.500e-07, grad_norm 4.200,        skipped 0
Iter 30: Train loss 1.4008, LR 1.450e-06, grad_norm 5044.074,     skipped 0
Iter 40: SKIPPED non-finite batch (loss 1.2308621406555176, norm nan) -- weights untouched
Iter 40: Train loss 1.2680, LR 1.950e-06, grad_norm 177650662.281, skipped 1
Iter 50: Train loss 1.1788, LR 2.450e-06, grad_norm 2.361,        skipped 1

Three details that I think matter for the design:

1. The loss was finite while the gradient norm was NaN. loss = 1.2309, norm = nan on the same batch. The overflow is in the backward pass, not the forward. This is why my earlier bisection found the forward pass, the backward pass in isolation, and the compiled step all clean. It also means the failure is undetectable by watching the loss: by the time the loss prints nan, the weights are already gone. A loss-based guard would not catch this.

2. Clipping alone would still have lost this run. The norm escalated 4.2 → 4.2 → 5044 → 1.78e8 → nan with clipping active the whole time. Clipping bounds the update that gets applied, not the norm computation itself, so the large-but-finite steps were handled and the next one overflowed anyway. Only the mx.isfinite guard prevented the run from ending here. This is the concrete case behind my first question above.

3. It recovered completely. Iter 50 came back at grad_norm 2.361 with the lowest loss of the run so far. In the unguarded version of this exact config, iter 10 was terminal and every subsequent iteration was nan. One skipped batch out of 50 iterations, no lasting effect.

The four rows in the skipped batch are unremarkable (4677-7189 chars, no degenerate content, longest repeated-character run is 8), which is further reason to think this is not a bad-data problem. They did all pack to 2030-2031 tokens against max_seq_length 2048, so the batch was at the length ceiling. That may be a hint about the mechanism, but it's a single sample and I don't want to over-read it.

Possibly related, and I'm not proposing to change it here: CacheDataset.itemlen does len(self._data[idx]), and for the {"text": ...} format self._data[idx] is a dict, so it returns 1 for every row. The length-based sort in iterate_batches therefore doesn't sort with that format, and a batch can mix a 420-token row with a 2048-token one and pad everything to the max. That's a memory and throughput issue rather than a correctness one, but it does mean long-sequence batches get assembled more or less at random.