Enhance optimizer state loading with runtime overrides by yhgalaxy · Pull Request #3720 · NVIDIA/Megatron-LM

yhgalaxy · 2026-03-05T14:16:13Z

Added logic to override optimizer parameter scheduler settings based on runtime arguments when option --override-opt-param-scheduler on. This includes restoring max_lr and min_lr values and synchronizing scheduler num_steps with consumed_train_samples.

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

copy-pr-bot · 2026-03-05T14:16:17Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Copilot

Pull request overview

Adds checkpoint-load behavior to re-apply runtime LR bounds and scheduler progress when --override-opt-param-scheduler/--override-opt_param-scheduler is enabled, preventing optimizer state dict values from pinning the schedule to checkpoint settings.

Changes:

Restore max_lr/min_lr on optimizer param groups from runtime args when scheduler override is enabled.
Align opt_param_scheduler.num_steps with args.consumed_train_samples and refresh LR via a no-increment scheduler step.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-05T14:23:30Z

+                for param_group in optimizer.param_groups:
+                    if (
+                        param_group.get('is_decoupled_lr', False)
+                        and args.decoupled_lr is not None
+                    ):
+                        param_group['max_lr'] = args.decoupled_lr
+                        param_group['min_lr'] = (
+                            args.decoupled_min_lr
+                            if args.decoupled_min_lr is not None
+                            else args.min_lr
+                        )
+                    else:
+                        param_group['max_lr'] = args.lr
+                        param_group['min_lr'] = args.min_lr
+                # Synchronize scheduler num_steps with consumed_train_samples
+                # to ensure lr calculation is based on current training progress
+                if opt_param_scheduler.num_steps != args.consumed_train_samples:
+                    print_rank_0(
+                        f" > WARNING: scheduler num_steps ({opt_param_scheduler.num_steps}) "
+                        f"differs from consumed_train_samples ({args.consumed_train_samples}). "
+                        f"Resetting scheduler num_steps to match consumed_train_samples."
+                    )
+                    opt_param_scheduler.num_steps = args.consumed_train_samples
+                opt_param_scheduler.step(increment=0)
+                print_rank_0(
+                    " > restored optimizer param_group max_lr/min_lr from runtime args "
+                    "because --override-opt_param-scheduler is set"
+                )


args.override_opt_param_scheduler can be true even when optimizer or opt_param_scheduler is None (many call sites invoke load_checkpoint(..., None, None)). This block will then raise (e.g., optimizer.param_groups / opt_param_scheduler.num_steps). Please guard the override logic with checks that optimizer is not None, not a stub optimizer, and opt_param_scheduler is not None before accessing them (or skip with a warning).

Suggested change

for param_group in optimizer.param_groups:

if (

param_group.get('is_decoupled_lr', False)

and args.decoupled_lr is not None

):

param_group['max_lr'] = args.decoupled_lr

param_group['min_lr'] = (

args.decoupled_min_lr

if args.decoupled_min_lr is not None

else args.min_lr

)

else:

param_group['max_lr'] = args.lr

param_group['min_lr'] = args.min_lr

# Synchronize scheduler num_steps with consumed_train_samples

# to ensure lr calculation is based on current training progress

if opt_param_scheduler.num_steps != args.consumed_train_samples:

print_rank_0(

f" > WARNING: scheduler num_steps ({opt_param_scheduler.num_steps}) "

f"differs from consumed_train_samples ({args.consumed_train_samples}). "

f"Resetting scheduler num_steps to match consumed_train_samples."

)

opt_param_scheduler.num_steps = args.consumed_train_samples

opt_param_scheduler.step(increment=0)

print_rank_0(

" > restored optimizer param_group max_lr/min_lr from runtime args "

"because --override-opt_param-scheduler is set"

)

if optimizer is None or optimizer.is_stub_optimizer or opt_param_scheduler is None:

print_rank_0(

" > WARNING: --override-opt_param-scheduler is set, but optimizer or "

"opt_param_scheduler is not available; skipping override of scheduler "

"and optimizer param_group max_lr/min_lr."

)

else:

for param_group in optimizer.param_groups:

if (

param_group.get('is_decoupled_lr', False)

and args.decoupled_lr is not None

):

param_group['max_lr'] = args.decoupled_lr

param_group['min_lr'] = (

args.decoupled_min_lr

if args.decoupled_min_lr is not None

else args.min_lr

)

else:

param_group['max_lr'] = args.lr

param_group['min_lr'] = args.min_lr

# Synchronize scheduler num_steps with consumed_train_samples

# to ensure lr calculation is based on current training progress

if opt_param_scheduler.num_steps != args.consumed_train_samples:

print_rank_0(

f" > WARNING: scheduler num_steps ({opt_param_scheduler.num_steps}) "

f"differs from consumed_train_samples ({args.consumed_train_samples}). "

f"Resetting scheduler num_steps to match consumed_train_samples."

)

opt_param_scheduler.num_steps = args.consumed_train_samples

opt_param_scheduler.step(increment=0)

print_rank_0(

" > restored optimizer param_group max_lr/min_lr from runtime args "

"because --override-opt_param-scheduler is set"

)

