DeepSpeed

3b641487 - Preserve tensor learning rates across scheduler updates (#8202)

Preserve tensor learning rates across scheduler updates (#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 #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>