add support for tensor learning rate (vs scalar) - #7633
Merged
tohtana merged 2 commits intoOct 20, 2025
Conversation
This change is intended to help enable support for using a tensor learning rate value vs a scalar ones. We found this helpful in cases where the Optimizer is torch.compiled (in such cases changing the scalar LR value could cause recompilation degrading the performance). The implementation allows the model script to determine the type of LR value used , by setting the initial value.
sfc-gh-truwase
approved these changes
Oct 17, 2025
Contributor
Author
|
Thanks @sfc-gh-truwase |
Contributor
#7634 attempts to fix that, but is blocked because the CI seems not testing the right branch (yet). |
tohtana
enabled auto-merge (squash)
October 20, 2025 05:08
pull Bot
pushed a commit
to davidsolomon21cn/DeepSpeed
that referenced
this pull request
Aug 2, 2026
…8202) ## What it is DeepSpeed's tensor learning-rate support currently replaces an optimizer's LR tensor on every scheduler update. A scalar `float64` LR tensor becomes a new one-dimensional tensor during scheduler initialization, changing its identity, shape, and initially its dtype. This also leaves any caller-held reference to the supplied LR tensor pointing at the stale value. The root cause is `update_lr()` constructing `tensor([lr])` instead of updating the tensor supplied through the optimizer. This is a follow-up to the tensor LR support added in deepspeedai#7633. ## How it works - Fill the optimizer's existing LR tensor in place, matching PyTorch scheduler behavior. - Squeeze a calculated one-element tensor to the scalar value expected by `fill_`, covering both zero-dimensional and one-element LR tensors. - Snapshot tensor base LRs in `WarmupCosineLR` before initialization updates the optimizer tensor, preserving the original value used by later schedule steps. - Leave the scalar LR assignment path unchanged. - Add `WarmupLR` and `WarmupCosineLR` regression tests that check object identity, shape, dtype, initialization, and scheduled values. ## E2E Top-hatting On current `master`, a zero-dimensional `float64` LR tensor is replaced by a different one-dimensional tensor. With this change, both zero-dimensional and one-element LR tensors retain their original identity, shape, and dtype while reaching the expected warmup values. `WarmupCosineLR` also retains an independent base value, preventing its schedule from remaining at zero after initialization. A clean GitHub-hosted Ubuntu 24.04 CPU run used Python 3.11 and PyTorch 2.10: - `pytest --forked -n 4 unit/runtime/test_lr_schedulers.py --torch_ver=2.10` - **70 passed** ([run](https://github.com/n33levo/DeepSpeed/actions/runs/30714252928)) ## Checks - `pre-commit run --files deepspeed/runtime/lr_schedules.py tests/unit/runtime/test_lr_schedulers.py` - All formatting, lint, license, spelling, and custom Torch checks passed. --------- Signed-off-by: n33levo <n33levo@users.noreply.github.com> Co-authored-by: n33levo <n33levo@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
4 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.
This change is intended to help enable support for using a tensor learning rate value vs a scalar ones.
We found this helpful in cases where the optimizer is torch.compiled (in such cases changing the scalar LR value could cause recompilation degrading the performance).
The implementation allows the model script to determine the type of LR value used by setting the initial value.