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

Đang mở Phù hợp với người mới
  • Có gì không ổn?
#8,951 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

@ai-yang đang làm issue này rồi.

Từ ngày 30/10/2026.

  • #8953 của @ai-yang — đang mở

Đánh giá

Độ khó
2/5
Thời gian dự kiến
1-3 giờ
Mức phù hợp với người mới
78/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
machine-learning

Hướng nghiên cứu

Bắt đầu với deepspeed/runtime/lr_schedules.py, đặc biệt là OneCycle._initialize_momentum() và bài kiểm tra momentum khởi đầu được nhắc đến trong issue. Dùng phép tái hiện trên CPU với float64 và hai lần cập nhật Adam để kiểm tra rằng một chu kỳ mới áp dụng momentum tối đa ở lần cập nhật đầu tiên và khớp với tham chiếu độc lập; xác nhận rằng việc khởi tạo lại sau khi tiếp tục giữ nguyên các thiết lập optimizer đã được khôi phục.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Python
Star
43.3k
Fork
5.1k
Merge trung bình
4 ngày 7 giờ
Pull request đã merge (30 ngày)
66

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của deepspeedai/DeepSpeed

Tất cả issue của deepspeedai/DeepSpeed

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.