DeepSpeed
90e30f4d - Guard LRRangeTest and OneCycle schedulers against zero step sizes (#8166)

Commit
33 days ago
Guard LRRangeTest and OneCycle schedulers against zero step sizes (#8166) ## 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): ```python 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. --------- Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Author
Parents
Loading