wehub-resource-sync
114 lines
3.4 KiB
Python
114 lines
3.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import os
|
|
import re
|
|
import torch
|
|
from collections import OrderedDict
|
|
from .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX, MODEL_FILE_PREFIX)
|
|
|
|
|
|
def basic_folder_validation(dir):
|
|
assert os.path.exists(dir), f'{dir} path does not exist'
|
|
assert os.path.isdir(dir), f'{dir} is not a folder'
|
|
|
|
|
|
def get_files_with_prefix(all_files, prefix):
|
|
file_list = []
|
|
for file_path in all_files:
|
|
_, fname = os.path.split(file_path)
|
|
if fname.startswith(prefix):
|
|
file_list.append(file_path)
|
|
|
|
return sorted(file_list)
|
|
|
|
|
|
def validate_files(file_list):
|
|
for file in file_list:
|
|
if not os.path.isfile(file):
|
|
print(f'Error: {file} is not existent')
|
|
|
|
|
|
def get_files(dir):
|
|
file_list = []
|
|
for root, _, files in os.walk(dir):
|
|
for file in files:
|
|
file_list.append(os.path.join(root, file))
|
|
return file_list
|
|
|
|
|
|
def sort_zero_files(files, prefix):
|
|
pattern = f"{prefix}([0-9]+)_{MODEL_FILE_PREFIX}([0-9]+)"
|
|
rank_pairs = []
|
|
for f in files:
|
|
m = re.search(pattern, f)
|
|
if m:
|
|
dp_rank = int(m.group(1))
|
|
mp_rank = int(m.group(2))
|
|
rank_pairs.append((dp_rank, mp_rank, f))
|
|
else:
|
|
raise ValueError(f"Cannot parse dp_rank and mp_rank from {f}")
|
|
|
|
sorted_files = sorted(rank_pairs, key=lambda x: (x[0], x[1]))
|
|
return [f for _, _, f in sorted_files]
|
|
|
|
|
|
def get_zero_files(dir):
|
|
file_list = get_files(dir)
|
|
for prefix in [ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX]:
|
|
zero_files = get_files_with_prefix(file_list, prefix)
|
|
if len(zero_files) > 0:
|
|
return sort_zero_files(zero_files, prefix)
|
|
|
|
return []
|
|
|
|
|
|
def partition_data(data_list, num_partitions):
|
|
num_elems = len(data_list)
|
|
assert num_elems % num_partitions == 0
|
|
partition_size = num_elems // num_partitions
|
|
partitions_list = [data_list[i:i + partition_size] for i in range(0, num_elems, partition_size)]
|
|
return partitions_list
|
|
|
|
|
|
def _key_list_to_string(key_list):
|
|
return '.'.join(key_list)
|
|
|
|
|
|
def merge_state_dict(dict_a, dict_b, key_list):
|
|
merged_dict = type(dict_a)({})
|
|
|
|
for key, value in dict_b.items():
|
|
if key in dict_a.keys():
|
|
merged_dict[key] = merge_state(dict_a[key], dict_b[key], [str(key)])
|
|
else:
|
|
merged_dict[key] = value
|
|
|
|
return merged_dict
|
|
|
|
|
|
def merge_state_list(list_a, list_b, key_list):
|
|
if len(list_a) != len(list_b):
|
|
print(f'{_key_list_to_string(key_list)}')
|
|
raise ValueError(f'Cannot merge lists of different lengths, a = {len(list_a)} b = {len(list_b)}')
|
|
|
|
return [merge_state(a, b, key_list) for a, b in zip(list_a, list_b)]
|
|
|
|
|
|
def merge_state(state_a, state_b, key_list=[]):
|
|
if type(state_a) != type(state_b):
|
|
key_list_string = _key_list_to_string(key_list)
|
|
print(f'key_list = {key_list_string}')
|
|
raise ValueError(f'Cannot merge two states of types {type(state_a)} and type {type(state_b)}')
|
|
|
|
if type(state_a) in (dict, OrderedDict):
|
|
return merge_state_dict(state_a, state_b, key_list)
|
|
elif type(state_a) in (list, tuple):
|
|
return type(state_a)(merge_state_list(state_a, state_b, key_list))
|
|
elif torch.is_tensor(state_a):
|
|
return torch.cat([state_a, state_b], 0)
|
|
else:
|
|
return state_a
|