From d069dac5724fb53bf2934a7d7d582506b99d484f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 9 Oct 2026 23:44:22 +0000 Subject: [PATCH] Cut the DFlash regression to 500 training steps and a 20-step resume test_dflash_training trained 3 epochs (1,500 steps) and test_dflash_resume then trained a fourth epoch (500 more steps). On the 1-GPU pull-request runner the recorded call durations solve to about 33 s of start-up plus 55-73 ms per step, so the two tests spend roughly 115-140 s and 60-70 s. The loss assertion only needs a final loss below 3.0, and the L40 baseline is already 1.78 at step 500. Train for 500 steps and resume for 20 more (max_steps 520). The resume still goes through the whole restore-and-train path; only its tail is shorter. It had no assertion besides the exit status, so it now reads trainer_state.json and checks that global_step is 520 and that a log entry follows step 500 (the training summary the resumed run appends). The step count is fixed with max_steps instead of num_train_epochs because the regression job runs on 2 GPUs off pull requests, and `accelerate launch` starts one process per visible GPU when there is no accelerate config. An epoch is then 250 steps, so num_train_epochs=1 would train 250 steps with the last loss logged at step 200, which may not clear the unchanged `< 3.0` bound, and 1.04 epochs would resume for 10 steps, not 20. logging_steps and the loss assertions (final < first, final < 3.0) are unchanged. Only the comment on the final-loss bound now cites the step-500 baseline. Estimated saving: about 80 s per pull-request regression job (55 s in training, 25 s in the resume) and about 30 s on the 2-GPU job. Risk: the 1.78 baseline at step 500 comes from a 1,500-step run. With max_steps=500 the linear learning-rate schedule now decays to zero at step 500, so the final loss can differ from that baseline. The bound leaves 1.2 of margin on it. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Read the transformers 4.57.1 sources (max_steps overrides num_train_epochs; with the steps save strategy a checkpoint is saved at step 500 and again at the end of training) and the accelerate 1.10.0 launcher defaults. The recipe's training schema accepts extra keys, and launcher examples already pass training.max_steps. Not run: pytest, because this machine has no GPU and the test downloads Qwen3-0.6B. GPU CI has to confirm that the final loss stays below 3.0 after 500 steps and that the resume run ends at global_step 520. Signed-off-by: Shengliang Xu --- .../regression/torch/speculative/test_dflash.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/regression/torch/speculative/test_dflash.py b/tests/regression/torch/speculative/test_dflash.py index 18e95148fcd..82a80a17153 100644 --- a/tests/regression/torch/speculative/test_dflash.py +++ b/tests/regression/torch/speculative/test_dflash.py @@ -61,6 +61,10 @@ "dflash.dflash_architecture_config.num_hidden_layers=2", ] +# Steps, not epochs: accelerate runs one process per GPU, so an epoch is 500 steps on 1 GPU, 250 on 2. +_TRAIN_STEPS = 500 +_RESUME_STEPS = 20 + @pytest.fixture(scope="session") def qwen3_model_name(): @@ -80,8 +84,8 @@ def test_dflash_training(qwen3_model_name, dflash_output_dir): overrides = [ f"model.model_name_or_path={qwen3_model_name}", f"training.output_dir={output_dir}", - "training.num_train_epochs=3", - "training.save_steps=500", + f"training.max_steps={_TRAIN_STEPS}", + f"training.save_steps={_TRAIN_STEPS}", *_DFLASH_OVERRIDES, ] @@ -108,7 +112,7 @@ def test_dflash_training(qwen3_model_name, dflash_output_dir): first_loss = float(logs[0]["loss"]) final_loss = float(logs[-1]["loss"]) assert final_loss < first_loss, f"Loss did not decrease: {first_loss:.3f} -> {final_loss:.3f}" - # Sanity: final loss should be reasonable (baseline: ~1.1 on L40) + # Sanity: final loss should be reasonable (baseline: ~1.8 at step 500 on L40) assert final_loss < 3.0, f"Final loss {final_loss:.3f} too high (expected < 3.0)" @@ -118,7 +122,7 @@ def test_dflash_resume(qwen3_model_name, dflash_output_dir): overrides = [ f"model.model_name_or_path={qwen3_model_name}", f"training.output_dir={output_dir}", - "training.num_train_epochs=4", + f"training.max_steps={_TRAIN_STEPS + _RESUME_STEPS}", "training.save_steps=5000", *_DFLASH_OVERRIDES, ] @@ -128,6 +132,11 @@ def test_dflash_resume(qwen3_model_name, dflash_output_dir): "speculative_decoding", ) + with open(os.path.join(output_dir, "trainer_state.json")) as f: + state = json.load(f) + assert state["global_step"] == _TRAIN_STEPS + _RESUME_STEPS, state["global_step"] + assert any(h["step"] > _TRAIN_STEPS for h in state["log_history"]), "no log entry after resume" + def test_dflash_export(dflash_output_dir): """Export DFlash checkpoint to deployment format."""