fix(train): skip optimizer and scheduler for eval-only - #2296
Merged
Merged
Conversation
--num-rollout 0 still loads actor weights for eval. It does not train, so do not construct Megatron's optimizer or LR scheduler (which asserts lr_decay_steps > 0). Dummy train_iters=1 would lie about the schedule. Co-authored-by: Cursor <cursoragent@cursor.com>
2 tasks
Eval-only still builds the actor and loads weights, but there is no optimizer to restore into. Cover the skip with CPU tests for the train.py control flow, critic allocation, and setup_model_and_optimizer. Co-authored-by: Cursor <cursoragent@cursor.com>
That file was a local proof, not a test slime should keep. The skip is the three-line early return in setup_model_and_optimizer. Co-authored-by: Cursor <cursoragent@cursor.com>
This was referenced Aug 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replaces #2109.
Desired vs bug
Eval-only is intended.
train.pyalready has:The train loop is
range(..., num_rollout), so this correctly does zero training steps. Eval still needs the actor loaded andupdate_weights()into the rollout engines, socreate_training_modelsruns first.The bug is that bring-up always constructed Megatron's optimizer and
OptimizerParamScheduler. With--num-rollout 0and default--lr-decay-iters, slime setstrain_iters = 0→lr_decay_steps = 0, and Megatron assertslr_decay_steps > 0. You never reach the eval special case.Why not
train_iters = 1That was #2109. It lies to the scheduler ("there is one training iteration") so the constructor succeeds, then still runs zero train steps.
train_itersstops meaning "how many steps we train." It also still dies iflr_warmup_steps >= lr_decay_steps.Eval-only does not train, so it should not build a train stack.
Change
setup_model_and_optimizer: afterget_model, ifnum_rollout == 0, setno_load_optimand return(model, None, None). Weight load is the same pathload_other_checkpointalready uses for ref/teacher (optimizer=None,no_load_optim=True).create_training_models: do not allocate a critic whennum_rollout == 0.num_rollout > 0is unchanged.Test plan
--num-rollout 0 --eval-interval 1loads the actor, skips optimizer/scheduler construction, hitsrollout_manager.eval