This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,209 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
"""NPU-only Megatron checkpoint compatibility helpers.
|
||||
|
||||
MindSpeed patches Megatron's distributed optimizer on NPU, but some Megatron-Core
|
||||
checkpoint formats still need the native Megatron param_state loaders.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from contextlib import contextmanager
|
||||
|
||||
from swift.utils import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _iter_optimizer_param_groups(optimizer):
|
||||
visited = set()
|
||||
|
||||
def visit(obj):
|
||||
if obj is None or id(obj) in visited:
|
||||
return
|
||||
visited.add(id(obj))
|
||||
|
||||
param_groups = getattr(obj, 'param_groups', None)
|
||||
if param_groups is not None:
|
||||
yield param_groups
|
||||
|
||||
inner_optimizer = getattr(obj, 'optimizer', None)
|
||||
if inner_optimizer is not obj:
|
||||
yield from visit(inner_optimizer)
|
||||
|
||||
for child in getattr(obj, 'chained_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
for child in getattr(obj, 'sub_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
|
||||
yield from visit(optimizer)
|
||||
|
||||
|
||||
def _step_to_int(step):
|
||||
if isinstance(step, torch.Tensor):
|
||||
if step.numel() != 1:
|
||||
raise RuntimeError(f'Optimizer step tensor must be scalar, got shape: {tuple(step.shape)}')
|
||||
return int(step.item())
|
||||
return int(step)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _canonicalize_optimizer_steps_for_checkpoint(optimizer):
|
||||
"""Normalize NPU scalar step tensors while Megatron builds optimizer checkpoint state.
|
||||
|
||||
Megatron-Core deduplicates param-group steps with set(). Equal NPU scalar
|
||||
tensors can still hash as distinct objects, so use their numeric value only
|
||||
while sharded_state_dict() is being built and restore the optimizer in place.
|
||||
"""
|
||||
saved_steps = []
|
||||
numeric_steps = set()
|
||||
for param_groups in _iter_optimizer_param_groups(optimizer):
|
||||
for param_group in param_groups:
|
||||
if len(param_group.get('params', [])) == 0 or 'step' not in param_group:
|
||||
continue
|
||||
step = param_group['step']
|
||||
numeric_step = _step_to_int(step)
|
||||
saved_steps.append((param_group, step))
|
||||
numeric_steps.add(numeric_step)
|
||||
|
||||
if len(numeric_steps) > 1:
|
||||
raise RuntimeError(f'Inconsistent optimizer steps before checkpoint save: {sorted(numeric_steps)}')
|
||||
|
||||
canonical_step = next(iter(numeric_steps), None)
|
||||
try:
|
||||
if canonical_step is not None:
|
||||
for param_group, _step in saved_steps:
|
||||
param_group['step'] = canonical_step
|
||||
if any(isinstance(step, torch.Tensor) for _param_group, step in saved_steps):
|
||||
logger.warning(f'Canonicalized optimizer param-group step to {canonical_step} for checkpoint save.')
|
||||
yield
|
||||
finally:
|
||||
for param_group, step in saved_steps:
|
||||
param_group['step'] = step
|
||||
|
||||
|
||||
def optimizer_sharded_state_dict(optimizer, state_dict, **optim_sd_kwargs):
|
||||
with _canonicalize_optimizer_steps_for_checkpoint(optimizer):
|
||||
return optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs)
|
||||
|
||||
|
||||
def _iter_distributed_optimizers(optimizer):
|
||||
visited = set()
|
||||
|
||||
def visit(obj):
|
||||
if obj is None or id(obj) in visited:
|
||||
return
|
||||
visited.add(id(obj))
|
||||
|
||||
if hasattr(obj, 'load_parameter_state_from_dp_reshardable') or hasattr(
|
||||
obj, 'load_parameter_state_from_fully_reshardable'):
|
||||
yield obj
|
||||
return
|
||||
|
||||
for child in getattr(obj, 'chained_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
for child in getattr(obj, 'sub_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
|
||||
yield from visit(optimizer)
|
||||
|
||||
|
||||
def _has_mindspeed_patched_load_state_dict(distributed_optimizer):
|
||||
load_state_dict = getattr(type(distributed_optimizer), 'load_state_dict', None)
|
||||
return getattr(load_state_dict, '__module__', '').startswith('mindspeed.')
|
||||
|
||||
|
||||
_MEGATRON_RESHARDABLE_PARAM_STATE_LOADERS = {
|
||||
'dp_reshardable': 'load_parameter_state_from_dp_reshardable',
|
||||
'fully_reshardable': 'load_parameter_state_from_fully_reshardable',
|
||||
}
|
||||
|
||||
|
||||
def _current_npu_device():
|
||||
if hasattr(torch, 'npu'):
|
||||
return torch.device('npu', torch.npu.current_device())
|
||||
return torch.cuda.current_device()
|
||||
|
||||
|
||||
def _restore_mindspeed_optimizer_step_tensors(optimizer):
|
||||
restored_count = 0
|
||||
for param_groups in _iter_optimizer_param_groups(optimizer):
|
||||
for param_group in param_groups:
|
||||
step = param_group.get('step')
|
||||
if isinstance(step, torch.Tensor):
|
||||
continue
|
||||
if isinstance(step, (int, float)):
|
||||
param_group['step'] = torch.tensor(int(step), dtype=torch.int64, device=_current_npu_device())
|
||||
restored_count += 1
|
||||
if restored_count:
|
||||
logger.warning(f'Restored {restored_count} MindSpeed optimizer param-group step values to NPU tensors.')
|
||||
|
||||
|
||||
def _split_chained_optimizer_state_dict(chained_optimizers, state_dict):
|
||||
if isinstance(state_dict, dict):
|
||||
state_dicts = [v for _k, v in sorted(state_dict.items())]
|
||||
else:
|
||||
state_dicts = list(state_dict)
|
||||
if len(chained_optimizers) != len(state_dicts):
|
||||
raise RuntimeError(
|
||||
f'Expected {len(chained_optimizers)} entries in optimizer state dict, but got {len(state_dicts)}.')
|
||||
return state_dicts
|
||||
|
||||
|
||||
def _load_chained_optimizer_state_dict(optimizer, state_dict):
|
||||
chained_optimizers = getattr(optimizer, 'chained_optimizers', None)
|
||||
if not chained_optimizers or len(chained_optimizers) <= 1:
|
||||
return False
|
||||
|
||||
state_dicts = _split_chained_optimizer_state_dict(chained_optimizers, state_dict)
|
||||
for child_optimizer, child_state_dict in zip(chained_optimizers, state_dicts):
|
||||
load_optimizer_state_dict(child_optimizer, child_state_dict)
|
||||
synchronize_steps = getattr(optimizer, '_synchronize_steps', None)
|
||||
if synchronize_steps is not None:
|
||||
synchronize_steps()
|
||||
return True
|
||||
|
||||
|
||||
def load_optimizer_state_dict(optimizer, state_dict):
|
||||
if _load_chained_optimizer_state_dict(optimizer, state_dict):
|
||||
return
|
||||
|
||||
distributed_optimizers = list(_iter_distributed_optimizers(optimizer))
|
||||
mindspeed_patched = any(
|
||||
_has_mindspeed_patched_load_state_dict(distributed_optimizer)
|
||||
for distributed_optimizer in distributed_optimizers)
|
||||
sharding_type = state_dict.get('param_state_sharding_type') if isinstance(state_dict, dict) else None
|
||||
native_loader_name = _MEGATRON_RESHARDABLE_PARAM_STATE_LOADERS.get(sharding_type)
|
||||
if native_loader_name is None:
|
||||
optimizer.load_state_dict(state_dict)
|
||||
if mindspeed_patched:
|
||||
_restore_mindspeed_optimizer_step_tensors(optimizer)
|
||||
return
|
||||
|
||||
if not mindspeed_patched:
|
||||
optimizer.load_state_dict(state_dict)
|
||||
return
|
||||
|
||||
if len(distributed_optimizers) != 1:
|
||||
raise RuntimeError(f'MindSpeed optimizer checkpoint compatibility supports exactly one distributed optimizer, '
|
||||
f'got {len(distributed_optimizers)}.')
|
||||
distributed_optimizer = distributed_optimizers[0]
|
||||
if not hasattr(distributed_optimizer, native_loader_name):
|
||||
raise RuntimeError(f'Distributed optimizer does not support sharding type {sharding_type}.')
|
||||
|
||||
state_dict_without_param_state = dict(state_dict)
|
||||
param_state = state_dict_without_param_state.pop('param_state', None)
|
||||
state_dict_without_param_state.pop('param_state_sharding_type', None)
|
||||
if param_state is None:
|
||||
raise RuntimeError(f'Optimizer checkpoint missing param_state for sharding type {sharding_type}.')
|
||||
|
||||
logger.warning(f'Loading optimizer param_state with ms-swift compatibility path because MindSpeed '
|
||||
f'DistributedOptimizer.load_state_dict does not support {sharding_type}.')
|
||||
# Let MindSpeed restore the generic optimizer state; load the missing
|
||||
# reshardable param_state with Megatron-Core's native implementation.
|
||||
optimizer.load_state_dict(state_dict_without_param_state)
|
||||
_restore_mindspeed_optimizer_step_tensors(optimizer)
|
||||
getattr(distributed_optimizer, native_loader_name)(param_state)
|
||||
|
||||
|
||||
__all__ = ['load_optimizer_state_dict', 'optimizer_sharded_state_dict']
|
||||
Reference in New Issue
Block a user