Add optional gradient clipping and non-finite guard to the LoRA trainer

[axiom-of-choice]

Closes #1658 (opened for discussion; opening the PR since the implementation is done and tested, happy to change either design decision).

Summary

The LoRA trainer has no bound on gradient size and no guard against non-finite gradients. One batch whose gradient overflows puts the LoRA weights into NaN permanently, and every later iteration reports nan with no warning and no recovery path.

This adds grad_clip to TrainingArgs and --grad-clip to mlx_lm.lora. Off by default, so no existing run changes behavior.

High-level changes

• mlx_lm/tuner/trainer.py: grad_clip: Optional[float] = None on TrainingArgs; clipping plus a non-finite guard inside the compiled step(); max_grad_norm and non_finite_grad_steps surfaced to training callbacks; a warning when a step is dropped.
• mlx_lm/lora.py: --grad-clip flag and its CONFIG_DEFAULTS entry.
• mlx_lm/LORA.md: a "NaN Loss" section.
• mlx_lm/examples/lora_config.yaml: grad_clip: null.
• tests/test_tuner_grad_clip.py: 8 tests driving the real train().

Clipping is applied to the already-accumulated and averaged gradient, immediately before optimizer.update, so it composes with grad_accumulation_steps and average_gradients rather than clipping each micro-batch separately.

Why clipping alone is not enough

clip_grad_norm propagates non-finite values rather than removing them. A NaN in one entry makes the norm NaN, which scales every parameter to NaN. An inf norm makes the normalizer 0, and inf * 0 is NaN, so infinite gradients also poison every parameter:

grads	total_norm	clipped result
[nan, 1.0], [2.0, 3.0]	nan	all four entries nan
[inf, 1.0], [2.0, 3.0]	inf	[nan, 0.0], [0.0, 0.0]

So clipping covers large-but-finite gradients, a different failure from the one that motivated the issue. The guard zeroes the gradient when the norm is not finite:

grad, grad_norm = clip_grad_norm(grad, args.grad_clip)
keep = mx.isfinite(grad_norm)
grad = tree_map(lambda g: mx.where(keep, g, 0), grad)

mx.where keeps the check inside the compiled graph. A Python branch on grad_norm would force a device sync every iteration.

This is not a full no-op: a momentum optimizer still moves the weights from its existing state, and AdamW applies weight decay independently of the gradient (measured 1.07e-07 relative at lr=5e-6; exactly 0 with Adam or weight_decay=0.0). The non-finite values never reach the weights, which is the guarantee. Documented in LORA.md and in a code comment.

Evaluation

The default is inert. 30 iterations, same seed, main at e5baded vs this branch with no --grad-clip:

main   losses: 4.169967 4.170533 4.171922 4.167962 4.160542 4.160848
branch losses: 4.169967 4.170533 4.171922 4.167962 4.160542 4.160848

Losses bit-identical. Weights differ by 2.98e-08, but main differs from itself across two runs by exactly the same 2.98e-08, so that is run-to-run nondeterminism, not this change.

It catches a real overflow and the run recovers. Qwen3-1.7B, LoRA rank 8 over 16 layers, batch 4, max_seq 2048, AdamW, cosine + warmup 100. This config previously died at iter 10 with Train loss nan and never recovered. With --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
Iter 60: Train loss 1.0549, LR 2.950e-06, grad_norm 1.511,         skipped 1

Two things worth noting. The loss was finite (1.2309) while the norm was NaN, so the overflow is in the backward pass and is invisible to loss-watching; by the time the loss prints nan the weights are already lost. And the norm escalated 4.2 → 5044 → 1.78e8 → nan with clipping active the whole time, so clipping alone would not have saved this run. One skip in 60 iterations, and the lowest loss of the run came after it.

Per-iteration cost is below measurement noise. 5 interleaved pairs of 40 iterations on a synthetic model: median 32.09 ms without, 33.78 ms with, stdev 15.33 and 10.19. I am not claiming a number from that; the difference is inside the noise floor. The added work is one square().sum() tree-reduce plus a where per parameter, all inside the existing compiled graph.

Environment: mlx 0.32.0, M2 Pro 32 GB, macOS 26.5.2.

Risks

• The main technical risk was whether clip_grad_norm survives mx.compile, since it introduces a data-dependent value. Verified it does, including across a step with a 1000x gradient spike, and the real run above is 60+ compiled iterations with no tracing error.
• Behavior change when the flag is unset: none, per the equivalence check above.
• The zeroed step is not a complete no-op with a momentum optimizer, as described above. This is the one behavior a user could find surprising, hence documenting it rather than hiding it.
• --grad-clip <= 0 raises ValueError rather than silently doing nothing.

