Skip to content

Fix OneCycle initial momentum: use cycle_max_mom, not cycle_min_mom - #8979

Open
siddhant-rajhans wants to merge 1 commit into
deepspeedai:masterfrom
siddhant-rajhans:fix/onecycle-initial-momentum
Open

siddhant-rajhans wants to merge 1 commit into
deepspeedai:masterfrom
siddhant-rajhans:fix/onecycle-initial-momentum

Conversation

@siddhant-rajhans

Copy link
Copy Markdown

Bug

OneCycle._initialize_momentum() sets the optimizer's initial beta1 to cycle_min_mom, but the inverse-cycling schedule starts at cycle_max_mom (at scale_factor=0, get_mom() returns max_mom - 0 = max_mom). This means the first optimizer.step() consumes a different momentum than the schedule intends, causing a training-correctness divergence from the second update onward.

Fixes #8951

Root cause

# Before (line 576):
for momentum, group in zip(self.min_moms, optimizer.param_groups):
    group['betas'] = momentum

# After:
for momentum, group in zip(self.max_moms, optimizer.param_groups):
    group['betas'] = momentum

The LR/momentum symmetry: LR initialises at min and ramps up; momentum initialises at max and ramps down. _initialize_lr correctly sets min_lrs; _initialize_momentum was incorrectly mirroring that instead of inverting.

Changes

  1. deepspeed/runtime/lr_schedules.py: Initialise optimizer betas from self.max_moms instead of self.min_moms
  2. Docstring: Remove "Initial momentum" wording from cycle_min_mom param (it's the lower bound, not the initial value)
  3. tests/unit/runtime/test_lr_schedulers.py: Fix test_one_cycle_accepts_per_group_lr_and_momentum_lists which asserted the wrong initial betas, and add test_one_cycle_initial_optimizer_momentum_matches_schedule verifying optimizer state matches the scheduler's reported momentum at construction

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Nov 2, 2026 •

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-11-02T00:23:06.467785Z ba1d132 PR opened
ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment on lines +872 to +876
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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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>
@siddhant-rajhans
siddhant-rajhans force-pushed the fix/onecycle-initial-momentum branch from ba1d132 to 5bcf1d2 Compare November 2, 2026 05:10

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] OneCycle applies the minimum momentum to the first Adam update

1 participant