Preserve tensor learning rates across scheduler updates - #8202
Merged
Conversation
Signed-off-by: n33levo <n33levo@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 80d036c974
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Signed-off-by: n33levo <n33levo@users.noreply.github.com>
sfc-gh-truwase
approved these changes
Aug 2, 2026
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Aug 3, 2026
Follow up deepspeedai#8202: When an optimizer starts with a tensor learning rate, OneCycle initialization replaces it with a Python scalar while applying `cycle_min_lr`. This loses the caller's tensor identity, shape, and dtype before later scheduler updates can preserve them. This PR initializes OneCycle learning rates through the existing tensor-aware update helper, matching the path used by subsequent scheduler steps. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Aug 4, 2026
…peedai#8206) ## What Replaces the three per-scheduler `*_preserves_tensor_lr` tests (`WarmupLR`, `WarmupCosineLR`, `OneCycle`) with a single parametrized contract test driven by a `TENSOR_LR_CONTRACTS` table, plus a registry-coverage assertion that fails collection when a scheduler listed in `VALID_LR_SCHEDULES` has no contract row. ## Why deepspeedai#8205 recently fixed `OneCycle._initialize_lr` bypassing `update_lr()`, which silently replaced a caller-supplied tensor LR with a Python scalar on construction. WarmupLR and WarmupCosineLR had the same class of bug fixed in deepspeedai#8202. Each fix was caught only because someone happened to write a hand-rolled regression test for that specific scheduler. This PR turns that pattern into a contract: every scheduler in `VALID_LR_SCHEDULES` is pinned by one table row asserting tensor identity / shape / dtype at init and after `step(1)`. Two failure modes are covered: 1. **Registered but buggy** (writes `group['lr'] = scalar` in its constructor) — the parametrized contract case fails on the identity assertion. 2. **Forgotten** (a new scheduler is added to `VALID_LR_SCHEDULES` without a contract row) — `test_all_schedulers_covered_by_tensor_lr_contract` fails at collection with the exact missing class name. Both guards were verified empirically against this branch: - Applied the contract test on top of `pr-8202` (i.e. before deepspeedai#8205): the two `OneCycle` cases fail with `assert 0.01 is tensor(0.1000, dtype=torch.float64)`, exactly the symptom deepspeedai#8205 fixes. The other 8 cases pass, so no false positives. - Removed the `OneCycle` row from `TENSOR_LR_CONTRACTS`: the coverage test fails with `missing tensor-LR contract for: ['OneCycle']`. ## Test footprint CPU-only, no GPU/distributed runner required: ``` pytest tests/unit/runtime/test_lr_schedulers.py -k "preserves_tensor or covered_by_tensor" # 11 passed (5 schedulers x 2 shapes + 1 coverage) ``` `pre-commit run --files tests/unit/runtime/test_lr_schedulers.py` passes (yapf, flake8, check-torchdist, check-license, codespell). --- Authored by an AI coding agent (opencode, model: glm-5.2) under interactive human direction and review. Signed-off-by: Guokai Ma <guokai.ma@intel.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
3 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 it is
DeepSpeed's tensor learning-rate support currently replaces an optimizer's LR tensor on every scheduler update. A scalar
float64LR 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()constructingtensor([lr])instead of updating the tensor supplied through the optimizer. This is a follow-up to the tensor LR support added in #7633.How it works
fill_, covering both zero-dimensional and one-element LR tensors.WarmupCosineLRbefore initialization updates the optimizer tensor, preserving the original value used by later schedule steps.WarmupLRandWarmupCosineLRregression tests that check object identity, shape, dtype, initialization, and scheduled values.E2E Top-hatting
On current
master, a zero-dimensionalfloat64LR 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.WarmupCosineLRalso 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.10Checks
pre-commit run --files deepspeed/runtime/lr_schedules.py tests/unit/runtime/test_lr_schedulers.py