NVIDIA / NVIDIA/Megatron-LM

[BUG] Scheduler override overwrites per-group LR bounds

Open
#7,173 0 comments 0 reactions 0 assignees View on GitHub
community-request waiting-on-maintainers
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.