Implementation of a more flexible optimizer/scheduler override system by jstjohn · Pull Request #2723 · NVIDIA/Megatron-LM

jstjohn · 2025-12-19T17:42:49Z

What does this PR do ?

This PR removes the hard-coded wd_mult, and brings back the flexibility of the previous lambda API while allowing users to override a number of different optimizer or scheduler configs per layer. This ideas were discussed in the design document "Optimizer param lookup refactor idea".

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

(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.

Signed-off-by: John St John <jstjohn@nvidia.com>

copy-pr-bot · 2025-12-19T17:42:53Z

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.

… (no longer 1-1 constraint) Signed-off-by: John St John <jstjohn@nvidia.com>

Signed-off-by: John St John <jstjohn@nvidia.com>

… but False Signed-off-by: John St John <jstjohn@nvidia.com>

Signed-off-by: John St John <jstjohn@nvidia.com>

deepakn94

This looks great. Thanks a lot, @jstjohn!

Just one minor comment.

… for the user to make use of Signed-off-by: John St John <jstjohn@nvidia.com>

Signed-off-by: John St John <jstjohn@nvidia.com>

ko3n1g · 2026-01-05T20:04:41Z

/ok to test 6ae2c41

Signed-off-by: John St. John <jstjohn@nvidia.com>

ko3n1g · 2026-01-05T20:57:20Z

/ok to test 7d1ecc3

jstjohn · 2026-01-06T20:00:48Z

/ok to test ac4470e

…NVIDIA#2723) Signed-off-by: John St John <jstjohn@nvidia.com> Signed-off-by: John St. John <jstjohn@nvidia.com>

Signed-off-by: John St John <jstjohn@nvidia.com> Signed-off-by: John St. John <jstjohn@nvidia.com> Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> Co-authored-by: John St John <jstjohn@nvidia.com> Co-authored-by: Boxiang Wang <boxiangw@nvidia.com>

…NVIDIA#2723) Signed-off-by: John St John <jstjohn@nvidia.com> Signed-off-by: John St. John <jstjohn@nvidia.com>

Signed-off-by: John St John <jstjohn@nvidia.com> Signed-off-by: John St. John <jstjohn@nvidia.com> Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> Co-authored-by: John St John <jstjohn@nvidia.com> Co-authored-by: Boxiang Wang <boxiangw@nvidia.com>

…DIA#2835)" This reverts commit b0a702b.

…e system (NVIDIA#2723)" This reverts commit 41cedd6. Signed-off-by: dimapihtar <dpihtar@gmail.com>

…e system (NVIDIA#2723)" This reverts commit 41cedd6.

…NVIDIA#2723) Signed-off-by: John St John <jstjohn@nvidia.com> Signed-off-by: John St. John <jstjohn@nvidia.com>

Implementation of a more flexible optimizer/scheduler override system

fc2fc3d

Signed-off-by: John St John <jstjohn@nvidia.com>

jstjohn requested review from a team as code owners December 19, 2025 17:42

jstjohn commented Dec 19, 2025

View reviewed changes

Comment thread megatron/training/training.py Outdated

jstjohn added 5 commits December 19, 2025 09:57

Add more comments to get param groups since the behavior is different…

9e81f47

… (no longer 1-1 constraint) Signed-off-by: John St John <jstjohn@nvidia.com>

Fix circular import in optimizer_param_scheduler

856e395

Signed-off-by: John St John <jstjohn@nvidia.com>

Add test coverage

6c816c6

Signed-off-by: John St John <jstjohn@nvidia.com>

Adding test coverage for ParamKey and ParamPredicte matching

253b9f7

Signed-off-by: John St John <jstjohn@nvidia.com>

Also add clarifying test around what happens if an attribute is there…

8c9725d

… but False Signed-off-by: John St John <jstjohn@nvidia.com>

jstjohn commented Dec 19, 2025

View reviewed changes

Comment thread tests/unit_tests/test_optimizer.py Outdated

Remove unused import and reformatting

640ee26

Signed-off-by: John St John <jstjohn@nvidia.com>

jstjohn commented Dec 19, 2025

View reviewed changes

Comment thread tests/unit_tests/test_utilities.py

jstjohn added 2 commits December 19, 2025 13:44

Add/update more comments

dbaafc5

Signed-off-by: John St John <jstjohn@nvidia.com>

Show that ordering does not impact the merging of the param group keys

d16efe4

Signed-off-by: John St John <jstjohn@nvidia.com>

deepakn94 reviewed Dec 23, 2025

View reviewed changes

jstjohn added 2 commits January 5, 2026 08:59

Add utility function for returning a standard set of config overrides…

0faf0ea

… for the user to make use of Signed-off-by: John St John <jstjohn@nvidia.com>

Have default behavior for users of mcore who do not provide any skips

6ae2c41

Signed-off-by: John St John <jstjohn@nvidia.com>

copy-pr-bot Bot temporarily deployed to nemo-ci January 5, 2026 20:04 Inactive

ko3n1g added this to the Core 0.16 milestone Jan 5, 2026

copy-pr-bot Bot had a problem deploying to nemo-ci January 5, 2026 20:05 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci January 5, 2026 20:05 Inactive

copy-pr-bot Bot temporarily deployed to test January 5, 2026 20:05 Inactive

Adding missing license headers

7d1ecc3

Signed-off-by: John St. John <jstjohn@nvidia.com>

copy-pr-bot Bot temporarily deployed to nemo-ci January 5, 2026 20:57 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci January 5, 2026 20:57 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci January 5, 2026 20:57 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci January 6, 2026 20:00 Inactive

jstjohn enabled auto-merge January 6, 2026 20:01

copy-pr-bot Bot temporarily deployed to nemo-ci January 6, 2026 20:01 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci January 6, 2026 20:01 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci January 6, 2026 20:01 Inactive

copy-pr-bot Bot temporarily deployed to test January 6, 2026 20:01 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci January 6, 2026 20:25 Inactive

jstjohn added this pull request to the merge queue Jan 6, 2026

yaoyu-33 mentioned this pull request Jan 6, 2026

Cherry-pick optimizer override refactor from #2723 #2835

Merged

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jan 6, 2026

ko3n1g added this pull request to the merge queue Jan 7, 2026

Merged via the queue into NVIDIA:main with commit 41cedd6 Jan 7, 2026
47 of 51 checks passed

BoxiangW pushed a commit to BoxiangW/Megatron-LM that referenced this pull request Jan 9, 2026

Implementation of a more flexible optimizer/scheduler override system (…

c9f9ab1

…NVIDIA#2723) Signed-off-by: John St John <jstjohn@nvidia.com> Signed-off-by: John St. John <jstjohn@nvidia.com>

ko3n1g pushed a commit to ko3n1g/Megatron-LM that referenced this pull request Jan 20, 2026

Implementation of a more flexible optimizer/scheduler override system (…

2a52e59

…NVIDIA#2723) Signed-off-by: John St John <jstjohn@nvidia.com> Signed-off-by: John St. John <jstjohn@nvidia.com>

ko3n1g added a commit to ko3n1g/Megatron-LM that referenced this pull request Feb 3, 2026

Revert "Cherry-pick optimizer override refactor from NVIDIA#2723 (NVI…

28938cc

…DIA#2835)" This reverts commit b0a702b.

dimapihtar added a commit to dimapihtar/Megatron-LM that referenced this pull request Feb 4, 2026

Revert "Implementation of a more flexible optimizer/scheduler overrid…

110152a

…e system (NVIDIA#2723)" This reverts commit 41cedd6. Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar mentioned this pull request Feb 4, 2026

Revert of #2723 #3246

Closed

6 tasks

ko3n1g added a commit to ko3n1g/Megatron-LM that referenced this pull request Feb 4, 2026

Revert "Implementation of a more flexible optimizer/scheduler overrid…

3760547

…e system (NVIDIA#2723)" This reverts commit 41cedd6.

ko3n1g mentioned this pull request Feb 4, 2026

Revert "Implementation of a more flexible optimizer/scheduler overrid… #3247

Closed

6 tasks

Navigation Menu

Uh oh!

Implementation of a more flexible optimizer/scheduler override system - #2723

Implementation of a more flexible optimizer/scheduler override system#2723
ko3n1g merged 17 commits into
NVIDIA:mainfrom
jstjohn:optimizer_param_lookup_refactor

jstjohn commented Dec 19, 2025

Uh oh!

copy-pr-bot Bot commented Dec 19, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

deepakn94 left a comment

Uh oh!

ko3n1g commented Jan 5, 2026

Uh oh!

ko3n1g commented Jan 5, 2026

Uh oh!

jstjohn commented Jan 6, 2026

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

jstjohn commented Dec 19, 2025

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 Dec 19, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

deepakn94 left a comment

Choose a reason for hiding this comment

Uh oh!

ko3n1g commented Jan 5, 2026

Uh oh!

ko3n1g commented Jan 5, 2026

Uh oh!

jstjohn commented Jan 6, 2026

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

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