Guard LRRangeTest and OneCycle schedulers against zero step sizes

[ebarkhordar]

Problem

Two learning-rate schedulers in deepspeed/runtime/lr_schedules.py divide by a step-size value taken directly from user config, with no validation, so a 0 step size crashes with a bare ZeroDivisionError instead of a clear configuration error:

• LRRangeTest divides the step index by self.step_size in _continuous_interval / _staircase_interval. With lr_range_test_step_size=0 the first step() raises ZeroDivisionError.
• OneCycle computes self.step_ratio = cycle_first_step_size / self.total_size in _initialize_cycle, where total_size = cycle_first_step_size + cycle_second_step_size. When both halves are 0, the constructor raises ZeroDivisionError.

Repro (CPU-only):

import torch
from deepspeed.runtime.lr_schedules import LRRangeTest, OneCycle

opt = lambda: torch.optim.SGD([torch.nn.Parameter(torch.zeros(1))], lr=0.1)

LRRangeTest(opt(), lr_range_test_step_size=0).step()          # ZeroDivisionError
OneCycle(opt(), cycle_min_lr=0.001, cycle_max_lr=0.1,
         cycle_first_step_size=0, cycle_second_step_size=0)   # ZeroDivisionError at construction

The sibling WarmupLR/WarmupCosineLR constructors already reject invalid warmup_num_steps this way (#8126, #8142, #8151); these two schedulers were skipped.

Fix

Validate at construction, before the division:

• LRRangeTest.__init__: reject a non-positive lr_range_test_step_size with a ValueError, mirroring the existing warmup_num_steps guard exactly.
• OneCycle._initialize_cycle: reject a non-positive total_size (cycle_first_step_size + cycle_second_step_size) with a ValueError.

No behavior change for valid configs: the guards only fire when the value is <= 0, which previously crashed (or, for a negative OneCycle total, produced a meaningless schedule).

Testing

Added CPU-only regression tests next to the existing scheduler-validation tests. They raise ZeroDivisionError (OneCycle) or silently accept the misconfig (LRRangeTest) on current master, and pass with this change:

pytest tests/unit/runtime/test_lr_schedulers.py -k "nonpositive or warmup_cosine or reject_invalid"
15 passed, 45 deselected

yapf, flake8, codespell clean via pre-commit run. DCO signed off.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1673e819f8

ℹ️ 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]

Reject zero first cycle step independently

When cycle_first_step_size=0 and cycle_second_step_size is positive, this sum-only check passes because total_size > 0, but self.step_ratio becomes 0. A get_lr() call before the first step() then enters _get_scale_factor() with x == 0 and evaluates x / self.step_ratio, raising the same ZeroDivisionError this guard is meant to prevent; DeepSpeed already exercises pre-training get_lr() via TestGetLrBeforeTrain. Please reject a zero first-step size separately, or explicitly handle a zero-length warm-up half.

Useful? React with 👍 / 👎.

[tohtana]

Hi @ebarkhordar,
Thank you for your contribution!

I found the comment from Codex bot is reasonable. Can you address it? I didn't see any other issue. Let's merge this after the issue has been addressed.

[ebarkhordar]

@tohtana Thanks for looking. The bot's note is right and it reproduces.

With cycle_first_step_size=0 and a positive cycle_second_step_size, total_size stays positive, so the sum-only check I added passes, step_ratio becomes 0.0, and _get_scale_factor then evaluates x / self.step_ratio with x == 0. That is the first get_lr() before any step(), and again at each cycle boundary, where x returns to 0.

Pushed 1e6da81. The two halves are now validated separately, first positive and second non-negative, which replaces the sum check. A zero second half is left working: step_ratio is then 1.0 and x stays below it, so the divide is never reached, and there is a test pinning that so the guard does not grow into rejecting a usable config.

How I checked it: clean python:3.11-slim container on this branch, pip install -e ., pytest tests/unit/runtime/test_lr_schedulers.py -k "nonpositive or zero_second", 8 passed. Loading lr_schedules.py directly from the checkout before the change, OneCycle(cycle_first_step_size=0, cycle_second_step_size=100) constructs and then raises ZeroDivisionError: float division by zero on get_lr(); after it, the same call raises ValueError: cycle_first_step_size must be positive, got 0.0. pre-commit run --files on both changed files is clean.

The workflows on this PR are still waiting for approval, so that evidence is from my container and not from this repo's CI.

[ebarkhordar]

Correcting my last line: five of the six workflows here are held for approval, but modal-torch-latest did run on 1e6da81 and it is red. It is not from this change.

The collect tests job checks out github.event.pull_request.base.sha and then runs ci/torch_latest.py checkout-candidate from that tree. This PR's base.sha is 886790b, where it was branched on 07-23, and ci/torch_latest.py at that commit has no checkout-candidate subcommand. It fails before argument parsing anyway, on the module level import modal at line 10, because modal is installed in the deploy job and not in collect. Job 89852805725: ModuleNotFoundError: No module named 'modal'.

The same job, running the same checkout-candidate step, is green on PRs whose base.sha is newer: #8171 at 22:03 today and #8168, both on d326520. My diff touches deepspeed/runtime/lr_schedules.py and tests/unit/runtime/test_lr_schedulers.py and nothing under ci/ or .github/.

base.sha is set when the PR is opened, so I do not think a rebase moves it, and I have not found a way to clear this leg from the branch side. Happy to do whatever is easiest for you, including reopening the same commits as a fresh PR.

[tohtana]

Thank you for the update, @ebarkhordar!
This looks good to me. The CI also looks working now.

[ebarkhordar]

Thanks for the review and the merge, @tohtana. The cycle_first_step_size=0 case the bot caught was a real hole in my first patch, so I am glad you held it until that was covered.