[ai-yang]

Fixes #8951.

A fresh OneCycle starts at minimum LR but previously applied minimum momentum to the optimizer, even though its computed schedule starts at maximum momentum. The first Adam moment update was wrong and subsequent parameter updates diverged from the inverse-cycle schedule.

Change

• Initialize beta1 from max_moms for last_batch_iteration == -1.
• Preserve resumed optimizer settings and the cycle_momentum=False path.
• Correct the contradictory minimum-momentum docstring and clarify the first-update behavior in the existing tutorial.
• Exercise applied optimizer momentum rather than only computed scheduler values. Add real-engine Adam/reference comparisons for scalar and per-group bounds, cycling enabled/disabled, and the resume-constructor guard. Remove the old assertion that pinned the incorrect initial minimum.

No public API, configuration key, checkpoint format, beta2, post-cycle decay, or LR-step indexing change. First-update numerics deliberately change to match the documented inverse cycle.

Validation

Base: 8d4f29c75f4a4b30d4db00a94b4aebfdff5107f3.
Python 3.11.15; torch 2.12.1+cu126; NVIDIA RTX 3090 (one GPU used), CUDA runtime 12.6, driver 580.178.04; Linux x86_64. CPU uses Gloo; GPU uses NCCL. PyTorch Adam, FP32, stage 0, no offload.

export PYTHONPATH="$PWD:$PWD/tests"
export OMP_NUM_THREADS=1 LOCAL_SIZE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
DS_ACCELERATOR=cpu python -m pytest -p pytest_forked --forked \
  tests/unit/runtime/test_lr_schedulers.py \
  -k 'one_cycle or OneCycle' \
  --torch_ver=2.12.1+cu126 --cuda_ver=12.6
DS_ACCELERATOR=cuda CUDA_VISIBLE_DEVICES=0 python -m pytest -p pytest_forked --forked \
  tests/unit/runtime/test_lr_schedulers.py \
  -k 'test_adam_updates or test_one_cycle' \
  --torch_ver=2.12.1+cu126 --cuda_ver=12.6
DS_ACCELERATOR=cuda CUDA_VISIBLE_DEVICES=0 python -m pytest -p pytest_forked --forked \
  tests/unit/runtime/test_lr_schedulers.py::TestOneCycle \
  --torch_ver=2.12.1+cu126 --cuda_ver=12.6

Run above	Result
CPU full OneCycle selection	30 passed
GPU targeted selection	17 passed
GPU TestOneCycle integration class	13 passed

The GPU selections overlap in the four new Adam cases; these counts must not be added as unique cases.

All applicable pre-commit hooks pass on the three changed files:

pre-commit run --files deepspeed/runtime/lr_schedules.py \
  tests/unit/runtime/test_lr_schedulers.py docs/_tutorials/one-cycle.md

The independent CPU float64 two-update reproducer fails before the fix (weight 0.7387770530174719 versus reference 0.7878449067299509) and matches exactly afterward. A separate GPU torchrun engine smoke changes from a 0.04906785488128662 error to 5.960464477539063e-08 (rtol=1e-6, atol=1e-7).

No full DeepSpeed suite, full-model convergence, throughput, or peak-memory benchmark is claimed. Multi-node, ZeRO optimizer wrappers, offload, mixed precision and non-CUDA accelerators were not exercised for this scheduler change. Initial parallel scheduler attempts timed out; the completed isolated runs above are the final validation. After committing, python ci/tests_fetcher.py --base origin/master selected 44 test files in its tests/unit/v1 scope. This is a selection preview, not executed upstream CI; the targeted regression commands above provide the local validation.

Sibling paths and related work

All engine-created/client-created OneCycle schedulers use this implementation. Pipeline training advances the same engine optimizer-step path; wrapper support is centralized in get_torch_optimizer. Other scheduler classes do not cycle momentum. Their code and checkpoint formats are unchanged; these are source-inspection statements, not coverage claims.

Duplicate searches and relevant open diffs were checked. #8495/#8751 address beta2, #8589 momentum decay, #8472 CLI forwarding, #8528 LR tensor aliasing, and #8753 initial get_last_lr(). None repairs the initial beta1 endpoint. The beta2 PRs touch adjacent initialization code and should retain this maximum-endpoint selection if merged first.

AI assistance was used for investigation, implementation.

[chatgpt-codex-connector]

Codex Review Summary

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

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-30T07:36:36.916502Z We couldn’t complete this review. Comment @codex review to try again.	ddad70e	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.