Fix issue where parameter groups with different min/max LRs get overridden at checkpoint load time

[jstjohn]

What does this PR do ?

Adds some missing parameter scheduler keys (like min/max lr) to the distributed optimizer load so that those PGs do not get settings from other PGs overwriting them.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• 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

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!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cspades]

To summarize, I believe the main problem this PR solves is that we have a hard-coded set of optimizer states that are used to find commonalities between parameter group states that include entries not in the hard-coded set.

This means, if you have a checkpoint that saves more than the hard-coded set:

# Checkpoint Param Group State Values
(a, b, c, d, e, f, g, ...)

they end up being hashed into a look-up table and overwriting each others values:

# Post-Hash
# param_group_identifier_keys = ('wd_mult', 'lr_mult', 'is_expert_parallel', 'is_decoupled_lr')
(a, b, c, d) 

when we need to consider ALL states, not just the previous 4:

# Current MLM Optimizer State
{'is_expert_parallel': False, 'default_config': True, 'wd_mult': 1.0, 'lr_mult': 1.0, 'is_decoupled_lr': False, 'max_lr': 0.00015, 'min_lr': 1e-05, 'lr': 0.0, 'bias_correction': True, 'betas': (0.9, 0.95), 'eps': 1e-08, 'weight_decay': 0.1}

Wouldn't having a union of all state values (as the common key) in both the initialized optimizer and loaded checkpoint be more comprehensive?

# Checkpoint State
(a, b, c, d, None, e, None, None)
# Optimizer State
(a, b, c, d, e, None, f, g)
# Intersection Size: 4

where we find the checkpoint state that matches the most entries in any particular parameter group and choose that state as the one to load into this parameter group of the optimizer?

@yuzhongw-nvidia @deepakn94 This code was previously touched by a rather old ADLR PR, am I completely off the dot here or what?

[jstjohn]

To summarize, I believe the main problem this PR solves is that we have a hard-coded set of optimizer states that are used to find commonalities between parameter group states that include entries not in the hard-coded set.

This means, if you have a checkpoint that saves more than the hard-coded set:

# Checkpoint Param Group State Values
(a, b, c, d, e, f, g, ...)

they end up being hashed into a look-up table and overwriting each others values:

# Post-Hash
# param_group_identifier_keys = ('wd_mult', 'lr_mult', 'is_expert_parallel', 'is_decoupled_lr')
(a, b, c, d) 

when we need to consider ALL states, not just the previous 4:

# Current MLM Optimizer State
{'is_expert_parallel': False, 'default_config': True, 'wd_mult': 1.0, 'lr_mult': 1.0, 'is_decoupled_lr': False, 'max_lr': 0.00015, 'min_lr': 1e-05, 'lr': 0.0, 'bias_correction': True, 'betas': (0.9, 0.95), 'eps': 1e-08, 'weight_decay': 0.1}

Wouldn't having a union of all state values (as the common key) in both the initialized optimizer and loaded checkpoint be more comprehensive?

# Checkpoint State
(a, b, c, d, None, e, None, None)
# Optimizer State
(a, b, c, d, e, None, f, g)
# Intersection Size: 4

where we find the checkpoint state that matches the most entries in any particular parameter group and choose that state as the one to load into this parameter group of the optimizer?

@yuzhongw-nvidia @deepakn94 This code was previously touched by a rather old ADLR PR, am I completely off the dot here or what?

My Response

TLDR it sounds like the risk here is that if we go with a "use all keys other than these X that we should not use" approach would be that if a new mutable key were added, it would not match and then result in groups getting dropped silently... Even worse than the current issue of groups getting over-merged and sharing the same state.

Claude's response

What the matching code is doing

Both MegatronOptimizer._filter_and_reorder_param_groups and DistributedOptimizer.load_state_dict need to pair saved param_groups with current (freshly-built) param_groups across a checkpoint load. Construction order isn't part of the checkpoint (it depends on runtime sharding / EP / parameter ordering), so they match by hashing a tuple of per-group config into a dict and looking up by that tuple.

Goal: the tuple has to cover every per-group field that distinguishes one group's behavior from another. If two groups produce the same tuple they collide in the lookup, and the later writer silently overwrites the earlier saved group's override state (max_lr, min_lr, wd_mult, etc.). The params themselves are stitched in from the current inner optimizer, so a collision doesn't lose parameters — it loads the wrong per-group config (peak LR / WD / etc.) onto one of the colliding groups, which then runs at the wrong LR at the next optimizer step. That's exactly the failure mode that produced the trunk-LR overshoot in our recipe.

Why not "all keys minus params" / "all hashable"

On a real saved checkpoint (stage2_lora_v3/iter_0014500 chained_0), every per-group dict has these keys:

betas, bias_correction, default_config, eps, is_decoupled_lr,
is_expert_parallel, lr, lr_mult, max_lr, min_lr, step, wd_mult,
weight_decay, params

