fix(train): don't rebuild resumed schedulers with last_epoch - #1049
fix(train): don't rebuild resumed schedulers with last_epoch#1049zihanlin-ai wants to merge 2 commits into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
Included review availability: Your plan provides up to 4 included reviews per hour; 3 remain after this review. 📝 WalkthroughWalkthroughChangesResume training
Merge Risk: ⚪ Minimal · up to The change fixes distributed checkpoint resume failures while preserving saved scheduler state and improving test coverage; no actionable merge-blocking risk remains beyond normal checks and review. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Full details: Linked Issues checkExplanation The changes satisfy issue ✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
On resume, setup_optimizer constructed fresh schedulers with last_epoch=checkpointer.previous_epoch. That argument is redundant: three lines later load_scheduler_state_dict() restores the scheduler's real state (last_epoch, base_lrs, _step_count). It is also wrong twice over: these are per-step schedulers, so an epoch index is the wrong unit, and constructing with last_epoch >= 0 makes LRScheduler.__init__ demand initial_lr in every param group. The single-device optimizer load happens to restore initial_lr inside its saved param groups, but the distributed path (set_optimizer_state_dict) does not round-trip it, so every resume under DDP/FSDP crashes with KeyError: "param 'initial_lr' is not specified in param_groups[0]". CI never sees this because its single-GPU runners skip the multi-GPU resume tests. Always construct with the default last_epoch=-1 and let load_scheduler_state_dict() restore the true position. Unmasking the crash exposed a second gap: the checkpoint_dir fixture saved a never-stepped AdamW, so test_distributed_resume[ddp]'s "optimizer state restored" assertion was checking an empty dict against an empty dict ([fsdp] only passed it because DCP initializes state as a side effect of set_optimizer_state_dict). The fixture now takes one real optimizer step before saving, then re-pins the weights, so the assertion has teeth on both paths. Fixes vllm-project#1048 Signed-off-by: Zihan Lin <linzihan.ai@gmail.com>
4381ad7 to
dc212e4
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dc212e400b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
… scheduler bases on them Loading the optimizer state overwrites each param group's lr with the decayed value from the checkpoint, and scheduler construction stamps the group lr as its base (initial_lr). With a scheduler state file present that base is corrected by load_scheduler_state_dict; without one (legacy checkpoints, or an interruption between the optimizer and scheduler writes) the decayed LR silently became the base forever - a linear schedule saved near LR zero would keep training at zero. Restore the configured LRs after the optimizer load, so the missing-state path restarts the schedule from the configured base. Raised by review on the PR. Signed-off-by: Zihan Lin <linzihan.ai@gmail.com>
1 participant
Purpose
Fixes #1048: every distributed resume (DDP and FSDP) crashes in
setup_optimizerwithKeyError: "param 'initial_lr' is not specified in param_groups[0]".The
last_epoch=checkpointer.previous_epochpassed at scheduler construction is redundant —load_scheduler_state_dict()restores the scheduler's real state three lines later — and wrong twice over: these are per-step schedulers, so an epoch index is the wrong unit, andlast_epoch >= 0makesLRScheduler.__init__demandinitial_lrin every param group, which the distributed optimizer load (set_optimizer_state_dict) does not round-trip. The fix is to drop it: construct with the defaultlast_epoch=-1and letload_scheduler_state_dict()restore the true position. Net −4 lines intrainer.py.Unmasking the crash exposed a second gap, also fixed here: the
checkpoint_dirfixture saved a never-stepped AdamW, sotest_distributed_resume[ddp]'s "optimizer state restored" assertion compared an empty dict to an empty dict ([fsdp]only passed it because DCP initializes state as a side effect). The fixture now takes one real optimizer step before saving.One behavioral note: a legacy checkpoint with no scheduler file previously got a scheduler positioned at
last_epoch = epoch_index— an epoch count where a step count is expected, so not a meaningful position. Such checkpoints now restart the schedule from step 0.Tests
test_distributed_resume[ddp/fsdp]go red → green on a multi-GPU machine. Fulltests/unit: