Fix megatron optimizer resume: retain lr_mult param-group key

[gusario]

Summary

Resuming a run with optimizer state — finetune=false together with
mcore_model or mcore_adapter — crashes on megatron-core 0.17.x with:

ValueError: Key lr_mult (or pre_lr_mult) not found in param_group {...}

Training and checkpoint saving are unaffected; only resume fails.

Root cause

_get_param_groups builds each param group with an lr_mult field, folds it
into max_lr/min_lr, and then drops it via param_group.pop('lr_mult').

On megatron-core ≤0.16 that pop was harmless cleanup — lr_mult had already
served its purpose and nothing read it afterwards. megatron-core 0.17 changed
that: it added lr_mult to param_group_identifier_keys, which
DistributedOptimizer.sharded_state_dict(is_loading=True) →
make_needed_groups uses to match param groups when loading optimizer state.
With the key removed, that match fails and raises.

This only surfaces on resume because the load path calls sharded_state_dict
with is_loading=True (which triggers the matching roundtrip), whereas saving
uses is_loading=False and never looks at the identifier keys.

Fix

Keep lr_mult on the param group instead of popping it.

The learning-rate schedule is unaffected: OptimizerParamScheduler derives the
LR from max_lr/min_lr and never reads lr_mult. Retaining the real value
(rather than dropping or normalizing it) also keeps groups that differ only by
LR scaling — e.g. vit_lr/aligner_lr vs the LLM — distinct, so optimizer
state is matched to the correct group on resume.

Verification

megatron-core 0.17.x, multi-GPU:

1. Train with --save_steps N to produce checkpoint-N.
2. Resume with --finetune false --mcore_model <output_dir>/checkpoint-N.

Before this change the resume crashes in make_needed_groups; after it, the
optimizer (and RNG) state loads and training continues.

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Fixes optimizer-state resume on megatron-core 0.17.x. See summary above.

Experiment results

N/A — behavioral fix; verification steps above.

[gemini-code-assist]

Code Review

This pull request modifies swift/megatron/trainers/base.py to retain the lr_mult key in param_group instead of popping it, as mcore 0.17 requires it as a param-group identity key when resuming optimizer state. There are no review comments, so I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Jintao-Huang]

#9524

[Jintao-Huang]

This issue has been fixed.