Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

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

Open Beginner friendly
  • Is something wrong?
#8,951 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

@ai-yang is already working on this.

Since Oct 30, 2026.

  • #8953 by @ai-yang — open

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
78/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start with deepspeed/runtime/lr_schedules.py, especially OneCycle._initialize_momentum() and the starting-momentum test mentioned in the issue. Use the CPU float64 reproduction with two Adam updates to check that a fresh cycle applies maximum momentum on the first update and matches the independent reference; confirm resumed construction preserves restored optimizer settings.

Written by the indexing model from the issue text.

Description

Description

On DeepSpeed master 8d4f29c75f4a4b30d4db00a94b4aebfdff5107f3, constructing OneCycle with momentum cycling enabled sets the optimizer's initial beta1 to cycle_min_mom, while scheduler.get_mom() reports cycle_max_mom. The first actual optimizer update therefore consumes a different momentum from the inverse LR/momentum schedule.

This changes Adam's accumulated first moment. Bias correction can hide the problem in the first parameter update, but the second update diverges from the intended schedule. This is a training-correctness issue, not a security issue.

Reproduction

Run the following with DeepSpeed imported from the revision above. It uses real forward/backward and optimizer updates; no mocks or GPU are required. Beta2 is explicitly 0.99 to isolate this issue from the existing beta2 reports.

# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import torch
from deepspeed.runtime.lr_schedules import OneCycle

actual = torch.nn.Linear(1, 1, bias=False).double()
reference = torch.nn.Linear(1, 1, bias=False).double()
with torch.no_grad():
    actual.weight.fill_(1.0)
    reference.weight.fill_(1.0)
optimizer = torch.optim.Adam(actual.parameters(), lr=0.1, betas=(0.9, 0.99), foreach=False)
reference_optimizer = torch.optim.Adam(reference.parameters(), lr=0.1, betas=(0.9, 0.99), foreach=False)
scheduler = OneCycle(optimizer, cycle_min_lr=0.1, cycle_max_lr=0.2,
                     cycle_first_step_size=4, cycle_second_step_size=4,
                     cycle_min_mom=0.8, cycle_max_mom=0.9)
print('optimizer initial beta1:', optimizer.param_groups[0]['betas'][0])
print('scheduler initial beta1:', scheduler.get_mom()[0][0])
inputs = torch.ones(1, 1, dtype=torch.float64)
for lr, beta1 in [(0.1, 0.9), (0.125, 0.875)]:
    reference_optimizer.param_groups[0]['lr'] = lr
    reference_optimizer.param_groups[0]['betas'] = (beta1, 0.99)
    for model, optim in [(actual, optimizer), (reference, reference_optimizer)]:
        optim.zero_grad()
        loss = 0.5 * model(inputs).square().sum()
        loss.backward()
        optim.step()
    scheduler.step()
print('actual weight:', actual.weight.item())
print('reference weight:', reference.weight.item())
torch.testing.assert_close(actual.weight, reference.weight, rtol=0, atol=0)
print('PASS: both Adam updates match the inverse-momentum reference')

Observed output:

optimizer initial beta1: 0.8
scheduler initial beta1: 0.9
actual weight: 0.7387770530174719
reference weight: 0.7878449067299509
AssertionError: Tensor-likes are not equal!

Expected behavior

A fresh inverse cycle must apply minimum LR and maximum momentum to the first optimizer update, consistently with get_mom(). The independent two-update reference should match. Resumed construction must not overwrite restored optimizer settings.

Root cause

OneCycle._initialize_momentum() writes self.min_moms into the optimizer when last_batch_iteration == -1. The cyclic momentum calculation starts at self.max_moms. The engine invokes the optimizer before stepping the scheduler, so the discrepancy affects training before the first scheduler step.

The parameter docstring incorrectly describes the lower bound as initial momentum; this contradicts the inverse-cycling description and the existing starting-momentum test.

Environment and validation

  • Linux x86_64, Python 3.11.15, PyTorch 2.12.1+cu126; DeepSpeed source reports 0.19.9.
  • Minimal reproduction: CPU, float64, PyTorch Adam.
  • Also reproduced through deepspeed.initialize() / backward() / step() on one RTX 3090, CUDA runtime 12.6, driver 580.178.04, FP32, ZeRO stage 0, no offload.
  • Real-engine two-update weight error before the fix: 0.04906785488128662. After initializing from the maximum momentum: 5.960464477539063e-08 versus the FP32 reference (rtol=1e-6, atol=1e-7). The CPU float64 reproduction matches exactly after the fix.
  • These are short numerical-correctness checks, not full-model convergence or throughput benchmarks.
    Environment report from the validated checkout:
Recorded ds_report output (local installation paths redacted) 
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
pin_memory ............. [NO] ....... [OKAY]
dc ..................... [NO] ....... [OKAY]
 [WARNING]  Unable to locate CUTLASS. Install nvidia-cutlass, clone CUTLASS next to DeepSpeed, or set CUTLASS_PATH to the CUTLASS checkout.
evoformer_attn ......... [NO] ....... [NO]
fp_quantizer ........... [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
 [WARNING]  gds requires the dev libaio .so object and headers but these were not found.
 [WARNING]  gds: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
gds .................... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
spatial_inference ...... [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... [local path redacted]
torch version .................... 2.12.1+cu126
deepspeed install path ........... [local path redacted]
deepspeed info ................... 0.19.9, [none], [none]
torch cuda version ............... 12.6
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0
shared memory (/dev/shm) size .... 125.77 GB

Related work / duplicate check

Live issue/open-PR searches found no initial-beta1-endpoint fix. #8750, #8751 and #8495 concern preservation of beta2; #8589 concerns post-cycle momentum decay; #8472 concerns CLI forwarding; #8528 concerns tensor LR aliasing; #8753 concerns initial get_last_lr(). Their changes do not correct the initial beta1 endpoint. This report does not request those unrelated fixes.

AI assistance was used to investigate and prepare this report. The report and proposed fix were reviewed before submission.

Dominant language
Python
Stars
43.3k
Forks
5.1k
Avg merge
4d 7h
Merged PRs (30d)
66

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepspeedai/DeepSpeed

All issues in deepspeedai/DeepSpeed

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.