Uh oh!
There was an error while loading. Please reload this page.
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>
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 👍 / 👎.
tohtana
commented
Jul 26, 2026
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>
ebarkhordar
commented
Jul 26, 2026
@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. |
ebarkhordar
commented
Jul 26, 2026
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
There was a problem hiding this comment.
Thank you for the update, @ebarkhordar!
This looks good to me. The CI also looks working now.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
…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>
ebarkhordar
commented
Jul 29, 2026
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>
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.