From 1673e819f84ced20f89e1b49e9ea3d4d61ff24fa Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Thu, 23 Jul 2026 05:47:53 +0000 Subject: [PATCH 1/2] Guard LRRangeTest and OneCycle schedulers against zero step sizes LRRangeTest divides the step index by self.step_size and OneCycle divides cycle_first_step_size by total_size (first + second step size), both taken unvalidated from user config. A zero step size raises a bare ZeroDivisionError instead of a clear configuration error. Reject a non-positive step size at construction with a ValueError, mirroring the existing warmup_num_steps guards (#8126, #8142, #8151). Valid configs are unaffected. Signed-off-by: Ehsan Barkhordar --- deepspeed/runtime/lr_schedules.py | 6 +++++ tests/unit/runtime/test_lr_schedulers.py | 30 +++++++++++++++++++++++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/lr_schedules.py b/deepspeed/runtime/lr_schedules.py index 62dbf9f4cc6e..e0de0ab240a0 100755 --- a/deepspeed/runtime/lr_schedules.py +++ b/deepspeed/runtime/lr_schedules.py @@ -331,6 +331,9 @@ def __init__(self, else: self.min_lr = [lr_range_test_min_lr] * len(self.optimizer.param_groups) + if not isinstance(lr_range_test_step_size, int) or lr_range_test_step_size <= 0: + raise ValueError(f"lr_range_test_step_size must be a positive integer, got {lr_range_test_step_size}") + self.step_size = lr_range_test_step_size self.step_rate = lr_range_test_step_rate self.last_batch_iteration = last_batch_iteration @@ -483,6 +486,9 @@ def _initialize_cycle(self, cycle_first_step_size, cycle_second_step_size, cycle cycle_second_step_size) if cycle_second_step_size is not None else cycle_first_step_size self.total_size = cycle_first_step_size + cycle_second_step_size + if self.total_size <= 0: + raise ValueError("cycle_first_step_size + cycle_second_step_size must be positive, got " + f"{cycle_first_step_size} + {cycle_second_step_size} = {self.total_size}") self.step_ratio = cycle_first_step_size / self.total_size self.first_stair_count = cycle_first_stair_count self.second_stair_count = cycle_first_stair_count if cycle_second_stair_count is None else cycle_second_stair_count diff --git a/tests/unit/runtime/test_lr_schedulers.py b/tests/unit/runtime/test_lr_schedulers.py index a2e0de02401c..ace01a7f1b1e 100644 --- a/tests/unit/runtime/test_lr_schedulers.py +++ b/tests/unit/runtime/test_lr_schedulers.py @@ -16,7 +16,7 @@ from deepspeed.runtime.lr_schedules import CYCLE_MIN_MOM, CYCLE_MAX_MOM, DECAY_MOM_RATE from deepspeed.runtime.lr_schedules import WARMUP_DECAY_LR, TOTAL_NUM_STEPS from deepspeed.runtime.lr_schedules import WARMUP_COSINE_LR, WARMUP_MIN_RATIO, COS_MIN_RATIO, WarmupCosineLR -from deepspeed.runtime.lr_schedules import WarmupLR, WarmupDecayLR +from deepspeed.runtime.lr_schedules import WarmupLR, WarmupDecayLR, LRRangeTest, OneCycle def _verify_continuous_decrease(values): @@ -627,3 +627,31 @@ def test_warmup_cosine_lr_linear_warmup_type_produces_linear_ratios(): for step in range(warmup_num_steps): scheduler.step(step) assert scheduler.get_lr_ratio() == pytest.approx(step / warmup_num_steps) + + +@pytest.mark.parametrize("bad_step_size", [0, -5]) +def test_lr_range_test_rejects_nonpositive_step_size(bad_step_size): + # lr_range_test_step_size divides the step index in _continuous_interval and + # _staircase_interval, so the first step() with a value of 0 raises ZeroDivisionError. + # Mirror the WarmupLR positive-integer guard and reject the misconfig at construction. + param = torch.nn.Parameter(torch.zeros(1)) + optimizer = torch.optim.SGD([param], lr=0.1) + + with pytest.raises(ValueError): + LRRangeTest(optimizer, lr_range_test_step_size=bad_step_size) + + +@pytest.mark.parametrize("first, second", [(0, 0), (0, None), (-1, None)]) +def test_one_cycle_rejects_nonpositive_total_step_size(first, second): + # OneCycle divides cycle_first_step_size by total_size (cycle_first_step_size + + # cycle_second_step_size) in _initialize_cycle; a total of 0 raises ZeroDivisionError + # at construction. Reject the degenerate cycle with a clear ValueError instead. + param = torch.nn.Parameter(torch.zeros(1)) + optimizer = torch.optim.SGD([param], lr=0.1) + + with pytest.raises(ValueError): + OneCycle(optimizer, + cycle_min_lr=0.001, + cycle_max_lr=0.1, + cycle_first_step_size=first, + cycle_second_step_size=second) From 1e6da81eba5acc8de67297b71f043a10cd198fde Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Sun, 26 Jul 2026 23:10:03 +0000 Subject: [PATCH 2/2] Reject a zero-length first cycle half in OneCycle A sum-only check on cycle_first_step_size + cycle_second_step_size lets cycle_first_step_size=0 through whenever the second half is positive, and step_ratio is then 0. _get_scale_factor divides x by step_ratio, and x is 0 at every cycle boundary including the first get_lr() before any step(), so the ZeroDivisionError the guard was meant to prevent still fires. Validate the two halves separately instead: the first must be positive, the second non-negative. A zero second half is left working, since step_ratio is then 1.0 and x stays below it, and a test pins that. Signed-off-by: Ehsan Barkhordar --- deepspeed/runtime/lr_schedules.py | 10 +++++--- tests/unit/runtime/test_lr_schedulers.py | 31 ++++++++++++++++++++---- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/deepspeed/runtime/lr_schedules.py b/deepspeed/runtime/lr_schedules.py index e0de0ab240a0..191f62d97b3a 100755 --- a/deepspeed/runtime/lr_schedules.py +++ b/deepspeed/runtime/lr_schedules.py @@ -485,10 +485,14 @@ def _initialize_cycle(self, cycle_first_step_size, cycle_second_step_size, cycle cycle_second_step_size = float( cycle_second_step_size) if cycle_second_step_size is not None else cycle_first_step_size + # Both halves are validated separately: a zero-length first half leaves total_size + # positive but makes step_ratio 0, and _get_scale_factor divides by step_ratio. + if cycle_first_step_size <= 0: + raise ValueError(f"cycle_first_step_size must be positive, got {cycle_first_step_size}") + if cycle_second_step_size < 0: + raise ValueError(f"cycle_second_step_size must be non-negative, got {cycle_second_step_size}") + self.total_size = cycle_first_step_size + cycle_second_step_size - if self.total_size <= 0: - raise ValueError("cycle_first_step_size + cycle_second_step_size must be positive, got " - f"{cycle_first_step_size} + {cycle_second_step_size} = {self.total_size}") self.step_ratio = cycle_first_step_size / self.total_size self.first_stair_count = cycle_first_stair_count self.second_stair_count = cycle_first_stair_count if cycle_second_stair_count is None else cycle_second_stair_count diff --git a/tests/unit/runtime/test_lr_schedulers.py b/tests/unit/runtime/test_lr_schedulers.py index ace01a7f1b1e..5a49d2f3d22b 100644 --- a/tests/unit/runtime/test_lr_schedulers.py +++ b/tests/unit/runtime/test_lr_schedulers.py @@ -641,11 +641,12 @@ def test_lr_range_test_rejects_nonpositive_step_size(bad_step_size): LRRangeTest(optimizer, lr_range_test_step_size=bad_step_size) -@pytest.mark.parametrize("first, second", [(0, 0), (0, None), (-1, None)]) -def test_one_cycle_rejects_nonpositive_total_step_size(first, second): - # OneCycle divides cycle_first_step_size by total_size (cycle_first_step_size + - # cycle_second_step_size) in _initialize_cycle; a total of 0 raises ZeroDivisionError - # at construction. Reject the degenerate cycle with a clear ValueError instead. +@pytest.mark.parametrize("first, second", [(0, 0), (0, None), (-1, None), (0, 100), (100, -1)]) +def test_one_cycle_rejects_nonpositive_step_sizes(first, second): + # _initialize_cycle divides cycle_first_step_size by total_size, and _get_scale_factor + # then divides by the resulting step_ratio. A total of 0 raises ZeroDivisionError at + # construction; a zero first half keeps total_size positive but sets step_ratio to 0, + # so the first get_lr() raises instead. Reject both shapes with a clear ValueError. param = torch.nn.Parameter(torch.zeros(1)) optimizer = torch.optim.SGD([param], lr=0.1) @@ -655,3 +656,23 @@ def test_one_cycle_rejects_nonpositive_total_step_size(first, second): cycle_max_lr=0.1, cycle_first_step_size=first, cycle_second_step_size=second) + + +def test_one_cycle_allows_zero_second_step_size(): + # The mirror case is not degenerate: a zero second half gives step_ratio 1.0, and x + # stays below 1.0 in _get_scale_factor, so no division by zero is reachable. Pin it so + # the guard above does not grow into rejecting a working configuration. + param = torch.nn.Parameter(torch.zeros(1)) + optimizer = torch.optim.SGD([param], lr=0.1) + + scheduler = OneCycle(optimizer, + cycle_min_lr=0.001, + cycle_max_lr=0.1, + cycle_first_step_size=100, + cycle_second_step_size=0) + + assert scheduler.step_ratio == 1.0 + assert scheduler.get_lr() == [pytest.approx(0.001)] + for _ in range(3): + scheduler.step() + assert scheduler.get_lr()[0] > 0.001