Fix WarmupLR collapsing multi-group base LRs to group 0's - #8171
Merged
tohtana merged 2 commits intoJul 29, 2026
Conversation
When warmup_max_lr is unspecified, WarmupLR inherits the optimizer's lr, but the fallback computed [group['lr'] for group in param_groups][0]. The trailing [0] reduced the per-group list to group 0's scalar, which _format_param then broadcast back to every group. On a multi-group optimizer with distinct base LRs, every group warmed up to group 0's lr and the other groups' configured LRs were silently discarded. Drop the [0] so _format_param consumes the 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. 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: aa6b0edebb
ℹ️ 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".
|
|
||
| if warmup_max_lr is None: | ||
| warmup_max_lr = [group['lr'] for group in self.optimizer.param_groups][0] | ||
| warmup_max_lr = [group['lr'] for group in self.optimizer.param_groups] |
There was a problem hiding this comment.
The reviewed commit message for 614749b9539b78b9b44a002a10367858a82c734f has no Signed-off-by: trailer, so it violates the repository's DCO/signoff requirement and is likely to be rejected by the project before merge; please recreate/amend the commit with --signoff using the configured git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
tohtana
approved these changes
Jul 26, 2026
tohtana
left a comment
tohtana
left a comment
Collaborator
There was a problem hiding this comment.
Hi @ebarkhordar,
Great catch! This is a very important fix. Let's merge it.
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Jul 28, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Jul 28, 2026
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Aug 11, 2026
…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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
When
warmup_max_lris left unspecified,WarmupLRinherits the optimizer's learning rate (added in #7360). The fallback computed:The trailing
[0]reduces the per-group list to group 0's scalar._format_paramthen 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_paramreceives the full per-group list and each group warms up to its own base lr. This mirrors #7969, which fixed the same multi-group collapse in the siblingWarmupCosineLR.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 andwarmup_max_lromitted:max_lrs == [0.1, 0.1](group 1 collapsed to group 0)max_lrs == [0.1, 0.2]Added
test_warmup_lr_inherits_per_group_lr_when_max_unspecifiedintests/unit/runtime/test_lr_schedulers.py, mirroring the existingtest_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.WarmupDecayLRdefaultswarmup_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 #8166 (a different scheduler class), kept to a one-line change plus one test.