THUDM / THUDM/slime

[Bug] Resume advances Megatron scheduler with rollout_id instead of train step

Open
#1,761 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
8.5k
Forks
1.3k
Avg merge
5h 36m
Merged PRs (30d)
22

Description

Bug Description

On resume, slime does:

iteration, _ = load_checkpoint(...)
opt_param_scheduler.step(increment=iteration * args.global_batch_size)

But in slime, the saved Megatron iteration is effectively used as rollout_id, not true train step:

  • save path uses save(rollout_id, ...)
  • load path uses loaded_rollout_id = iteration
  • then start_rollout_id = loaded_rollout_id + 1

This is fine for rollout/data-order resume, but not for scheduler resume.

Steps to Reproduce

opt_param_scheduler is stepped per successful optimizer update:

opt_param_scheduler.step(increment=args.global_batch_size)

A rollout may contain multiple train steps (num_steps_per_rollout > 1), so rollout_id * global_batch_size is not the
correct scheduler position.

Also, Megatron checkpoints already save/load scheduler state, so this extra step(...) can double-advance the scheduler.

Expected Behavior

Resume can produce right LR/WD schedule.

Actual Behavior

Resume can produce wrong LR/WD schedule, especially when:

  • num_steps_per_rollout > 1
  • non-constant LR schedule is used

Current ckpt tests do not catch this because they use constant LR and num_steps_per_rollout == 1.

Environment
  • slime version:
  • Python version:
  • PyTorch version:
  • CUDA/ROCm version:
  • GPU type and count:
  • OS:
  • SGLang version (if relevant):
  • Megatron-LM version (if relevant):
Logs

Additional Context

No response

Pre-submission Checklist
  • I have read the CONTRIBUTING.md and understand the collaboration scope.
  • I have read the documentation and my issue is not addressed there.
  • I have searched for existing issues and this is not a duplicate.
  • I have provided a minimal, reproducible example.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at the resume path around load_checkpoint and opt_param_scheduler.step, then inspect the existing checkpoint tests. Add coverage for num_steps_per_rollout greater than 1 and a non-constant LR schedule; done means resume preserves the correct LR/WD schedule without advancing it twice.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.