[BUG] Scheduler override overwrites per-group LR bounds
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**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](https://github.com/NVIDIA/Megatron-LM/blob/c6be9199750845bd99a7d1a438862bb9ab29b630/megatron/training/checkpointing.py#L3044-L3054) replaces the group values with global arguments.
The fix PR will extend this existing test with group-bound assertions:
```sh
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](https://github.com/NVIDIA/Megatron-LM/pull/3720#discussion_r2890305214) identified this concern. [#5213](https://github.com/NVIDIA/Megatron-LM/pull/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
Contributor guide
Research direction
Start with the post-load assignment in megatron/training/checkpointing.py around lines 3044-3054, then run tests/unit_tests/test_checkpointing.py::test_load_checkpoint_override_opt_param_scheduler. Extend the existing test with per-group bound assertions; done means runtime max_lr and min_lr survive override loading, including absent overrides, while ordinary checkpoint restoration remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, testing-qa
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100