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)