DeepSpeed

6c58fb75 - Fix WarmupLR collapsing multi-group base LRs to group 0's (#8171)

Fix WarmupLR collapsing multi-group base LRs to group 0's (#8171) ## What When `warmup_max_lr` is left unspecified, `WarmupLR` inherits the optimizer's learning rate (added in #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 #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 #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>