wehub-resource-sync
42 lines
1.7 KiB
Python
42 lines
1.7 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import torch
|
|
from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory
|
|
from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy
|
|
|
|
from swift.utils import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def patch_torch_dist_shard(thread_count):
|
|
__init__ = TorchDistSaveShardedStrategy.__init__
|
|
|
|
def __new_init__(*args, **kwargs):
|
|
kwargs['thread_count'] = thread_count
|
|
return __init__(*args, **kwargs)
|
|
|
|
TorchDistSaveShardedStrategy.__init__ = __new_init__
|
|
|
|
|
|
def patch_merge_fn(state_dict_model):
|
|
# https://github.com/NVIDIA/Megatron-LM/issues/1380
|
|
|
|
def sh_ten_merge_fn(sub_state_dict):
|
|
with torch.no_grad():
|
|
shared_storage = sub_state_dict[0].untyped_storage()
|
|
if all(shared_storage.data_ptr() == tensor.untyped_storage().data_ptr() for tensor in sub_state_dict):
|
|
element_size = sub_state_dict[0].element_size()
|
|
total_numel = sum(tensor.numel() for tensor in sub_state_dict)
|
|
if shared_storage.nbytes() == total_numel * element_size:
|
|
dim_0 = sum(tensor.shape[0] for tensor in sub_state_dict)
|
|
shape = (dim_0, ) + sub_state_dict[0].shape[1:]
|
|
combined_tensor = torch.empty(
|
|
shape, dtype=sub_state_dict[0].dtype,
|
|
device=sub_state_dict[0].device).set_(shared_storage, 0, shape)
|
|
return combined_tensor
|
|
return torch.cat(sub_state_dict)
|
|
|
|
for v in state_dict_model.values():
|
|
if isinstance(v, ShardedTensorFactory) and 'apply_swiglu_sharded_factory' in v.merge_fn.__qualname__:
|
|
v.merge_fn = sh_ten_merge_fn
|