THUDM / THUDM/slime

[BUG] LR scheduler double-counted when resuming from checkpoint

Open Beginner friendly
#1,546 0 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

When resuming training from a checkpoint, the LR scheduler's num_steps is incremented twice, causing the scheduler position to be doubled.

Root Cause

In slime/backends/megatron_utils/model.py, function initialize_model_and_optimizer() (line 786):

iteration, _ = load_checkpoint(model, optimizer, opt_param_scheduler, ...)

opt_param_scheduler.step(increment=iteration * args.global_batch_size)  # ← BUG

Megatron's load_checkpoint() already calls opt_param_scheduler.load_state_dict(), which internally calls self.step(increment=num_steps) with the checkpoint's saved num_steps. Line 786 then adds iteration * global_batch_size again.

Result: scheduler.num_steps is doubled on every resume.

Fix

Remove line 786:

# opt_param_scheduler.step(increment=iteration * args.global_batch_size)

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 in slime/backends/megatron_utils/model.py at initialize_model_and_optimizer() around line 786. Read how load_checkpoint() restores opt_param_scheduler state and compare it with the following scheduler step. Done means resuming from a checkpoint advances num_steps once rather than twice.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
1/5
Estimated time
Under an hour
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.