[BUG] Scheduler override overwrites per-group LR bounds
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:
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_schedulerThose 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
Source: NVIDIA/Megatron-LM