Preserve tensor learning rates in OneCycle - #8205
Merged
delock merged 1 commit intoAug 3, 2026
Merged
Conversation
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1058ade958
ℹ️ 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".
| if last_batch_iteration == -1: | ||
| for lr, group in zip(self.min_lrs, optimizer.param_groups): | ||
| group['lr'] = lr | ||
| update_lr(optimizer.param_groups, self.min_lrs) |
There was a problem hiding this comment.
Commit 21836484eba3765a36a9c467cebed3808d0943a7 is a non-merge commit, but its message has no Signed-off-by trailer. The required DCO workflow in .github/workflows/dco.yml explicitly fails commits missing this trailer, so this change cannot pass CI until the commit is recreated with --signoff.
AGENTS.md reference: AGENTS.md:L6-L8
Useful? React with 👍 / 👎.
delock
approved these changes
Aug 3, 2026
Collaborator
|
Hi @tohtana, I think a reflective test might help to prevent new lr scheduler not using |
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
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.
Follow up #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.