Guard WarmupCosineLR against total_num_steps == warmup_num_steps (ZeroDivisionError) (#8142)
`WarmupCosineLR.get_lr_ratio` computes the cosine denominator as:
```python
real_total_steps = self.total_num_steps - self.warmup_num_steps
...
ratio = (1 + math.cos(math.pi * real_last_step / real_total_steps)) / 2
```
with no guard against `real_total_steps == 0`. When `total_num_steps ==
warmup_num_steps` (or `total_num_steps < warmup_num_steps`, which the
constructor only warns about, not rejects), the first step past warmup
makes `real_total_steps == 0` and the division raises
`ZeroDivisionError`.
Repro (CPU-only):
```python
import torch
from deepspeed.runtime.lr_schedules import WarmupCosineLR
opt = torch.optim.Adam([{"params": [torch.nn.Parameter(torch.zeros(1))], "lr": 0.01}])
sched = WarmupCosineLR(opt, total_num_steps=10, warmup_num_steps=10, cos_min_ratio=0.1)
sched.step(10) # first step past warmup -> real_total_steps == 0
sched.get_lr_ratio() # ZeroDivisionError: float division by zero
```
The sibling `WarmupDecayLR._get_gamma` already floors this exact
denominator with `max(1.0, self.total_num_steps -
self.warmup_num_steps)`. This mirrors that guard:
```python
real_total_steps = max(1, self.total_num_steps - self.warmup_num_steps)
```
There is no behavior change in the normal case (`total_num_steps >
warmup_num_steps` leaves the value unchanged). In the degenerate case
the ratio now floors to `cos_min_ratio` via the existing `max(0.0, ...)`
clamp, instead of crashing.
Added a CPU-only regression test
`test_warmup_cosine_lr_total_num_steps_equals_warmup_num_steps` next to
the existing plain `WarmupCosineLR` tests; it raises `ZeroDivisionError`
before this change and passes after. yapf and flake8 clean; DCO signed
off.
Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
DeepSpeed
2163a013 - Guard WarmupCosineLR against total_num_steps == warmup_num_steps (ZeroDivisionError) (#8142)
