[Bug] Resume advances Megatron scheduler with rollout_id instead of train step
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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