Tests

python -m unittest tests.test_tuner_grad_clip tests.test_tuner_trainer tests.test_finetune
Ran 24 tests in 0.285s -- OK

The new tests drive the real train() rather than a reimplementation, and include a control asserting that without the flag the weights do go non-finite, so the guard is tested against the actual current behavior. Note that a bad batch has to be marked in the batch data rather than by a Python-side counter, since mx.compile traces the loss once and would otherwise bake the overflow into every later iteration.

Open questions

1. Do you want the non-finite guard in this change, or clipping only? Clipping alone would not have saved the run above, which is why both are here, but the guard is the less standard half and I would rather ask.
2. One flag for both, or a separate flag for the guard?

Happy to split, rename, or drop either half.

Pre-merge

Nothing required beyond review. No new dependencies, no migration, no config changes for existing users.

Post-merge

If you want, --grad-clip 1.0 could be mentioned in the LoRA docs as a first thing to try when a fine-tune reports nan, since #49 and #361 both drew several reports of that symptom.

Not in this PR

Two adjacent things I found while diagnosing this and deliberately left out:

• iterate_batches gates its seeding on if seed:, so an explicit seed=0 is silently ignored, and 0 is the default in CONFIG_DEFAULTS. Filed as iterate_batches ignores seed=0 due to if seed:, and 0 is the default seed #1660. Note batch order for a normal CLI run is still reproducible, since run() seeds global numpy before training; this only bites when calling the API directly with seed=0, or when something else draws from numpy in between.
• CacheDataset.itemlen does len(self._data[idx]), and all three dataset classes return the raw record from __getitem__, so for {"text": ...} it returns 1 for every row and the length sort in iterate_batches never sorts. That is a throughput and memory issue, not a correctness one. I measured 31.6% extra padded compute on one 8786-row dataset and will open it separately.

[axiom-of-choice]

Additional diagnostics, because I wanted to rule out the obvious competing explanation before leaning on "a safety net is the right fix".

The batch that overflowed was 4x2048 with all four sequences at ~2030 tokens, i.e. about as long as this dataset gets. So "long sequences overflow" was the natural alternative. It does not hold up. Re-running forward and backward on those exact rows with weights at init and no optimizer updates, batch size 1:

=== the exact rows that exploded, one at a time ===
row 647  (len 2030)  loss=1.2015  norm=4.0593  |g|max=0.9331  ok
row 1493 (len 2030)  loss=2.1618  norm=6.4187  |g|max=1.6782  ok
row 1554 (len 2031)  loss=1.4870  norm=5.3354  |g|max=1.5249  ok
row 6634 (len 2031)  loss=0.8357  norm=2.8471  |g|max=0.7074  ok

=== the 8 longest rows in the dataset (2049 tokens, longer than the culprits) ===
norms 2.7646 - 6.0173, all finite

=== control: the 4 shortest rows (420-480 tokens) ===
norms 6.4126 - 9.8167, all finite

Two conclusions. The culprit rows are perfectly healthy in isolation: the row that produced a NaN mid-run gives norm 4.06 at init. And length does not correlate in the direction the hypothesis needs, in fact it inverts, the shortest rows produce the highest gradient norms here.

Not a context-window effect either: Qwen3-1.7B has max_position_embeddings: 40960 with sliding_window: null and rope_scaling: null, so a 2048-token sequence uses 5% of the supported window.

So the overflow requires the accumulated weight state of a real training run. It's an interaction between the adapted LoRA weights and a batch, not a property of any batch or of any sequence length. No single row reproduces it in isolation.

That is the case for a guard rather than for finding and removing the offending data: there is nothing wrong with the data to remove. The run that skipped the batch kept training and the loss kept falling.

I mention this mainly because #361 was closed after moving toward data format, and #49 after an mlx upgrade. This failure mode is neither, and without something in the trainer there is nothing a user can do about it beyond lowering the learning rate and hoping, which as noted in #1658 did not work here either (NaN at lr 4.5e-7).

[axiom-of-choice]

Note: this and #1665 both touch TrainingArgs and the train() loop in mlx_lm/tuner/trainer.py. They are independent and branch off main separately, so whichever lands second will need a rebase. Happy to do that as soon as either merges, in whichever order suits you.