Repository navigation
Fix OneCycle initial momentum: use cycle_max_mom, not cycle_min_mom - #8979
siddhant-rajhans wants to merge 1 commit into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ba1d132235
ℹ️ 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".
| reported = scheduler.get_mom()[0][0] | ||
| actual = optimizer.param_groups[0]['betas'][0] | ||
| assert actual == pytest.approx(reported), ( | ||
| f"optimizer beta1 ({actual}) != scheduler.get_mom() ({reported})") | ||
| assert actual == pytest.approx(0.9) |
There was a problem hiding this comment.
Add a training-loop regression for initial momentum
When cycle_min_mom != cycle_max_mom, this fix changes Adam’s first-moment state on the first update and subsequent parameter updates, but this test only checks constructor metadata and never performs an optimizer update. The existing TestOneCycle.test_mom asserts scheduler-reported momenta rather than optimizer state or parameter values, so its assertions also pass with the faulty initialization. Root AGENTS.md requires integration coverage for training-loop contract changes; extend that training test with deterministic updates and compare optimizer state or parameters against an independent Adam reference initialized with cycle_max_mom.
AGENTS.md reference: AGENTS.md:L81-L81
Useful? React with 👍 / 👎.
_initialize_momentum() set the optimizer's initial beta1 to cycle_min_mom, but the inverse-cycling schedule starts at cycle_max_mom (scale_factor=0 yields max_mom - 0 = max_mom). This caused the first optimizer.step() to use the wrong momentum, diverging from the intended schedule. The fix mirrors the LR/momentum symmetry: LR initialises at min and ramps up; momentum initialises at max and ramps down. Also corrects the cycle_min_mom docstring which described the lower bound as "initial momentum", contradicting the inverse-cycling behaviour, and updates the per-group list test that asserted the wrong initial betas. Fixes deepspeedai#8951 Signed-off-by: Siddhant Rajhans <51508945+siddhant-rajhans@users.noreply.github.com>
ba1d132 to
5bcf1d2
Compare
Bug
OneCycle._initialize_momentum()sets the optimizer's initialbeta1tocycle_min_mom, but the inverse-cycling schedule starts atcycle_max_mom(atscale_factor=0,get_mom()returnsmax_mom - 0 = max_mom). This means the firstoptimizer.step()consumes a different momentum than the schedule intends, causing a training-correctness divergence from the second update onward.Fixes #8951
Root cause
The LR/momentum symmetry: LR initialises at
minand ramps up; momentum initialises atmaxand ramps down._initialize_lrcorrectly setsmin_lrs;_initialize_momentumwas incorrectly mirroring that instead of inverting.Changes
deepspeed/runtime/lr_schedules.py: Initialise optimizer betas fromself.max_momsinstead ofself.min_momscycle_min_momparam (it's the lower bound, not the initial value)tests/unit/runtime/test_lr_schedulers.py: Fixtest_one_cycle_accepts_per_group_lr_and_momentum_listswhich asserted the wrong initial betas, and addtest_one_cycle_initial_optimizer_momentum_matches_scheduleverifying optimizer state matches the scheduler's reported momentum at construction