Fix WarmupLR collapsing multi-group base LRs to group 0's - #8171
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
left a comment
tohtana
left a comment
There was a problem hiding this comment.
Hi @ebarkhordar,
Great catch! This is a very important fix. Let's merge it.
2 participants
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.