From fceebf181414bf9655efcb202221b0941a9a5319 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 22 Sep 2026 19:38:33 -0700 Subject: [PATCH] fix(train): hand off to the engine before an off-cadence final eval #3182 skips the last rollout's engine handoff (weight onload, update_weights, KV onload) unless a final eval follows, but its guard called should_run_periodic_action without args.num_rollout. That helper only treats the last rollout as due when num_rollout is passed, so whenever num_rollout is not a multiple of eval_interval the guard skipped the handoff while the eval gate below, which does pass num_rollout, still ran the final eval. With --colocate the eval then reached an engine left in the released state: the 2026-09-22 nightly lost test_qwen3_vl_4B_fsdp and test_qwen3_4B_fsdp_true_on_policy (--num-rollout 3 --eval-interval 20) to "Pointer argument cannot be accessed from Triton (cpu tensor?)" in write_req_to_token_pool_triton. Without colocation the final eval would score the previous rollout's weights. Compute eval_due once and use it for both the handoff guard and the eval gate, so the handoff runs exactly when a final eval needs it. Co-Authored-By: Claude Opus 5.5 --- tests/fast/test_train.py | 23 +++++++++++++++++++++++ train.py | 8 ++++---- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/tests/fast/test_train.py b/tests/fast/test_train.py index 32208b19b97..6d8e8c7a8bd 100644 --- a/tests/fast/test_train.py +++ b/tests/fast/test_train.py @@ -123,6 +123,29 @@ async def test_the_last_rollout_is_always_evaluated_even_off_cadence(self, monke assert [event for event in events if event.startswith("eval:")] == ["eval:0", "eval:1", "eval:2"] + async def test_an_off_cadence_final_eval_follows_the_last_handoff(self, monkeypatch: pytest.MonkeyPatch): + """The last rollout skips its handoff only when no final eval follows; an off-cadence + final eval must still see the engine onloaded with the trained weights.""" + events: list[str] = [] + args = _make_args(num_rollout=3, eval_interval=2, offload_rollout=True) + _install_driver_fakes(monkeypatch, args, events) + + await train_driver.train(args) + + final = events[events.index("actor_train:2") : events.index("eval:2")] + assert final.index("onload_weights") < final.index("update_weights:2") < final.index("onload_kv") + + async def test_the_last_rollout_skips_the_handoff_without_an_eval(self, monkeypatch: pytest.MonkeyPatch): + events: list[str] = [] + args = _make_args(num_rollout=3, offload_rollout=True) + _install_driver_fakes(monkeypatch, args, events) + + await train_driver.train(args) + + assert "update_weights:1" in events + assert "update_weights:2" not in events + assert "onload_weights" not in events[events.index("actor_train:2") :] + class TestWeightEqualityCheck: async def test_weight_equality_check_is_routed_to_the_inference_controller(self, monkeypatch: pytest.MonkeyPatch): diff --git a/train.py b/train.py index c665827a7ad..e4316c02187 100644 --- a/train.py +++ b/train.py @@ -146,9 +146,9 @@ async def save_training_model(model): if external_save: os.remove(args.save_trigger_sentinel) - if rollout_id + 1 < args.num_rollout or should_run_periodic_action( - rollout_id, args.eval_interval, num_rollout_per_epoch - ): + # One predicate for both blocks: the handoff below exists to feed this eval on the last rollout. + eval_due = should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch, args.num_rollout) + if rollout_id + 1 < args.num_rollout or eval_due: if args.colocate_memory_peak_device == "gpu": await actor_model.clear_memory() if lora_rollout_enabled(args): @@ -163,7 +163,7 @@ async def save_training_model(model): if args.offload_rollout: await inference_controller.onload_kv() - if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch, args.num_rollout): + if eval_due: await inference_controller.prepare_eval() await eval_dispatcher.dispatch(rollout_id, force=rollout_id == args.num_rollout - 1)