Skip to content

fix(train): skip optimizer and scheduler for eval-only - #2296

Merged
zhuzilin merged 3 commits into
THUDM:mainfrom
EazyReal:fix/eval-only-skip-train-stack
Aug 21, 2026
Merged

zhuzilin merged 3 commits into
THUDM:mainfrom
EazyReal:fix/eval-only-skip-train-stack

Conversation

@EazyReal

@EazyReal EazyReal commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

Replaces #2109.

Desired vs bug

Eval-only is intended. train.py already has:

# special case for eval-only
if args.num_rollout == 0 and args.eval_interval is not None:
    ray.get(rollout_manager.eval.remote(rollout_id=0))

The train loop is range(..., num_rollout), so this correctly does zero training steps. Eval still needs the actor loaded and update_weights() into the rollout engines, so create_training_models runs first.

The bug is that bring-up always constructed Megatron's optimizer and OptimizerParamScheduler. With --num-rollout 0 and default --lr-decay-iters, slime sets train_iters = 0lr_decay_steps = 0, and Megatron asserts lr_decay_steps > 0. You never reach the eval special case.

Why not train_iters = 1

That was #2109. It lies to the scheduler ("there is one training iteration") so the constructor succeeds, then still runs zero train steps. train_iters stops meaning "how many steps we train." It also still dies if lr_warmup_steps >= lr_decay_steps.

Eval-only does not train, so it should not build a train stack.

Change

  • setup_model_and_optimizer: after get_model, if num_rollout == 0, set no_load_optim and return (model, None, None). Weight load is the same path load_other_checkpoint already uses for ref/teacher (optimizer=None, no_load_optim=True).
  • create_training_models: do not allocate a critic when num_rollout == 0.

num_rollout > 0 is unchanged.

Test plan

  • --num-rollout 0 --eval-interval 1 loads the actor, skips optimizer/scheduler construction, hits rollout_manager.eval
  • A normal training run still builds optimizer + scheduler

--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>
EazyReal and others added 2 commits August 21, 2026 01:52
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants