DeepSpeed
886790b5 - Validate warmup_type in WarmupCosineLR like WarmupLR (#8151)

Commit
31 days ago
Validate warmup_type in WarmupCosineLR like WarmupLR (#8151) ## Problem `WarmupLR.__init__` validates `warmup_type` and falls back to `log` with a warning for unknown values, and `WarmupDecayLR` inherits that behavior. `WarmupCosineLR` documents the same `{'log', 'linear'}` contract but stores `warmup_type` unvalidated, and `get_lr_ratio()`'s warmup branch only assigns `ratio` for the two known types. Any other value (e.g. a typo like `'Linear'` or `'cosine'`) crashes on the first `step()`: ``` File "deepspeed/runtime/lr_schedules.py", line 850, in get_lr_ratio ratio = self.warmup_min_ratio + ratio * ratio_delta UnboundLocalError: cannot access local variable 'ratio' where it is not associated with a value ``` a confusing crash deep inside the scheduler instead of the documented warn-and-default behavior its sibling classes give. ## Fix Normalize `warmup_type` in `WarmupCosineLR.__init__` exactly as `WarmupLR.__init__` already does (same warning text, same fallback to `log`). Behavior for valid `log`/`linear` values is unchanged. ## Testing Added `test_warmup_cosine_lr_unknown_warmup_type_falls_back_to_log`, which fails with the `UnboundLocalError` above on current master and passes with this change; it asserts an unknown `warmup_type` produces the same lr-ratio trajectory as an explicit `log` scheduler through warmup and into cosine decay. ``` pytest tests/unit/runtime/test_lr_schedulers.py -k "warmup_cosine or reject_invalid" 9 passed, 45 deselected ``` yapf/flake8 clean on both touched files. Follows up on the recent scheduler hardening in #8126 and #8142, which did not cover `warmup_type`. --------- Signed-off-by: Sohum Trivedi <trivsohum@gmail.com>
Author
Parents
Loading