Copilot · 2026-03-05T14:23:30Z

+                for param_group in optimizer.param_groups:
+                    if (
+                        param_group.get('is_decoupled_lr', False)
+                        and args.decoupled_lr is not None
+                    ):
+                        param_group['max_lr'] = args.decoupled_lr
+                        param_group['min_lr'] = (
+                            args.decoupled_min_lr
+                            if args.decoupled_min_lr is not None
+                            else args.min_lr
+                        )
+                    else:
+                        param_group['max_lr'] = args.lr
+                        param_group['min_lr'] = args.min_lr


This overwrites max_lr/min_lr for all non-decoupled param groups with args.lr/args.min_lr, which will clobber any runtime per-group LR bounds (e.g., ParamGroupOverride / muP-generated max_lr/min_lr). If the intent is to restore the runtime config, consider capturing the pre-checkpoint optimizer.param_groups LR bounds before optimizer.load_state_dict(...) and reapplying those per-group values after loading (instead of setting everything to args.*).

Copilot · 2026-03-05T14:23:30Z

                    opt_param_scheduler.load_state_dict(state_dict['lr_scheduler'])
                else:
                    opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler'])
+


There appears to be trailing whitespace on this blank line. Please remove it to avoid formatter/lint noise.

Suggested change

Copilot · 2026-03-05T14:23:31Z

