[BUG] Scheduler override overwrites per-group LR bounds

Author: jeyaletheiaCreated Sep 9, 2026Updated Sep 18, 2026
Labelscommunity-request

Describe the bug

On upstream c6be9199750845bd99a7d1a438862bb9ab29b630, --override-opt-param-scheduler overwrites explicit parameter-group LR bounds during checkpoint loading. A group configured with max_lr=5.0, min_lr=0.5 receives the global values 1.0/0.1. Engram is not involved.

Steps/Code to reproduce bug

Load optimizer and scheduler state with the override flag enabled, global bounds 1.0/0.1, and runtime group bounds 5.0/0.5. The post-load assignment replaces the group values with global arguments.

The fix PR will extend this existing test with group-bound assertions:

bash
python -m torch.distributed.run --standalone --nproc-per-node=8 -m pytest -q \
  tests/unit_tests/test_checkpointing.py::test_load_checkpoint_override_opt_param_scheduler

Those assertions are not yet upstream. With the test changes alone, upstream fails with assert 1.0 == 5.0; adding only the fix passes on every rank. The test writes and reads a real legacy torch checkpoint using mock model/optimizer fixtures.

Expected behavior

Override should preserve each group's runtime bounds, including absent overrides, while leaving ordinary checkpoint restoration unchanged.

Additional context

The proposed fix captures runtime group bounds before loading and restores them afterward. PR #3720's review identified this concern. #5213 addresses scheduler loading; this failure occurs in the later assignment.

Local test environment: 8×H800, Python 3.12.3, PyTorch 2.12.0a0+5aff3928d8.nv26.05, CUDA 13.2. This is checkpoint-path coverage, not full distributed-training restart or official CI.

@NVIDIA/mcore-oncall