DeepSpeed
6498adc5 - Consolidate LR scheduler tensor-LR tests into a contract suite (#8206)

Commit
1 day ago
Consolidate LR scheduler tensor-LR tests into a contract suite (#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 #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 #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 #8205): the two `OneCycle` cases fail with `assert 0.01 is tensor(0.1000, dtype=torch.float64)`, exactly the symptom #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>
Author
Parents
Loading