Skip to content

Validate warmup_type in WarmupCosineLR like WarmupLR#8151

Open
sohumt123 wants to merge 1 commit into
deepspeedai:masterfrom
sohumt123:fix-warmupcosine-warmup-type-validation
Open

Validate warmup_type in WarmupCosineLR like WarmupLR#8151
sohumt123 wants to merge 1 commit into
deepspeedai:masterfrom
sohumt123:fix-warmupcosine-warmup-type-validation

Conversation

@sohumt123

Copy link
Copy Markdown

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.

WarmupLR warns and falls back to the log warmup curve when given an
unrecognized warmup_type, but WarmupCosineLR stored the value
unvalidated. Since get_lr_ratio() only assigns ratio for the two known
types, any other value (e.g. a typo like 'Linear') crashed with
UnboundLocalError on the first step() instead of the documented
warn-and-default behavior.

Normalize warmup_type in WarmupCosineLR.__init__ exactly as
WarmupLR.__init__ does, and add a unit test asserting the fallback
matches the log warmup curve.

Signed-off-by: Sohum Trivedi <trivsohum@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant