Skip to content

Fix zero multimodal learning rates in Megatron #9948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tastelikefeet merged 1 commit into from
Sep 3, 2026
+22 −24
Merged

Fix zero multimodal learning rates in Megatron #9948

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 22 additions & 24 deletions swift/megatron/trainers/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,7 @@ def _get_param_groups(
logger.info_once(f'vit_lr: {vit_lr}, aligner_lr: {aligner_lr}, llm_lr: {args.lr}')
use_decoupled_learning_rate = decoupled_lr is not None

# Map (wd_mult, lr_mult, is_expert_parallel, is_decoupled_lr) to params.
# Map (wd_mult, max_lr, min_lr, is_expert_parallel, is_decoupled_lr) to params.
params_map = {}
for model_chunk in model_chunks:
visual = model_chunk.module.module.visual if is_multimodal else None
Expand All @@ -373,6 +373,7 @@ def _get_param_groups(
name.endswith('.bias') or len(param.shape) == 1
or (default_skip_embedding_weight_decay and 'embedding' in name))
_lr_mult = lr_mult
lr_override = None
if scale_lr_cond is not None:
scale_lr = scale_lr_cond(name, param)
else:
Expand All @@ -385,12 +386,12 @@ def _get_param_groups(
for k in visual._vision_tower) and not is_aligner
else:
is_aligner, is_vit = False, False
if is_vit and args.vit_lr:
if is_vit and args.vit_lr is not None:
scale_lr = True
_lr_mult = args.vit_lr / lr
elif is_aligner and args.aligner_lr:
lr_override = args.vit_lr
elif is_aligner and args.aligner_lr is not None:
scale_lr = True
_lr_mult = args.aligner_lr / lr
lr_override = args.aligner_lr

if not no_wd and not scale_lr:
wd_mult, _lr_mult = 1.0, 1.0
Expand All @@ -407,7 +408,18 @@ def _get_param_groups(
if use_decoupled_learning_rate and getattr(param, 'is_embedding_or_output_parameter', False):
is_decoupled_lr = True

key = (wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr)
if lr_override is not None:
_max_lr = lr_override
_min_lr = 0. if lr == 0. or lr_override == 0. else min_lr * lr_override / lr
elif is_decoupled_lr:
assert decoupled_lr is not None
_max_lr = decoupled_lr * _lr_mult
_min_lr = decoupled_min_lr * _lr_mult
else:
_max_lr = lr * _lr_mult
_min_lr = min_lr * _lr_mult

key = (wd_mult, _max_lr, _min_lr, is_expert_parallel, is_decoupled_lr)
if key not in params_map:
params_map[key] = []
params_map[key].append(param)
Expand All @@ -425,36 +437,22 @@ def _get_param_groups(

param_groups = []
for key in params_key:
wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr = key
wd_mult, _max_lr, _min_lr, is_expert_parallel, is_decoupled_lr = key
params = params_map[key] if key in params_map else []
param_group = {
'params': params,
'wd_mult': wd_mult,
'lr_mult': _lr_mult,
'lr_mult': 1.,
'is_expert_parallel': is_expert_parallel,
'is_decoupled_lr': is_decoupled_lr,
}
# Ensure param_group has required keys for matching when loading optimizer state
# See MegatronOptimizer._filter_and_reorder_param_groups.
if param_group_identifier_keys is not None:
assert set(param_group.keys()) - set(param_group_identifier_keys) == {'params'}
param_group['max_lr'] = _max_lr
param_group['min_lr'] = _min_lr
param_groups.append(param_group)

# Update min and max lr in param groups
# These changes are compatible with mcore 0.16.
for param_group in param_groups:
if param_group['is_decoupled_lr']:
assert decoupled_lr is not None
param_group['max_lr'] = decoupled_lr
param_group['min_lr'] = decoupled_min_lr
else:
param_group['max_lr'] = lr
param_group['min_lr'] = min_lr
lr_mult = param_group.pop('lr_mult')
# Instead of using lr_mult to control the learning rate, we directly use max_lr/min_lr.
param_group['lr_mult'] = 1.
param_group['max_lr'] *= lr_mult
param_group['min_lr'] *= lr_mult
return param_groups

@contextmanager
Expand Down
Loading