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: sohumt123

deepspeed/runtime/lr_schedules.py

@@ -822,6 +822,11 @@ def __init__(self,
         self.last_batch_iteration = last_batch_iteration
         self.cos_min_ratio = cos_min_ratio
 
+        # Currently only support linear and log function
+        if warmup_type not in {WARMUP_LOG_RATE, WARMUP_LINEAR_RATE}:
+            logger.warning(f"Using unknown warmup_type: {warmup_type}. The increasing function "
+                           f"is set to default (log)")
+            warmup_type = WARMUP_LOG_RATE
         self.warmup_type = warmup_type
         self.warmup_min_ratio = warmup_min_ratio
         self.warmup_num_steps = max(2, warmup_num_steps)

tests/unit/runtime/test_lr_schedulers.py

@@ -577,3 +577,53 @@ def test_warmup_schedulers_reject_invalid_warmup_num_steps(scheduler_cls, bad_wa
 
     with pytest.raises(ValueError):
         scheduler_cls(**kwargs)
+
+
+def test_warmup_cosine_lr_unknown_warmup_type_falls_back_to_log():
+    # WarmupLR warns and falls back to the log warmup curve for an unrecognized
+    # warmup_type; WarmupCosineLR must do the same instead of crashing with an
+    # UnboundLocalError in get_lr_ratio on the first step.
+    param = torch.nn.Parameter(torch.zeros(1))
+    optimizer = torch.optim.Adam([param], lr=0.001)
+
+    scheduler = WarmupCosineLR(optimizer=optimizer,
+                               total_num_steps=100,
+                               warmup_num_steps=10,
+                               warmup_type="not_a_warmup_type")
+
+    assert scheduler.warmup_type == WARMUP_LOG_RATE
+
+    param_ref = torch.nn.Parameter(torch.zeros(1))
+    optimizer_ref = torch.optim.Adam([param_ref], lr=0.001)
+    scheduler_ref = WarmupCosineLR(optimizer=optimizer_ref,
+                                   total_num_steps=100,
+                                   warmup_num_steps=10,
+                                   warmup_type=WARMUP_LOG_RATE)
+
+    for step in range(15):
+        scheduler.step(step)
+        scheduler_ref.step(step)
+        assert scheduler.get_lr_ratio() == pytest.approx(scheduler_ref.get_lr_ratio())
+
+
+def test_warmup_cosine_lr_linear_warmup_type_produces_linear_ratios():
+    # No other test exercises WarmupCosineLR with warmup_type=WARMUP_LINEAR_RATE,
+    # so a regression that silently routed 'linear' through the log curve would
+    # keep the suite green. Pin the per-step warmup ratios: with
+    # warmup_min_ratio=0.0 the linear curve is step / warmup_num_steps, which
+    # clearly diverges from the log curve (e.g. 0.1 vs ~0.301 at step 1).
+    param = torch.nn.Parameter(torch.zeros(1))
+    optimizer = torch.optim.Adam([param], lr=0.001)
+    warmup_num_steps = 10
+
+    scheduler = WarmupCosineLR(optimizer=optimizer,
+                               total_num_steps=100,
+                               warmup_num_steps=warmup_num_steps,
+                               warmup_min_ratio=0.0,
+                               warmup_type=WARMUP_LINEAR_RATE)
+
+    assert scheduler.warmup_type == WARMUP_LINEAR_RATE
+
+    for step in range(warmup_num_steps):
+        scheduler.step(step)
+        assert scheduler.get_lr_ratio() == pytest.approx(step / warmup_num_steps)