[davrens101]

Fixes #8750
Fixes #8951
Fixes #8752

What was wrong

Three small defects in OneCycle (deepspeed/runtime/lr_schedules.py) all sit in the same two initializer methods, so I fixed them together.

1. _initialize_momentum built every param group's betas as (mom, 0.99). Only betas[0] is meant to cycle, so the optimizer's own beta2 (Adam/AdamW default 0.999, or a per-group value) was replaced with 0.99 when the scheduler was constructed and again on every step() ([BUG] OneCycle overwrites the optimizer's beta2 with 0.99 #8750).
2. At construction the initial betas[0] was set to cycle_min_mom. The schedule at step 0 is at cycle_min_lr, where get_mom() returns cycle_max_mom, so the first optimizer update used a momentum the schedule never reports ([BUG] OneCycle applies the minimum momentum to the first Adam update #8951).
3. _initialize_lr applied cycle_min_lr to the optimizer but threw the returned lrs away, so get_last_lr() raised AssertionError: need to call step() first until the first step(). WarmupLR, WarmupDecayLR, WarmupCosineLR and LRRangeTest all return the initial lr ([BUG] OneCycle.get_last_lr() raises before the first step() while other schedules return the initial lr #8752).

How I saw it

On master (6f96b4a), Adam with betas=(0.9, 0.999), OneCycle(opt, cycle_min_lr=1e-4, cycle_max_lr=1e-3, cycle_first_step_size=10):

before: (0.9, 0.999)
after init: (0.8, 0.99) get_mom(): [(0.9, 0.99)]
get_last_lr before step: AssertionError need to call step() first
after step: (0.89, 0.99)

With the change, after init is (0.9, 0.999), matching get_mom() apart from the unchanged beta2, get_last_lr() returns [0.0001], and beta2 stays 0.999 (and 0.95 for a group created with 0.95) through the cycle.

What changed

• _initialize_momentum reads each group's existing betas[1] and builds min_moms/max_moms from it, then starts the optimizer at max_moms (only when last_batch_iteration == -1, as before).
• _initialize_lr stores the result of update_lr in _last_lr.
• Tests in tests/unit/runtime/test_lr_schedulers.py: two new non-distributed tests (per-group beta2 kept and initial momentum matches get_mom(); get_last_lr() before the first step). One existing assertion in test_one_cycle_accepts_per_group_lr_and_momentum_lists expected the initial beta1 to be cycle_min_mom (0.8, 0.85); it now expects cycle_max_mom (0.9, 0.95). Its later peak/bottom assertions are unchanged.

What I ran

CPU only, Windows 11, torch 2.15.0+cpu (as reported by torch.__version__ in my environment), Python 3.12, CPU accelerator. deepspeed/accelerator and two ops entries are git symlinks that Windows checks out as text files, so I replaced them locally with directory junctions (marked skip-worktree, not part of the commit).

• Before the fix (source file reverted, new tests kept): pytest tests/unit/runtime/test_lr_schedulers.py -k one_cycle gave 3 failed, 12 passed (the two new tests plus the updated per-group test).
• After the fix: 15 passed, 75 deselected.
• Whole test_lr_schedulers.py module: 45 passed, 45 failed. All 45 failures are the DistributedTest classes failing with ValueError: cannot find context for 'forkserver', which is a Windows limitation of the test launcher, not this change. I read TestOneCycle::test_lr/test_mom: they use get_lr()/get_mom() and do not depend on the initial optimizer betas, but I could not run them.
• yapf and flake8 on the two touched files: no output.

Unpublished follow-up: restoring _last_lr from an older checkpoint

The local follow-up commit 3df189a6 is not pushed; this PR remains at 2b8b59d34add1d6225d53ac49cc54d61e12c2f1e. It remains held for human review, DCO attestation and the outstanding validation described below. It adds test_one_cycle_get_last_lr_after_checkpoint_restore covering load_state_dict directly: get_last_lr() on a scheduler restored from a checkpoint written before _last_lr was stored, and the same thing with last_batch_iteration reused from the checkpoint.

I ran the four parameter combinations against the real scheduler/optimizer/state-dict code (isolated harness, exact source):

legacy_checkpoint	resume_iteration	pre-fix (master lr_schedules.py)	fixed
False	False	PASS	PASS
False	True	AttributeError (eager self._last_lr default)	PASS
True	False	AssertionError (legacy LR restored)	PASS
True	True	AttributeError	PASS

So 3 of 4 configurations fail without the follow-up; all 4 pass with it. The module-level run needs the Windows symlink workaround described above; this isolated proof uses the same classes and logic.

Honest limits

Not validated: any GPU, ZeRO stage, offload or precision configuration, and the distributed TestOneCycle cases. This only changes scheduler bookkeeping that runs on the host before any ZeRO or precision handling. WarmupLR/WarmupDecayLR/WarmupCosineLR/LRRangeTest do not touch betas, so there was no sibling momentum path to check. Users who relied on the old beta2 of 0.99 under OneCycle will now get the optimizer's configured value; that is the documented behavior but it is a numerical change for them.

Correction: this PR duplicates earlier open work. #8495 and #8751 preserve beta2, #8753 records the initial learning rate, and #8953 fixes the initial momentum endpoint. My earlier duplicate search missed those PRs. I am withdrawing this combined PR rather than competing with them; the local checkpoint-restore follow-up remains unpublished.

I used an AI coding assistant to write and test this change, and the numbers above come from its runs on my machine.

[chatgpt-codex-connector]

Codex Review Summary: ✅ Completed

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

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-10T19:27:59.980198Z	b45fddd	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]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b45fddd920

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Exercise the beta changes in an integration training loop

This change alters the Adam beta values consumed by the first and subsequent model updates, but this loop only invokes scheduler.step() on a bare optimizer and never performs an optimizer/DeepSpeed training step. The existing TestOneCycle.test_mom integration loop only records get_mom()—which already returned the expected maximum before this patch—so it would not catch either fixed optimizer-beta defect. Add and execute a minimal SimpleModel training test that verifies the optimizer betas or resulting parameters through actual updates.

AGENTS.md reference: AGENTS.md:L81-L82

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Refresh _last_lr after restoring OneCycle state

When a stage-0 checkpoint taken after the first scheduler step is restored and get_last_lr() is queried before the next step, this constructor value is stale: DeepSpeedEngine._load_checkpoint restores the optimizer at engine.py:4735-4738 and then OneCycle.load_state_dict() restores only last_batch_iteration, leaving _last_lr at cycle_min_lr. The method therefore silently reports the minimum instead of the optimizer LR saved in the checkpoint; persist/restore _last_lr or refresh it from the restored param groups.

Useful? React with 👍 / 👎.

[davrens101]

Pushed fixes to this branch: persisted _last_lr in OneCycle.state_dict()/load_state_dict() so a restored checkpoint reports the optimizer LR, and added test_mom_betas_applied_through_training_step verifying the beta fix through a real DeepSpeed optimizer step (commit 2b8b59d).

[davrens101]

I missed the earlier open fixes in #8495, #8751, #8753 and #8953. They already cover the three issues combined here, so I am closing this duplicate. The local checkpoint-restore follow-up remains unpublished.