+                opt_param_scheduler.step(increment=0)
+                print_rank_0(
+                    " > restored optimizer param_group max_lr/min_lr from runtime args "
+                    "because --override-opt_param-scheduler is set"


The user-facing message references --override-opt_param-scheduler, but the documented/aliased flag also uses the dashed form --override-opt-param-scheduler. Consider printing the dashed form (or both) to match typical CLI usage and the PR description.

Suggested change

"because --override-opt_param-scheduler is set"

"because --override-opt-param-scheduler is set"

Copilot · 2026-03-05T14:23:31Z

+            if args.override_opt_param_scheduler:
+                for param_group in optimizer.param_groups:


This adds new runtime-override behavior for optimizer param groups / scheduler state during checkpoint load, but there are no tests covering the override_opt_param_scheduler path. Please add a unit test (likely extending tests/unit_tests/test_checkpointing.py::test_load_checkpoint) that sets the flag and verifies max_lr/min_lr and num_steps behavior (and that load_checkpoint(..., None, None) remains safe).

Suggested change

if args.override_opt_param_scheduler:

for param_group in optimizer.param_groups:

if (

args.override_opt_param_scheduler

and optimizer is not None

and not getattr(optimizer, "is_stub_optimizer", False)

and opt_param_scheduler is not None

):

for param_group in getattr(optimizer, "param_groups", []):

Phlip79 · 2026-03-05T18:55:02Z

We are changing our review process and marking all open, unlabeled PRs as draft. This change will go in effect starting once #3659 is merged.

Moving forward, all PRs will be required to start as draft PRs. If you wish to get your PR merged, mark your PR as “Ready for review”. Read more about the new process at submit.md.

gautham-kollu · 2026-04-07T23:42:10Z

@yhgalaxy Requesting to add unit-tests

dimapihtar

LGTM. Thank you!

yhgalaxy · 2026-04-17T09:45:42Z

What should I do?

Phlip79 · 2026-07-16T16:35:31Z

/claude fix

dimapihtar · 2026-07-17T11:47:52Z

Hi @yhgalaxy

Could you fix unit tests, please?

guihong-nv · 2026-07-22T19:34:05Z

@yhgalaxy please fix failing tests

guihong-nv · 2026-07-24T17:55:02Z

/ok to test b5e4908

guihong-nv · 2026-07-24T18:22:40Z

/ok to test a285c04

The optimizer state dict can overwrite per-group max_lr/min_lr values when loading a checkpoint. Previously the override path captured and restored the runtime param_group's existing max_lr/min_lr, which never applied the user's new lr settings from the command line. Re-apply runtime lr bounds directly from args after loading: regular groups get args.lr/args.min_lr, decoupled groups get args.decoupled_lr/args.decoupled_min_lr when set. Also synchronize scheduler num_steps with consumed_train_samples so lr calculation is based on current training progress. Add a unit test covering the override path, and fix the test mocks so checkpoint param_groups actually flow through save/load (MockOptimizer) and use_distributed_optimizer is disabled for the torch-format case. Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

guihong-nv · 2026-07-25T15:12:04Z

/ok to test 40234ed

guihong-nv · 2026-07-26T00:55:56Z

/ok to test 28dc589

guihong-nv · 2026-07-27T18:28:01Z

/ok to test 715ca42

ilml · 2026-08-03T20:10:35Z

/ok to test 715ca42

yaox12 · 2026-08-05T03:06:11Z

/ok to test c4904ea

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

janEbert · 2026-08-06T11:32:10Z

Minor nit: could you please make the usage of quotes consistent? I.e., in your changes, only use double quotes (which seem preferred across the code base) instead of mixing double and single quotes? Note that I don't want you to change the entire file, just your changes.

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

janEbert · 2026-08-06T22:18:14Z

Thanks!

janEbert · 2026-08-06T22:18:34Z

/ok to test 38c5d09

maanug-nv · 2026-08-12T23:02:03Z

/ok to test 30aa43f

maanug-nv · 2026-08-12T23:05:18Z

Hi @yhgalaxy . I just resolved the merge conflict to unblock, but looks like there's a formatting issue now. can you just run uv run --only-group linting bash tools/autoformat.sh to resolve this?

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

yaox12 · 2026-08-13T09:25:47Z

/ok to test efb7eb0

maanug-nv · 2026-08-14T02:19:53Z

/ok to test 9ddbe0d

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

chtruong814 · 2026-08-17T21:02:54Z

/ok to test 6086b8a

Copilot AI review requested due to automatic review settings March 5, 2026 14:16

github-actions Bot added the community-request label Mar 5, 2026

svcnvidia-nemo-ci requested a review from a team March 5, 2026 14:16

Copilot started reviewing on behalf of yhgalaxy March 5, 2026 14:16 View session

Copilot AI reviewed Mar 5, 2026

View reviewed changes

Phlip79 marked this pull request as draft March 5, 2026 18:55

yhgalaxy marked this pull request as ready for review March 6, 2026 03:38

chtruong814 added the needs-follow-up Issue needs follow-up label Mar 21, 2026

dimapihtar assigned ahmadki Mar 25, 2026

gautham-kollu requested a review from dimapihtar April 7, 2026 23:41

gautham-kollu requested review from a team as code owners April 7, 2026 23:42

chtruong814 removed the needs-follow-up Issue needs follow-up label Apr 8, 2026

dimapihtar approved these changes Apr 8, 2026

View reviewed changes

chtruong814 added the needs-follow-up Issue needs follow-up label Apr 10, 2026

ilml added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Apr 10, 2026

jaredcasper approved these changes Apr 10, 2026

View reviewed changes

chtruong814 added needs-follow-up Issue needs follow-up and removed needs-follow-up Issue needs follow-up labels Apr 11, 2026

chtruong814 added waiting-on-customer Waiting on the original author to respond needs-follow-up Issue needs follow-up and removed needs-follow-up Issue needs follow-up waiting-on-customer Waiting on the original author to respond labels Apr 18, 2026

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed needs-follow-up Issue needs follow-up labels Apr 21, 2026

maanug-nv approved these changes Jun 11, 2026

View reviewed changes

Merge branch 'main' into main

28dc589

Merge branch 'main' into main

715ca42

Merge branch 'main' into main

c4904ea

yhgalaxy and others added 2 commits August 6, 2026 19:02

Handle missing optimizer scheduler override flag

2920eb0

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

Merge branch 'main' into main

8f0330f

Use consistent quotes in scheduler override

38c5d09

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

Merge branch 'main' into main

30aa43f

yhgalaxy and others added 2 commits August 13, 2026 16:51

Remove trailing whitespace from checkpoint test

37c5f93

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

Merge branch 'main' into main

efb7eb0

Merge branch 'main' into main

9ddbe0d

Fix checkpoint override test after tokenizer update

6086b8a

Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

	"because --override-opt_param-scheduler is set"
	"because --override-opt-param-scheduler is set"

		if args.override_opt_param_scheduler:
		for param_group in optimizer.param_groups:

-            if args.override_opt_param_scheduler:
-                for param_group in optimizer.param_groups:
+            if (
+                args.override_opt_param_scheduler
+                and optimizer is not None
+                and not getattr(optimizer, "is_stub_optimizer", False)
+                and opt_param_scheduler is not None
+            ):
+                for param_group in getattr(optimizer, "param_groups", []):

Uh oh!

Conversation

yhgalaxy commented Mar 5, 2026

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Mar 5, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Phlip79 commented Mar 5, 2026

Uh oh!

gautham-kollu commented Apr 7, 2026

Uh oh!

dimapihtar left a comment

Choose a reason for hiding this comment

Uh oh!

yhgalaxy commented Apr 17, 2026

Uh oh!

Phlip79 commented Jul 16, 2026

Uh oh!

dimapihtar commented Jul 17, 2026

Uh oh!

guihong-nv commented Jul 22, 2026

Uh oh!

guihong-nv commented Jul 24, 2026

Uh oh!

guihong-nv commented Jul 24, 2026

Uh oh!

guihong-nv commented Jul 25, 2026

Uh oh!

guihong-nv commented Jul 26, 2026

Uh oh!

guihong-nv commented Jul 27, 2026

Uh oh!

ilml commented Aug 3, 2026

Uh oh!

yaox12 commented Aug 5, 2026

Uh oh!

janEbert commented Aug 6, 2026

Uh oh!

janEbert commented Aug 6, 2026

Uh oh!

janEbert commented Aug 6, 2026

Uh oh!

maanug-nv commented Aug 12, 2026

Uh oh!

maanug-nv commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yaox12 commented Aug 13, 2026

Uh oh!

maanug-nv commented Aug 14, 2026

Uh oh!

chtruong814 commented Aug 17, 2026

Uh oh!

(Step 1): Add PR label `Expert Review`

maanug-nv commented Aug 12, 2026 •

edited

Loading