tests/studio: pin max_grad_value=0 in MLX smoke so max_grad_norm=1.0 wins

unsloth_zoo PR #5340 added per-element gradient clipping to MLXTrainer
and defaulted ``MLXTrainingConfig.max_grad_value = 5.0``. When both
``max_grad_norm`` and ``max_grad_value`` are set, the trainer warns:

  Unsloth: max_grad_norm and max_grad_value are both enabled;
  ignoring max_grad_norm in favor of max_grad_value.

and silently drops the test's ``max_grad_norm=1.0``. +-5.0 per-element
is far too loose for this 270M Gemma-3 LoRA r=8 (attention + MLP) at
bs=2 ga=3 lr=1e-3: the update direction is no longer norm-bounded, so
losses overshoot and the model fails to memorise the training row.

Reproduced on a CUDA mirror (scripts/cuda_mlx_mirror_sim.py):

  norm_1       (max_grad_norm=1.0, no clip): losses 7.64 -> 0.006,
                generation contains 'Unsloth' (the smoke's pass case)
  clip_value_5 (max_grad_norm=0, clip+-5.0): losses 7.29 -> 8.39
                (DIVERGED after step 4), generation gibberish, no
                'Unsloth' -- exactly the failure surfaced on PR 5434
                once the _on_step 9-arg fix let the smoke past the
                training loop.

Pin ``max_grad_value=0.0`` so the smoke uses the same ``max_grad_norm=
1.0`` clipping it was designed against. Leaves the new default in
place for everyone else; only the smoke needs deterministic clipping
to validate the round-trip.
This commit is contained in:
Daniel Han 2026-05-15 09:47:03 +00:00 committed by danielhanchen
commit f93e918bee

View file

@ -278,6 +278,16 @@ def cmd_train(args) -> int:
optim = "adamw",
weight_decay = 0.0,
max_grad_norm = 1.0,
# Explicitly disable the new per-element clip introduced in
# #5340 (default max_grad_value=5.0). When both are set the
# MLX trainer silently drops max_grad_norm in favour of the
# per-element clip, but +-5.0 is far too loose for this
# 270M LoRA setup -- losses diverge after step 4 and the
# model never memorises "Unsloth!" (verified via the CUDA
# mirror at scripts/cuda_mlx_mirror_sim.py). Pinning
# max_grad_value=0 makes the smoke depend on the same
# max_grad_norm=1.0 the test was originally written for.
max_grad_value = 0.0,
logging_steps = 1,
max_seq_length = 64,
seed = SEED,