Guard LRRangeTest and OneCycle schedulers against zero step sizes

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

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)