Only params is non-hashable, so "hashable filter" reduces to "use all keys minus params." That's the deny-list framing.

The problem: lr, weight_decay, and step are rewritten every optimizer step:

• lr and weight_decay: OptimizerParamScheduler.step at optimizer_param_scheduler.py:297-300 writes both every iter.
• step: the inner Adam writes it every call (optimizer.py:1524).

At save time these reflect iter N. The freshly-built optimizer at load time is at iter 0 with lr not yet scheduled. So a strict hash-equality match using these keys would match zero groups across save/load. The "best partial match by intersection count" workaround that came up in the thread fixes that but reintroduces silent-mismatch risk in a new dress: ties between near-matches, ordering-dependent winners, no clean threshold.

Why a "deny-list of mutating keys" still loses

Excluding {lr, step, weight_decay, params} gives a clean self-extending rule on paper, but:

1. NeMo's pre_<key> aliasing breaks. The matcher today dealiases pre_wd_mult → wd_mult per canonical name (test: test_filter_reorder_handles_nemo_pre_prefix). With a deny-list iterating raw key names, pre_wd_mult and wd_mult hash to different tuples and NeMo-saved checkpoints stop loading. Fixing that needs either a hard-coded alias map (now you're maintaining a list again, just inverted) or a generic pre_ strip (which silently collapses any future pre_foo that's semantically distinct — re-introducing the original collision bug class).
2. Upgrade hazard. If a future mcore version adds a new per-step-mutating key (some stat-tracking field, last_grad_norm, anything) and the deny-list isn't updated, every group fails to match on the next checkpoint load.
3. Includes non-discriminating Adam-ctor keys (betas, eps, bias_correction). They don't add information inside one chained inner optimizer (same Adam ctor → same values everywhere). But they make the load fail to match if a user reloads with intentionally adjusted Adam config (e.g., new betas for fine-tuning) instead of carrying moments forward under the new betas.

What the MR does

Derive the identifier from a declarative contract:

param_group_identifier_keys = (
    *sorted(ParamGroupOverride.__annotations__.keys()),
    # = (end_wd, max_lr, min_lr, optimizer, start_wd, wd_mult)
    'lr_mult', 'is_expert_parallel', 'is_decoupled_lr',
)

The first chunk = every user-overridable per-group field, sourced from ParamGroupOverride (the TypedDict the scheduler reads via param_group.get(...)). Anyone adding a new per-group behavioral knob has to put it on ParamGroupOverride for the scheduler to see it, and the identifier picks it up automatically via __annotations__. The second chunk = the three structural flags set by _get_param_groups itself; these aren't user-overridable so they're listed explicitly.

The per-key fallback inside the matcher handles NeMo aliasing (pre_<key> falls back to <key>) because the allow-list knows the canonical name set.

On ParamGroupOverride containing a mutating key

Checked — it doesn't, today. All six fields (max_lr, min_lr, start_wd, end_wd, wd_mult, optimizer) are written only by _get_param_groups at construction time (megatron/core/optimizer/__init__.py:124,249,251,266,268); zero writes from the scheduler or inner Adam. The scheduler reads max_lr/min_lr and writes its computed value to pg['lr'], which is a separate key not on ParamGroupOverride. So the contract is naturally clean:

• ParamGroupOverride = scheduling inputs (read-only at runtime).
• pg['lr'], pg['weight_decay'], pg['step'] = scheduling/stepping outputs (written each iter), not part of the override TypedDict.

That convention isn't pinned by a test, though. So I added one:

def test_identifier_excludes_mutable_per_step_keys():
    mutating = {"lr", "weight_decay", "step"}
    overlap = mutating.intersection(param_group_identifier_keys)
    assert not overlap, ...

If a future maintainer adds lr (or anything else mutating) to ParamGroupOverride for some new use case, this test fails immediately — rather than silently re-introducing the "matches zero groups across save/load" failure mode.

Summary

• Failure mode of the current 4-tuple = silent collision → wrong config loaded → divergence at next step (the bug).
• Failure mode of a deny-list = matches break across NeMo aliasing, and silently no-match on any future mutating key mcore adds.
• Failure mode of the allow-list = a new per-group user knob added outside ParamGroupOverride is missed. Mitigation: the scheduler only reads from param_group.get(...), so by convention ParamGroupOverride is the place to add per-group config, and __annotations__ picks it up automatically; the existing test_identifier_keys_cover_all_param_group_override_fields test pins this.

Net: same self-extending property the deny-list claims, NeMo aliasing falls out for free, and the symmetric guard test on mutating keys catches the only remaining drift class.

[jiemingz]

/claude review

[claude]

LGTM

[kunlunl]

Could this break backward compatibility with older checkpoints that don’t contain max_lr or min_lr? Missing keys are normalized to None, but newly created parameter groups always have concrete default values for these fields, so their identifiers may still fail to match. Is there another migration path that handles this case?

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29541881254

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29577749870