From c91fa2615a35b533a02b3cb9b370aa2479d4ce28 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 17:56:19 +0000 Subject: [PATCH] tests/studio: assert losses_per_step matches max_steps, not stale 7 PR #5537 bumped max_steps from 7 to 30 but the post-train assertion still hardcoded the old count, so every fresh run that reaches the post-train phase fails on `expected 7 logged steps, got [30 floats]`. Derive the expected count from `config.max_steps` and add a `train_result["train_steps"]` cross-check so the gate self-updates with future sweep changes. --- tests/studio/run_real_mlx_smoke.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 27f682ee4e..72e7b74c58 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -390,7 +390,15 @@ def cmd_train(args) -> int: ) if k in train_result } - assert len(losses_per_step) == 7, f"expected 7 logged steps, got {losses_per_step}" + expected_logged_steps = int(config.max_steps) + assert ( + len(losses_per_step) == expected_logged_steps + ), f"expected {expected_logged_steps} logged steps, got {losses_per_step}" + if "train_steps" in train_result: + assert int(train_result["train_steps"]) == expected_logged_steps, ( + f"expected train_steps={expected_logged_steps}, got " + f"{train_result['train_steps']}" + ) for i, l in enumerate(losses_per_step): # Allow exact 0.0: fp16 per-step loss underflows to 0.0 after # the LoRA reaches loss=0 around step ~10 with this fixture +