Validate warmup_type in WarmupCosineLR like WarmupLR

[sohumt123]

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.

[ebarkhordar]

Reproduced both directions in a clean python:3.11-slim container against the real installed package (torch==2.5.1+cpu, DS_BUILD_OPS=0 pip install -e ., lr_schedules.__file__ confirmed as the checkout), not a stub:

• master 4c275f9e: WarmupCosineLR(optimizer=o, total_num_steps=100, warmup_num_steps=10, warmup_type="Linear").step(0) raises UnboundLocalError: cannot access local variable 'ratio' where it is not associated with a value, exactly as described.
• this branch a81d943b: warns and falls back to log. pytest tests/unit/runtime/test_lr_schedulers.py -k "warmup_cosine or reject_invalid" gives 9 passed, 45 deselected.
• with only lr_schedules.py reverted to master and your new test kept, that test fails, so it is pinning the fix rather than passing either way.

The asymmetry predates the class, which supports the framing: WarmupLR got this validation in #1530, and #4563 later added WarmupCosineLR copying the warmup_type usage but not the check. WarmupLR._get_gamma also ends with return 1.0 after its if/elif, which is why it degrades quietly where get_lr_ratio falls through to an unbound ratio.

One suggestion, non-blocking. The new check is a whitelist sitting directly upstream of the WARMUP_LINEAR_RATE branch in get_lr_ratio, and parsing test_lr_schedulers.py for tests that construct WarmupCosineLR shows none of them ever passes warmup_type="linear". So that curve is unpinned for this class, and a later edit to the normalization set could collapse linear into log with the file still green. On this branch the two are easy to tell apart over the first ten steps:

linear: 0.0, 0.1,   0.2,    0.3,    0.4,   0.5,    0.6,    0.7,    0.8,    0.9
log:    0.0, 0.301, 0.4771, 0.6021, 0.699, 0.7782, 0.8451, 0.9031, 0.9542, 1.0

A case asserting warmup_type=WARMUP_LINEAR_RATE still produces the linear ratios would close it, and it sits naturally beside the test you already added.

[tohtana]

Hi @sohumt123,
Thank you for submitting this PR!
@ebarkhordar, I appreciate that you help review a PR on DeepSpeed repository!

The fix looks good to me. However, the new commit added an almost duplicated test. Can you remove one?

[tohtana]

Looks good to me. Thank you @sohumt123 for the update!