Guard LRRangeTest and OneCycle schedulers against zero step sizes - #8166
Conversation
LRRangeTest divides the step index by self.step_size and OneCycle divides cycle_first_step_size by total_size (first + second step size), both taken unvalidated from user config. A zero step size raises a bare ZeroDivisionError instead of a clear configuration error. Reject a non-positive step size at construction with a ValueError, mirroring the existing warmup_num_steps guards (deepspeedai#8126, deepspeedai#8142, deepspeedai#8151). Valid configs are unaffected. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1673e819f8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| cycle_second_step_size) if cycle_second_step_size is not None else cycle_first_step_size | ||
|
|
||
| self.total_size = cycle_first_step_size + cycle_second_step_size | ||
| if self.total_size <= 0: |
There was a problem hiding this comment.
When cycle_first_step_size=0 and cycle_second_step_size is positive, this sum-only check passes because total_size > 0, but self.step_ratio becomes 0. A get_lr() call before the first step() then enters _get_scale_factor() with x == 0 and evaluates x / self.step_ratio, raising the same ZeroDivisionError this guard is meant to prevent; DeepSpeed already exercises pre-training get_lr() via TestGetLrBeforeTrain. Please reject a zero first-step size separately, or explicitly handle a zero-length warm-up half.
Useful? React with 👍 / 👎.
|
Hi @ebarkhordar, I found the comment from Codex bot is reasonable. Can you address it? I didn't see any other issue. Let's merge this after the issue has been addressed. |
A sum-only check on cycle_first_step_size + cycle_second_step_size lets cycle_first_step_size=0 through whenever the second half is positive, and step_ratio is then 0. _get_scale_factor divides x by step_ratio, and x is 0 at every cycle boundary including the first get_lr() before any step(), so the ZeroDivisionError the guard was meant to prevent still fires. Validate the two halves separately instead: the first must be positive, the second non-negative. A zero second half is left working, since step_ratio is then 1.0 and x stays below it, and a test pins that. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
|
@tohtana Thanks for looking. The bot's note is right and it reproduces. With Pushed 1e6da81. The two halves are now validated separately, first positive and second non-negative, which replaces the sum check. A zero second half is left working: How I checked it: clean The workflows on this PR are still waiting for approval, so that evidence is from my container and not from this repo's CI. |
|
Correcting my last line: five of the six workflows here are held for approval, but The The same job, running the same
|
tohtana
left a comment
tohtana
left a comment
There was a problem hiding this comment.
Thank you for the update, @ebarkhordar!
This looks good to me. The CI also looks working now.
…i#8171) ## What When `warmup_max_lr` is left unspecified, `WarmupLR` inherits the optimizer's learning rate (added in deepspeedai#7360). The fallback computed: ```python warmup_max_lr = [group['lr'] for group in self.optimizer.param_groups][0] ``` The trailing `[0]` reduces the per-group list to group 0's scalar. `_format_param` then broadcasts that scalar back to every group (`[value] * len(param_groups)`). So on an optimizer with multiple parameter groups that have distinct base LRs, every group warms up to group 0's lr and the other groups' configured LRs are silently discarded. ## Fix Drop the trailing `[0]` so `_format_param` receives the full per-group list and each group warms up to its own base lr. This mirrors deepspeedai#7969, which fixed the same multi-group collapse in the sibling `WarmupCosineLR`. ## Verification Reproduced and verified on a CPU-only container against this branch (real `import deepspeed`, module resolved from the checkout). With two param groups at lr 0.1 and 0.2 and `warmup_max_lr` omitted: - before: `max_lrs == [0.1, 0.1]` (group 1 collapsed to group 0) - after: `max_lrs == [0.1, 0.2]` Added `test_warmup_lr_inherits_per_group_lr_when_max_unspecified` in `tests/unit/runtime/test_lr_schedulers.py`, mirroring the existing `test_warmup_cosine_lr_initializes_all_param_groups`. It fails on master (`assert [0.1, 0.1] == [0.1, 0.2]`) and passes with this change. `WarmupDecayLR` defaults `warmup_max_lr=0.001`, so this path only changes behavior when the value is left unspecified. Ran the repo's formatting hooks (yapf, flake8, codespell, license, end-of-file) on the changed files; all pass. Note: this is a small follow-on in the same file as my open deepspeedai#8166 (a different scheduler class), kept to a one-line change plus one test. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
|
Thanks for the review and the merge, @tohtana. The |
…ai#8201) ## The bug `OneCycle` documents four of its arguments as accepting a per-param-group list: ``` cycle_min_lr (float or list): Initial learning rate which is the lower boundary in the cycle for each parameter group. cycle_max_lr (float or list): Upper learning rate boundaries in the cycle for each parameter group. cycle_min_mom (float or list): Initial momentum which is the lower boundary in the cycle for each parameter group. cycle_max_mom (float or list): Upper momentum boundaries in the cycle for each parameter group. ``` `_initialize_lr` and `_initialize_momentum` only ever broadcast a scalar: ```python self.min_lrs = [cycle_min_lr] * len(optimizer.param_groups) ... self.min_moms = [(cycle_min_mom, 0.99)] * len(optimizer.param_groups) ``` so the documented list is written whole into every param group, and the optimizer is left holding a list where it expects a number: ``` param_group lrs after construction: [[0.001, 0.002], [0.001, 0.002]] param_group betas after construction: [([0.8, 0.85], 0.99), ([0.8, 0.85], 0.99)] scheduler.step() -> TypeError: unsupported operand type(s) for -: 'list' and 'list' optimizer.step() -> TypeError: unsupported operand type(s) for -: 'int' and 'list' ``` The second line matters: the optimizer is corrupt from construction, so even a plain `optimizer.step()` fails before the scheduler is stepped at all. This is reachable from a plain JSON config, not just the Python API. `engine.py:1550` does `scheduler(optimizer, **scheduler_params)`, so `"cycle_min_lr": [0.001, 0.002]` in `ds_config` deserializes to a Python list and lands directly in `OneCycle.__init__`. A wrong-length list is also accepted silently, where the siblings raise: ``` OneCycle: accepted 3 values for 2 param groups, no error LRRangeTest: ValueError expected 2 lr_range_test_min_lr, got 3 WarmupLR: ValueError expected 2 value for min_lr, got [0.0, 0.1, 0.2] ``` ## Why implement it rather than delete the docstring lines Deleting the four "or list" claims would be a smaller diff, but the rest of `OneCycle` is already per-group end to end: `_get_cycle_lr` zips `min_lrs` with `max_lrs`, `_get_cycle_mom` zips `min_moms` with `max_moms`, and `update_lr` walks the param groups. Only the two initializers collapse the input. Both sibling schedulers in this file implement the same documented contract, and the two most recent multi-group fixes here (deepspeedai#7969 for `WarmupCosineLR`, deepspeedai#8171 for `WarmupLR`) went in the same direction. This reads as an unfinished port rather than a design decision. ## The fix Reuse `_format_param`, which is how the siblings already honour this contract. It was defined twice, identically: as a method on `WarmupLR`, and again on `WarmupCosineLR` where nothing calls it (`_format_param` appears in only two files repo-wide, and in the test file only inside a comment). I promoted the single copy to module level next to `update_lr` and `get_torch_optimizer`, dropped the dead one, and pointed `WarmupLR` and `OneCycle` at it. Net result is 19 added, 22 removed, and one implementation of this logic instead of two. I chose promoting over leaving one-line delegate methods behind because `_format_param` is private and has no callers outside this file, so a delegate would be indirection with no consumer; happy to switch to delegates if you would rather not remove the methods. Three details worth calling out rather than leaving for review: **The momentum call has to wrap the scalar, not the tuple.** `_format_param` accepts tuples, and the default `cycle_min_mom` pairs with `0.99` into a length-2 tuple, so wrapping the existing `(cycle_min_mom, 0.99)` expression would raise at construction for 1 and 3 param groups, and for exactly 2 groups would silently write `group['betas'] = 0.8` as a float and blow up later in `_get_cycle_mom`. The correct form, which is what this PR uses, formats the scalar first: ```python self.min_moms = [(mom, 0.99) for mom in _format_param(optimizer, cycle_min_mom, 'cycle_min_mom')] ``` **Both bounds are now validated before the optimizer is touched.** `_initialize_lr` used to compute `min_lrs`, write `group['lr']`, and only then look at `cycle_max_lr`, so a bad-length `cycle_max_lr` left the param groups half updated. Moving the second `_format_param` call above the mutation loop makes the constructor all-or-nothing: ``` before: lrs after a failed ctor = [[0.001, 0.002], [0.001, 0.002]] after: ValueError, lrs after a failed ctor = [0.1, 0.2] (untouched) ``` **One token in `_format_param`'s error message.** Both copies interpolate `FileNotFoundError(param_value)` where the wording promises a count, so `WarmupLR` currently reports `expected 2 value for min_lr, got [0.0, 0.1, 0.2]`. Since the two copies are collapsing into one shared helper, I corrected it to `len(param_value)` rather than carry the typo into the surviving copy. It is the only change to `WarmupLR`'s behaviour and nothing asserts on that message (no `pytest.raises(..., match=...)` anywhere in the file); say the word and I will drop it back to verbatim. **Not claiming this is strictly safer for momentum.** Because `_format_param` accepts tuples, a betas-shaped `cycle_min_mom=(0.8, 0.999)` on a two-group optimizer goes from a loud `TypeError` to silently training with per-group momenta. That hazard already exists identically in `WarmupLR`, so I kept the behaviour symmetric rather than diverging, but it is a real trade rather than a pure win. ## Tests Added to `tests/unit/runtime/test_lr_schedulers.py` as module-level functions, matching the existing plain tests there: - `test_one_cycle_accepts_per_group_lr_and_momentum_lists`: two param groups, per-group lists for all four arguments, asserting the constructor sets each group's own lr and `betas[0]`, that the cycle peak reaches each group's own `cycle_max_lr` with momentum at its own `cycle_min_mom`, and that the bottom of the cycle returns each group to its own `cycle_max_mom`. - `test_one_cycle_rejects_wrong_length_per_group_lists`, parametrized over all four arguments. It uses `Adam` rather than `SGD` on purpose: `_initialize_momentum` returns early when `'betas' not in optimizer.defaults`, so the momentum half of the test would silently never run under SGD. `pytest` cannot start on my machine (no GPU, and the `tests/unit` conftest pulls in the distributed harness), so I ran the module-level tests in this file directly against the real `lr_schedules.py`, with the `DistributedTest` classes stripped and only `deepspeed.utils.logger` stubbed. Three runs: ``` control upstream lr_schedules.py + upstream tests 21 passed, 0 failed before upstream lr_schedules.py + these tests 21 passed, 5 failed after this branch 26 passed, 0 failed ``` All 5 failures before are the new tests, and the 21 pre-existing ones are unchanged by this diff. The `DistributedTest` OneCycle coverage (`TestOneCycle.test_lr`, `test_mom`) and the other scalar-momentum users (`test_fp16.py`, `test_bf16.py`, `test_pipeline.py`, `test_other_optimizer.py`) all pass scalars, which take the unchanged broadcast path; I am relying on CI for those since they need a GPU. Lint: `yapf` 0.40.0 with the repo's `.style.yapf` reports no diff on both files, and `flake8` with the repo's `.flake8` is clean on both (also confirmed clean on the unmodified files, so that is a real result rather than a config that checks nothing). ## Prior art No open or closed PR implements list support here. `--search` over `lr_schedules`, `_format_param`, `OneCycle`, `cycle_min_lr` and `lr scheduler list param groups` turns up deepspeedai#8151, deepspeedai#8166, deepspeedai#8171, deepspeedai#7969, deepspeedai#8179, deepspeedai#1455 and deepspeedai#4563, all merged and none touching these two initializers. No open issue covers it either; the only open `OneCycle` issue is deepspeedai#3492, a request for `CosineAnnealingLR` support. This follows deepspeedai#8179 in the same class, so to be upfront about it: that one was about the cycle shape (`_initialize_cycle` and `_get_scale_factor`), this one is about the two value initializers, and I did not see it while in there. If you would rather batch further `lr_schedules.py` work, tell me and I will hold the rest. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com> Co-authored-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
3 participants
Problem
Two learning-rate schedulers in
deepspeed/runtime/lr_schedules.pydivide by a step-size value taken directly from user config, with no validation, so a0step size crashes with a bareZeroDivisionErrorinstead of a clear configuration error:LRRangeTestdivides the step index byself.step_sizein_continuous_interval/_staircase_interval. Withlr_range_test_step_size=0the firststep()raisesZeroDivisionError.OneCyclecomputesself.step_ratio = cycle_first_step_size / self.total_sizein_initialize_cycle, wheretotal_size = cycle_first_step_size + cycle_second_step_size. When both halves are0, the constructor raisesZeroDivisionError.Repro (CPU-only):
The sibling
WarmupLR/WarmupCosineLRconstructors already reject invalidwarmup_num_stepsthis way (#8126, #8142, #8151); these two schedulers were skipped.Fix
Validate at construction, before the division:
LRRangeTest.__init__: reject a non-positivelr_range_test_step_sizewith aValueError, mirroring the existingwarmup_num_stepsguard exactly.OneCycle._initialize_cycle: reject a non-positivetotal_size(cycle_first_step_size + cycle_second_step_size) with aValueError.No behavior change for valid configs: the guards only fire when the value is
<= 0, which previously crashed (or, for a negativeOneCycletotal, produced a meaningless schedule).Testing
Added CPU-only regression tests next to the existing scheduler-validation tests. They raise
ZeroDivisionError(OneCycle) or silently accept the misconfig (LRRangeTest) on current master, and pass with this change:yapf, flake8, codespell clean via
pre-commit run. DCO signed off.