зеркало из https://github.com/microsoft/DeepSpeed.git
460 строки
15 KiB
Python
460 строки
15 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import re
|
|
import collections.abc
|
|
import os
|
|
import json
|
|
from deepspeed.runtime.constants import GRADIENT_ACCUMULATION_STEPS, TRAIN_MICRO_BATCH_SIZE_PER_GPU
|
|
import itertools
|
|
import copy
|
|
|
|
from ..utils import logger
|
|
|
|
|
|
def search_error(filename):
|
|
if not os.path.exists(filename):
|
|
return "stderr.log does not exist"
|
|
with open(filename) as f:
|
|
for line in f:
|
|
for s in ["Error", "error", "ERROR"]:
|
|
idx = line.find(s)
|
|
if idx != -1:
|
|
return line[idx + len(s):].lstrip(": ")
|
|
return None
|
|
|
|
|
|
def was_interruptted(filename):
|
|
if not os.path.exists(filename):
|
|
return "stderr.log does not exist"
|
|
with open(filename) as f:
|
|
for line in f:
|
|
s = "KeyboardInterrupt"
|
|
idx = line.find(s)
|
|
if idx != -1:
|
|
return True
|
|
return False
|
|
|
|
|
|
def find_replace_str(value, replace_dict):
|
|
if not isinstance(value, str):
|
|
return str(value)
|
|
|
|
matches = re.findall(r"\$[\w]+", value)
|
|
for var in matches:
|
|
var_key = var.replace("$", "").lower()
|
|
if var_key == "nvme_path":
|
|
continue
|
|
assert var_key in replace_dict, f"unknown var key: {var_key}, in {replace_dict}"
|
|
if isinstance(replace_dict[var_key], str):
|
|
value = value.replace(var, replace_dict[var_key])
|
|
else:
|
|
assert len(matches) == 1, "unable to replace multiple non-string matches"
|
|
value = replace_dict[var_key]
|
|
return value
|
|
|
|
|
|
def find_replace(target, replace_dict):
|
|
if isinstance(target, dict):
|
|
for key, value in target.items():
|
|
if isinstance(value, str):
|
|
target[key] = find_replace_str(value, replace_dict)
|
|
if isinstance(value, list):
|
|
for i in range(len(value)):
|
|
value[i] = find_replace_str(value[i], replace_dict)
|
|
if isinstance(value, dict):
|
|
find_replace(value, replace_dict)
|
|
elif isinstance(target, list):
|
|
for i in range(len(target)):
|
|
target[i] = str(find_replace_str(target[i], replace_dict))
|
|
|
|
|
|
def get_list(val):
|
|
if not isinstance(val, list):
|
|
return [val]
|
|
else:
|
|
return val
|
|
|
|
|
|
def combine_dict(d, u):
|
|
for k, v in u.items():
|
|
if isinstance(v, collections.abc.Mapping):
|
|
d[k] = combine_dict(d.get(k, {}), v)
|
|
else:
|
|
if k not in d:
|
|
d[k] = v
|
|
else:
|
|
if not isinstance(d[k], list):
|
|
d[k] = [d[k]]
|
|
d[k].extend(i for i in get_list(v) if i not in d[k])
|
|
return d
|
|
|
|
|
|
def del_if_exists(t, d):
|
|
"""Deletes a key from a dictionary if it exists.
|
|
|
|
Args:
|
|
t (string): target key to delete
|
|
d (dict): dictionary to delete from
|
|
"""
|
|
if t in d:
|
|
del d[t]
|
|
return
|
|
for k, v in d.items():
|
|
if isinstance(v, collections.abc.Mapping):
|
|
del_if_exists(t, v)
|
|
|
|
|
|
def replace_dict(d, u, ignored_keys=[]):
|
|
"""Replaces values in dict d with values in dict u.
|
|
|
|
Args:
|
|
d (dict): the target dict to overwrite
|
|
u (dict): the dict containing the values to overwrite the target dict
|
|
|
|
Returns:
|
|
dict d with values overwritten by the corresponding ones in dict u.
|
|
"""
|
|
if u is not None:
|
|
for k, v in u.items():
|
|
if k not in ignored_keys:
|
|
if v is None:
|
|
del_if_exists(k, d)
|
|
continue
|
|
if isinstance(v, collections.abc.Mapping):
|
|
d[k] = replace_dict(d.get(k, {}), v, ignored_keys)
|
|
else:
|
|
d[k] = v
|
|
return d
|
|
|
|
|
|
def get_val_by_key(d: dict, k):
|
|
if k in d:
|
|
return d[k]
|
|
for v in d.values():
|
|
if isinstance(v, dict):
|
|
return get_val_by_key(v, k)
|
|
return None
|
|
|
|
|
|
def set_val_by_key(d: dict, k, vv):
|
|
if k in d:
|
|
d[k] = vv
|
|
for v in d.values():
|
|
if isinstance(v, dict):
|
|
set_val_by_key(v, k, vv)
|
|
|
|
|
|
def fetch_hostfile(hostfile_path):
|
|
if not os.path.isfile(hostfile_path):
|
|
logger.warning("Unable to find hostfile, will proceed with training "
|
|
"with local resources only.")
|
|
return None
|
|
|
|
# e.g., worker-0 slots=16
|
|
with open(hostfile_path, 'r') as fd:
|
|
resource_pool = collections.OrderedDict()
|
|
for line in fd.readlines():
|
|
line = line.strip()
|
|
if line == '':
|
|
# skip empty lines
|
|
continue
|
|
try:
|
|
hostname, slots = line.split()
|
|
_, slot_count = slots.split("=")
|
|
slot_count = int(slot_count)
|
|
except ValueError as err:
|
|
logger.error("Hostfile is not formatted correctly, unable to "
|
|
"proceed with training.")
|
|
raise err
|
|
if hostname in resource_pool:
|
|
logger.error("Hostfile contains duplicate hosts, unable to "
|
|
"proceed with training.")
|
|
raise ValueError("host {} is already defined".format(hostname))
|
|
resource_pool[hostname] = slot_count
|
|
|
|
return resource_pool
|
|
|
|
|
|
def validate_ds_config(config: dict):
|
|
|
|
def is_False(config: dict, key):
|
|
if config is None:
|
|
return False
|
|
return bool(config.get(key))
|
|
|
|
config_zero = config.get("zero_optimization", {})
|
|
if not config_zero:
|
|
return True
|
|
stage = config_zero.get("stage")
|
|
offload = False
|
|
if stage == 1:
|
|
return True
|
|
elif stage == 2:
|
|
if is_False(config_zero, "cpu_offload") and is_False(config_zero, "cpu_offload_params"):
|
|
return False
|
|
elif stage == 3:
|
|
offload_devices = ["cpu", "nvme"]
|
|
if config_zero.get("offload_optimizer", {}).get("device") in offload_devices:
|
|
offload = True
|
|
if config_zero.get("offload_param", {}).get("device") in offload_devices:
|
|
offload = True
|
|
else:
|
|
return True
|
|
|
|
# HF requires that "ZeRO Offload can only work with DeepSpeed optimizers"
|
|
if offload and not config.get("optimizer"):
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def remove_dupe_dicts(l):
|
|
""" Removes duplicate dictionaries from a list. Uses list comprehension and the json library to sort and stringify each dictionary and the set data type to ensure unique values. Works with nested data structures.
|
|
|
|
Args:
|
|
l (list): a list of (nested) data structures.
|
|
|
|
Returns:
|
|
A list of unique values.
|
|
"""
|
|
list_of_strings = [json.dumps(d, sort_keys=True) for d in l]
|
|
list_of_strings = set(list_of_strings)
|
|
return [json.loads(s) for s in list_of_strings]
|
|
|
|
|
|
def prune_config(config, ignored_keys=[]):
|
|
""" Prunes the input configurations
|
|
|
|
Args:
|
|
configs (dict): A configuration dictionary.
|
|
ignored_keys (list, optional): the keys of the sections to delete. Defaults to [].
|
|
|
|
Returns:
|
|
A configuration dictionary.
|
|
"""
|
|
if ignored_keys:
|
|
for k in ignored_keys:
|
|
|
|
def find_del_key(d: dict, k: str):
|
|
if k in d:
|
|
del d[k]
|
|
else:
|
|
for dd in d.values():
|
|
if isinstance(dd, dict):
|
|
find_del_key(dd, k)
|
|
|
|
find_del_key(config, k)
|
|
|
|
|
|
def prune_configs(configs, ignored_keys=[]):
|
|
""" Prunes the input list of configurations
|
|
|
|
Args:
|
|
configs (list): A list of configuration dictionaries.
|
|
ignored_keys (list, optional): the keys of the sections to delete. Defaults to [].
|
|
|
|
Returns:
|
|
A list of valid and unique configuration dictionaries.
|
|
"""
|
|
pruned_list = []
|
|
for config in configs:
|
|
prune_config(config, ignored_keys)
|
|
pruned_list.append(config)
|
|
|
|
return remove_dupe_dicts(pruned_list)
|
|
|
|
|
|
def get_tuning_keys(tuning_space: dict):
|
|
"""Outputs the list of tunable parameters in the tuning space dict.
|
|
|
|
Args:
|
|
tuning_space (dict): a configuration dictionary containing tunable parameters as lists of values.
|
|
|
|
Returns:
|
|
A list of strings
|
|
"""
|
|
tuning_keys = []
|
|
for key, val in tuning_space.items():
|
|
if isinstance(val, dict):
|
|
tuning_keys.extend(get_tuning_keys(val))
|
|
if isinstance(val, list) and len(val) > 1:
|
|
tuning_keys.append(key)
|
|
return tuning_keys
|
|
|
|
|
|
def get_all_configs(tuning_space: dict, ignore_keys=None):
|
|
""" Splits the tuning space dictionary to result in all combinations of values.
|
|
|
|
Args:
|
|
tuning_space (dict): the tuning space where tunable parameters are lists of values.
|
|
"""
|
|
|
|
def gen_combinations(d: dict):
|
|
keys, values = d.keys(), d.values()
|
|
for v in values:
|
|
if not isinstance(v, list):
|
|
v = [v]
|
|
values_choices = (gen_combinations(v) if isinstance(v, dict) else get_list(v) for v in values)
|
|
for comb in itertools.product(*values_choices):
|
|
yield dict(zip(keys, comb))
|
|
|
|
all_configs = []
|
|
ignored_key_vals = {}
|
|
for ik in ignore_keys:
|
|
ignored_key_vals[ik] = tuning_space.get(ik, {})
|
|
del_if_exists(ik, tuning_space)
|
|
for c in gen_combinations(tuning_space):
|
|
replace_dict(c, ignored_key_vals)
|
|
all_configs.append(c)
|
|
return all_configs
|
|
|
|
|
|
def canonical_name(config: dict, tuning_keys=None, prefix="", omit_val=False):
|
|
""" Generates a name from the acronyms of the tuning keys in the config dict. TRAIN_MICRO_BATCH_SIZE_PER_GPU is always included in the tuning keys.
|
|
Args:
|
|
config (dict): the config dict used to generate the name
|
|
tuning_keys (list, optional): the tuning keys used to generate the name. Defaults to None.
|
|
prefix (str, optional): a string added to the beginning of the name. Defaults to None.
|
|
"""
|
|
if TRAIN_MICRO_BATCH_SIZE_PER_GPU not in tuning_keys:
|
|
tuning_keys.append(TRAIN_MICRO_BATCH_SIZE_PER_GPU)
|
|
if GRADIENT_ACCUMULATION_STEPS not in tuning_keys:
|
|
tuning_keys.append(GRADIENT_ACCUMULATION_STEPS)
|
|
tuning_keys.sort()
|
|
|
|
def get_offload_name(offload_config):
|
|
cname = ""
|
|
if offload_config is None:
|
|
return "None_"
|
|
for key, val in offload_config.items():
|
|
key = "".join(map(lambda c: c[0], key.split('_')))
|
|
if (isinstance(val, int) or isinstance(val, float)) and val > 9000:
|
|
cname += key + '{:.1e}'.format(val) + "_"
|
|
else:
|
|
if isinstance(val, bool):
|
|
val = "T" if val else "F"
|
|
cname += f"{key}{val}_"
|
|
return cname
|
|
|
|
def get_name_by_keys(config: dict, tuning_keys=None, omit_val=False):
|
|
cname = ""
|
|
if not tuning_keys or config is None:
|
|
return cname
|
|
for key, val in config.items():
|
|
# skip the arg_mappings section when naming the exp file
|
|
if key == "arg_mappings":
|
|
continue
|
|
if key == "offload_param":
|
|
cname += "op_"
|
|
if not omit_val:
|
|
cname += get_offload_name(val)
|
|
continue
|
|
if key == "offload_optimizer":
|
|
cname += "oo_"
|
|
if not omit_val:
|
|
cname += get_offload_name(val)
|
|
continue
|
|
# recursively call the func to get name for the child dicts
|
|
if isinstance(val, dict):
|
|
n = get_name_by_keys(val, tuning_keys, omit_val=omit_val)
|
|
if n != "":
|
|
cname += n + "_"
|
|
if tuning_keys and key not in tuning_keys:
|
|
continue
|
|
|
|
key_str = "".join(map(lambda c: c[0], key.split('_')))
|
|
|
|
if not omit_val:
|
|
if (isinstance(val, int) or isinstance(val, float)) and val > 9000:
|
|
cname += key_str + '{:.1e}'.format(val) + "_"
|
|
else:
|
|
if isinstance(val, bool):
|
|
val = "T" if val else "F"
|
|
cname += f"{key_str}{val}_"
|
|
else:
|
|
cname += key_str + "_"
|
|
|
|
return cname[:-1]
|
|
|
|
name = get_name_by_keys(config, tuning_keys, omit_val=omit_val)
|
|
|
|
return prefix + (name if name != "" else "exp")
|
|
|
|
|
|
def get_first_config(config: dict):
|
|
if not config:
|
|
return None
|
|
cfg = copy.deepcopy(config)
|
|
|
|
for key, val in cfg.items():
|
|
if isinstance(val, dict):
|
|
if key == "optimizer": # use user defined optimizer which might have lists of values as params
|
|
cfg[key] = val
|
|
else:
|
|
cfg[key] = get_first_config(val)
|
|
if isinstance(val, list) and len(val) > 0:
|
|
cfg[key] = val[0]
|
|
return cfg
|
|
|
|
|
|
def write_experiments(exps: list, exps_dir: str):
|
|
exp_paths = []
|
|
for exp in exps:
|
|
exp_name = exp['name']
|
|
# write the expr config to a json file
|
|
exp_path = os.path.join(exps_dir, f'{exp_name}.json')
|
|
with open(exp_path, 'w') as fd:
|
|
|
|
json.dump(exp, fd)
|
|
exp_paths.append(exp_path)
|
|
return exp_paths
|
|
|
|
|
|
def memory_to_string(n, postfix="", units=None, precision=2):
|
|
if units is None:
|
|
if n // 10**12 > 0:
|
|
return str(round(n / 1024**4, precision)) + " T" + postfix
|
|
if n // 10**9 > 0:
|
|
return str(round(n / 1024**3, precision)) + " G" + postfix
|
|
elif n // 10**6 > 0:
|
|
return str(round(n / 1024**2, precision)) + " M" + postfix
|
|
elif n // 10**3 > 0:
|
|
return str(round(n / 1014, precision)) + " K" + postfix
|
|
else:
|
|
return str(n) + " "
|
|
else:
|
|
if units == "T":
|
|
return str(round(n / 1024**4, precision)) + " " + units
|
|
if units == "G" + postfix:
|
|
return str(round(n / 1024**3, precision)) + " " + units
|
|
elif units == "M" + postfix:
|
|
return str(round(n / 1024**2, precision)) + " " + units
|
|
elif units == "K" + postfix:
|
|
return str(round(n / 1024, precision)) + " " + units
|
|
else:
|
|
return str(n) + " "
|
|
|
|
|
|
def number_to_string(n, postfix="", units=None, precision=2):
|
|
if units is None:
|
|
if n // 10**9 > 0:
|
|
return str(round(n / 1000**3, precision)) + " B" + postfix
|
|
if n // 10**6 > 0:
|
|
return str(round(n / 1000**2, precision)) + " M" + postfix
|
|
elif n // 10**3 > 0:
|
|
return str(round(n / 1000**1, precision)) + " K" + postfix
|
|
else:
|
|
return str(n) + " "
|
|
else:
|
|
if units == "B" + postfix:
|
|
return str(round(n / 1000**3, precision)) + " " + units
|
|
elif units == "M" + postfix:
|
|
return str(round(n / 1000**2, precision)) + " " + units
|
|
elif units == "K" + postfix:
|
|
return str(round(n / 1000**1, precision)) + " " + units
|
|
else:
|
|
return str(n) + " "
|