* tests/studio: accept new grad_norm arg in MLX smoke _on_step callback
The MLX trainer's step callback now passes a ninth positional argument
(grad_norm) per unsloth_zoo/mlx/trainer.py's documented signature
``fn(step, total_steps, loss, lr, tokens_sec, peak_gb, elapsed,
num_tokens, grad_norm=None)``. The smoke's local ``_on_step`` was still
defined with eight, so every per-step invocation raised
``TypeError: _on_step() takes 8 positional arguments but 9 were given``,
``losses_per_step`` never got populated, and the post-train
``assert len(losses_per_step) == 7`` failed.
Add the ninth parameter with a default and surface the gradient norm in
the per-step log line when present.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* 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.
* tests/studio: clarify why MLX smoke pins max_grad_value=0
Refresh the rationale comment to reflect the new default landing in
unslothai/unsloth-zoo#652 (max_grad_value=1.0, not 5.0). The smoke
still needs the explicit pin because neither default value reliably
converges in 7 steps at seed=3407:
max_grad_value=5.0 -- diverges after step 4 (loss 7.3 -> 8.4)
max_grad_value=1.0 -- stalls (loss ~3.2 plateau across seeds)
max_grad_value=0.5/0.25/0.1 -- noisier still
max_grad_norm=1.0 -- cleanly drops loss to <0.01, emits "Unsloth!"
Mention both the historical 5.0 default and the new 1.0 default in
the comment so future readers do not assume the smoke is dead code
referencing a removed knob, and point to the CUDA mirror scripts
(cuda_mlx_mirror_sim.py + cuda_mlx_clip1_vs_norm1.py) for the
empirical evidence.
No behaviour change; comment-only refresh.
* tests/studio: replace fragile substring gate with loss + round-trip gates
The MLX smoke's three "EXPECT in completion" assertions assume the
trained model will greedy-emit the exact "Unsloth" token after the
prompt. On MLX a single near-zero-loss adamw step at the smoke's
fixed seed=3407 can perturb the final-step logits enough that greedy
decoding picks a wrong first token even while the teacher-forced loss
on the training row stays essentially zero (the smoke captures this
exact state -- step 6 loss=0.049, step 7 grad=36.7, step 7 loss=0.17;
completion goes from "Unsloth!" to "5 lbs!"). Reproduced extensively
on CUDA via scripts/cuda_mlx_step7_*.py: at seed=3407 only one config
in a 9-cell sweep lands inside the "Unsloth"-emitting basin, and only
1/3 seeds at that config pass. This is a property of the assertion,
not of save/reload correctness.
Refactor the three assertions to gate on what the smoke is actually
trying to verify:
in_memory:
- hard gate: post_train_loss < 1.0 (training memorised the row).
- soft check: log whether completion contains EXPECT_IN_OUTPUT
into metrics["in_memory_generation_has_expected"]; print a
WARN when missing instead of failing.
lora / merged reload:
- hard gate: reload output must equal the in-memory completion
saved in train_metrics.json. This is the actual save/reload
invariant -- the reloaded weights have to reproduce whatever
the in-memory model produced. Falls back to the original
gibberish gate if train_metrics.json is unavailable.
gguf reload:
- hard gate: llama.cpp produced usable, non-empty output after
the prompt (>=4 chars). llama.cpp's tokenizer + sampling differ
from mlx_lm so byte-exact match isn't sound. Log
gguf_has_expected for visibility.
Result: the smoke still gates on the real failure modes (training
didn't memorise, save/reload corrupted weights, llama.cpp produced
no output), without depending on the brittle "Unsloth as first
greedy-decoded token" guarantee that MLX's step-7 numerics can break
without harming any save/reload semantics.
Cross-version constraint: no transformers / trl API touched.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* tests/studio: gate MLX reload on training-row loss, not greedy text
The strict reload assertion (out == in_mem_out) failed on macOS:
in-memory completion was '5 lbs!' and the reloaded completion was
'_________________________'. Both are corrupted by the same MLX
step-7 grad spike (see scripts/cuda_mlx_step7_*), but greedy decoding
can pick a different first token at near-zero teacher-forced loss
even when weights are byte-identical, so exact text equality is not
the right round-trip invariant.
Replace with teacher-forced loss equality on TRAIN_TEXT: the
reloaded model must reach essentially the same post_train_loss the
in-memory model recorded. That is the real save/reload correctness
gate, robust to MLX's near-zero-loss adamw greedy-decode
perturbation. Falls back to a non-empty-body check when
train_metrics.json is missing.
CUDA mirror at this seed converges cleanly to ~0.006 loss; on MLX
post_train_loss < 1.0 still holds via the existing memorisation
gate. The completion text and "matches in-memory" flag are still
recorded in metrics for visibility, just not gated on.
* tests/studio: align MLX smoke with elementwise-clip + 30-step gates
Two corrections to the earlier f93e918b / e05d6c7d direction:
1. max_grad_value=0.0, max_grad_norm=1.0 picked the memory-heavy
norm clip. On MLX, max_grad_norm requires a cross-tree
reduction and materializing every grad tensor at full
precision; max_grad_value is tree_map(mx.clip) per leaf with
no reduction. MLXTrainingConfig defaults to max_grad_value=1.0
for exactly this reason. Flip the smoke to
max_grad_norm=0.0, max_grad_value=1.0 so the configured clip
matches what actually runs (the trainer prints a "both
enabled, value wins" notice otherwise).
13-seed empirical pass rates at this fixture also favor the
elementwise mode: value=1.0 62%, norm=1.0 46%, value=5.0 33%,
value=0.5 77%. Cheaper default = higher pass rate, no
tradeoff. (See PR #5498 / staging-2#119 rounds A-AT.)
2. max_steps=7 was below the convergence horizon at every clip
tested. At 30 steps every seed hits post_train_loss=0 across
all clip configurations; that's the seed-robust gate. Bump
max_steps 7 -> 30, tighten the memorisation gate from
post_loss < 1.0 to post_loss < 0.1.
3. Relax per-step lower bound from 0 < l to 0 <= l: with
max_steps=30 + bs=2 + grad_accum=3 the LoRA collapses loss
to 0 by ~step 10 and the fp16 per-step loss underflows to
exact 0.0 from then on. That's the success signal, not a bug.
Keeps the e7ec2f52 EXPECT_IN_OUTPUT demotion-to-warning and the
e7347643 reload teacher-forced-loss round-trip invariant -- those
are the right gates regardless of the clip / steps choice.
* tests/studio: hard gate via teacher-forced completion loss
The prior "soft warn + metric" was a step back from the original
hard assert: regressions could land silently if greedy decode
happened to pass on seed=3407 but post_train_loss diverged.
A true hard gate is needed.
Greedy decode is empirically fragile -- a 47-round, 13-seed sweep
on this fixture (see danielhanchen/unsloth-staging-2#119) showed
contains-Unsloth lands in 46-77% across MLX clip configs even
when post_train_loss is zero, because fp16 noise on the first
generated token after PROMPT perturbs the argmax. Teacher-forced
loss on the completion does not have this problem: it just reads
back the probability mass the model assigns to the trained
continuation. In every config where post_train_loss < 0.1, the
completion loss is essentially zero.
Add `_teacher_forced_completion_loss(model, tokenizer, prompt,
completion)` that scores the next-token CE only on the completion
positions (no decoding involved) and assert it < 0.5. This gate
is 100% reliable across (seed, clip, bc) combinations tested,
while the greedy substring check remains as a soft metric so
regressions there are still visible.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
525b3b4a43