From 92f9d4bda07f066acc0839b863494595d9876161 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 15 May 2026 07:54:22 +0000 Subject: [PATCH] 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. --- tests/studio/run_real_mlx_smoke.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index f0c90dd9c6..1a11a91c75 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -296,11 +296,16 @@ def cmd_train(args) -> int: args = config, ) - def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens): + def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed, + num_tokens, grad_norm = None): losses_per_step.append(round(float(loss), 4)) + grad_text = ( + f" grad={grad_norm:.4f}" + if grad_norm is not None else "" + ) print( f" step {step}/{total} loss={loss:.4f} lr={lr:.2e} " - f"tok/s={tok_s:.0f} peak={peak_gb:.2f}GB", + f"tok/s={tok_s:.0f} peak={peak_gb:.2f}GB{grad_text}", flush = True, )