Skip to content

[lr_scheduler] Keep completed linear schedules at zero - #6543

Open
Afloat16 wants to merge 1 commit into
hpcaitech:mainfrom
Afloat16:fix/linear-warmup-rate-floor-20261112
Open

Afloat16 wants to merge 1 commit into
hpcaitech:mainfrom
Afloat16:fix/linear-warmup-rate-floor-20261112

Conversation

@Afloat16

Copy link
Copy Markdown

What does this PR do?

Fixes #6540.

LinearWarmupLR currently continues linearly below zero after total_steps. An extra training step or continuation of a completed schedule therefore reverses parameter updates. For total_steps=4 and warmup_steps=1, the scheduler reaches zero at step 4 and then returns negative rates.

Clamp the decay numerator at zero, preserving the existing warmup and decay trajectory through the configured horizon. Completed schedules now keep a zero learning rate with both ordinary and explicit epoch stepping.

Validation

  • 17 native regression cases: upstream base 14 failed / 3 passed; patch 17 passed.
  • Hard-coded learning-rate trajectories cover zero/one/three warmup steps and both stepping forms; two real SGD parameter groups with momentum are compared with independent manually assigned reference rates.
  • Explicit epoch jumps and real optimizer/scheduler checkpoint round trips before and after completion verify the zero tail.
  • python -m pytest tests/test_optimizer/test_linear_warmup_floor.py tests/test_optimizer/test_lr_scheduler.py -q: 18 passed, including the existing scheduler save/load regression.
  • Applicable configured pre-commit hooks and git diff --check passed.
  • Python 3.12 and PyTorch 2.9.1+cpu. GPU/distributed training and the full repository suite were not run. Validation covers valid configurations with warmup_steps < total_steps.

@Afloat16
Afloat16 requested a review from a team as a code owner November 12, 2026 16:14
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]: LinearWarmupLR becomes negative after the configured training horizon

1 participant