Repository navigation
Fix stale weight decay after OptimizerParamScheduler.load_state_dict - #7468
Nicholas022400701 wants to merge 3 commits into
Conversation
OptimizerParamScheduler.load_state_dict called step() before it restored start_wd, end_wd, wd_incr_steps and wd_incr_style from the checkpoint. The weight_decay written into the optimizer param groups therefore came from the constructor arguments while the learning rate already came from the checkpoint, so the first optimizer step after a resume ran with the wrong weight decay whenever the two disagreed. Call step() at the end of load_state_dict instead. Signed-off-by: 区梓灏 <116372750+Nicholas022400701@users.noreply.github.com>
… param groups Signed-off-by: 区梓灏 <116372750+Nicholas022400701@users.noreply.github.com>
black --check already fails on this file on main because of this one line, so the formatter run for this change could not pass without it. Signed-off-by: 区梓灏 <116372750+Nicholas022400701@users.noreply.github.com>
svcnvidia-nemo-ci
added
Final Review
|
@skyw, could you review this scheduler resume fix? It moves the initial step until after weight-decay settings are restored, so the first resumed optimizer update uses the checkpoint schedule. The added regression checks the optimizer parameter-group values immediately after loading. |
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
I'm not familiar with how this function works and can't provide meaningful review. The syntax LGTM based on the description of the PR. |
svcnvidia-nemo-ci
added
the
waiting-on-customer
|
This PR was labelled waiting-on-customer on Sep 22 after the comment above, but I cannot find a request addressed to me. If something is needed from my side, a short note on what to change or provide is enough and I will do it. |
svcnvidia-nemo-ci
added
waiting-on-maintainers
|
/claude review |
|
@dimapihtar Add |
svcnvidia-nemo-ci
added
waiting-on-maintainers
|
/review model=claude mode=strict |
svcnvidia-nemo-ci
added
waiting-on-maintainers
What does this PR do?
Fixes a resume bug in
OptimizerParamScheduler.load_state_dict. The method calledself.step(increment=num_steps)before it restoredstart_wd,end_wd,wd_incr_stepsandwd_incr_stylefrom the checkpoint.step()writes bothlrandweight_decayinto every param group, so right after a resume the param groups held the learning rate from the checkpoint schedule but the weight decay from the constructor arguments. The scheduler object itself ended up with the right values, only the param groups were stale.This matters because
train_stepinmegatron/training/training.pycallsoptimizer.step()beforeopt_param_scheduler.step(), so the first optimizer step after a resume ran with the wrong weight decay. Every later step is fine because the next scheduler step rewrites the param groups. The bug only shows up when--use-checkpoint-opt_param-scheduleris set and the weight decay arguments differ from the ones stored in the checkpoint. With matching arguments the two schedules agree and nothing is visible.The fix moves the
num_iters/num_stepslookup and thestep()call to the end ofload_state_dict, after every value has been restored. Legacy checkpoints that only carry the old learning rate keys still load the same way. No other behaviour changes.The change touches two files.
megatron/core/optimizer_param_scheduler.pymoves thestep()call to the end ofload_state_dict, adds a two line comment saying why, and adds the blank line after the module docstring that black 26.3.0 requires. That blank line is the only formatting change. black already fails on this file on main without it, so the formatter run for this PR could not pass otherwise.tests/unit_tests/test_optimizer_param_scheduler.pyaddstest_load_state_dict_applies_checkpoint_weight_decay. It loads a checkpoint whose weight decay schedule differs from the constructor and checks that the param group holds the checkpoint value and that lr and weight decay agree with the scheduler. On main the param group holds 0.05 while the scheduler reports 0.2, so the test fails there.Tests
Run on a CPU GitHub Actions runner with torch 2.14.0+cpu and Python 3.12.
tests/unit_tests/test_optimizer_param_scheduler.pypass and black, isort and ruff are clean on both changed files: https://github.com/Nicholas022400701/oss-ci/actions/runs/35307716702assert 0.05 == 0.2: https://github.com/Nicholas022400701/oss-ci/actions/runs/35306817760Issue tracking
Linked issue: none. I searched the issues and pull requests for this method and did not find a report. #4705 is about matching param groups on load and does not touch this ordering.
Contribution process
Pre-checks
jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it intests/unit_tests/determinism/kernels/manifest.py(guide)Notes on the boxes. No functional test because the change is covered by the unit test. No kernel is touched. The method keeps its existing type hints and no documentation describes this ordering.
tools/autoformat.shneeds a remote and pylint, so I ran the same black, isort and ruff commands on the two changed files in the CI run linked above instead.AI disclosure: I used an AI coding agent to help write this patch, the test and this description. I have read every line of the change myself and I will answer review comments personally.
The three commits are signed off. I can squash them into one if you prefer.