diff --git a/lib/python3.12/site-packages/deepspeed/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f1d85b8dd974f5c45c6a733f60bdd1c37282bb0 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6426923f70703c619176b03c58ae560e7d151d74 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/__pycache__/env_report.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/__pycache__/env_report.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20983f2878e64af7be1f490ebe39ee9085ee5cba Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/__pycache__/env_report.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/__pycache__/git_version_info.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/__pycache__/git_version_info.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..846cf31cabd6e0d52c0fdbc2bc9dd4abb52021fb Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/__pycache__/git_version_info.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd38fb8e830cf09bd7949409c9f210b368cf4ed5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/__init__.py b/lib/python3.12/site-packages/deepspeed/autotuning/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..73d8153bab920a8e6865a7fd4e59a068e9a2993b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .autotuner import Autotuner diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d883f61afe7c064346e1940a0f1f42a9badaa480 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/autotuner.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/autotuner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eff84728b0c0dd602445b167162832979136917a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/autotuner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e32dbd4b0fb83fd99abbce33c6f952d8c0464d0 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87c74bb67ec4237756109fe80da18b188805f80a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/scheduler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3961212579b099e06d7e5a52166cdda0a2da26eb Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..349b47018f116736eb6ce2669486f383ad082990 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/autotuner.py b/lib/python3.12/site-packages/deepspeed/autotuning/autotuner.py new file mode 100644 index 0000000000000000000000000000000000000000..419e92bd44ae75b9dd7aecc70fe7786c3d26e715 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/autotuner.py @@ -0,0 +1,1113 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import shutil +import subprocess +import time +import datetime +import math +import hjson + +from ..runtime.config_utils import dict_raise_error_on_duplicate_keys +from ..runtime.constants import * + +from ..runtime.zero.config import ZERO_OPTIMIZATION, ZeroStageEnum +from ..utils import logger +from .config import DeepSpeedAutotuningConfig +from .constants import * +from .scheduler import ResourceManager +from .tuner import GridSearchTuner, RandomTuner, ModelBasedTuner +from .utils import * +from deepspeed.accelerator import get_accelerator + +try: + from tabulate import tabulate +except ImportError: + tabulate = None + +try: + import mlflow + has_mlflow = True +except Exception as e: + has_mlflow = False + +ZERO_OPTIMIZATION_STAGE = "stage" +OFFLOAD_OPTIMIZER = "offload_optimizer" +OFFLOAD_PARAM = "offload_param" +ZERO_OPTIMIZATION_STAGE_DEFAULT = ZeroStageEnum.disabled + + +class Autotuner: + """The DeepSpeed Autotuner automatically discovers the optimal DeepSpeed configuration that delivers good training speed. The Autotuner uses model information, system information, and heuristics to efficiently tune system knobs that affect compute and memory efficiencies, such as ZeRO optimization stages, micro-batch sizes, and many other ZeRO optimization configurations. It not only reduces the time and resources user spend on tuning, but also can discover configurations better than hand-tuned methods. + Autotuning with DeepSpeed requires no code change from DeepSpeed users. Please refer to the README for usage details. + """ + + def __init__(self, args, active_resources): + self.args = args + self.selected_exp_dir = None + + assert tabulate is not None, "Missing required package `tabulate`, please install with `pip install deepspeed[autotuning]`." + + logger.debug(f"autotuning args={args}") + + self.user_config = self._get_user_config(args.user_args) + assert self.user_config is not None, "DeepSpeed configuration is not provided" + + self.autotuning_config = DeepSpeedAutotuningConfig(self.user_config) + if self.user_config[AUTOTUNING]: + if AUTOTUNING_EXPS_DIR in self.user_config[AUTOTUNING].keys(): + del self.user_config[AUTOTUNING][AUTOTUNING_EXPS_DIR] + if AUTOTUNING_RESULTS_DIR in self.user_config[AUTOTUNING].keys(): + del self.user_config[AUTOTUNING][AUTOTUNING_RESULTS_DIR] + + self.exps_dir = self.autotuning_config.exps_dir + if self.autotuning_config.overwrite and os.path.exists(self.exps_dir): + shutil.rmtree(self.exps_dir, ignore_errors=True) + if not os.path.exists(self.exps_dir): + try: + os.makedirs(self.exps_dir, exist_ok=True) + logger.info(f"Created autotuning experiments directory: {self.exps_dir}") + except: + logger.error( + f"Failed to create {self.exps_dir}, please check exps_dir in the autotuning config file is accessible by all the nodes in the job." + ) + exit(-1) + + self.results_dir = self.autotuning_config.results_dir + if self.autotuning_config.overwrite and os.path.exists(self.results_dir): + shutil.rmtree(self.results_dir, ignore_errors=True) + if not os.path.exists(self.results_dir): + try: + os.makedirs(self.results_dir, exist_ok=True) + logger.info(f"Created autotuning results directory: {self.exps_dir}") + except: + logger.error( + f"Failed to create {self.results_dir}, please check results_dir in the autotuning config file is accessible by all the nodes in the job." + ) + exit(-1) + + # set the active resource for the autotuner resource manager + self.rm = self._get_resource_manager(active_resources) + + # get resource requirement for each autotuning experiment + self.exp_num_nodes, self.exp_num_gpus = self._get_exp_resources(args) + + assert self.exp_num_gpus <= self.rm.num_gpus_per_node, "num_gpus in the autotuning configuration must not be less than the --num_gpus value in the train script if any" + assert self.exp_num_nodes <= len( + self.rm.nodes + ), "num_nodes in the autotuning configuration must not be less than the --num_nodes value in the train script if any" + + self.records = {} + self.optimal_cmd = None + self.optimal_ds_config = None + + self.mlflow_parent_id = None + + def print_tuning_results(self): + """Print the autotuning results in tabular format. + """ + best_space_records = self.get_best_space_records() + tab = [] + if best_space_records: + for key, val in best_space_records.items(): + if not val: + continue + row = [] + row.append(key) + num_exps = 0 + if key == GLOBAL_TUNING_SPACE: + cnt = 0 + for k, v in best_space_records.items(): + if k != GLOBAL_TUNING_SPACE: + cnt += v[2] + num_exps = cnt + else: + num_exps = val[2] + row.append(num_exps) + row.append(val[1]) + row.append(val[0]['name']) + tab.append(row) + summary = tabulate(tab, + headers=["tuning_space", "num_experiments", "best_metric_val", "best_exp_name"], + tablefmt="pipe") + print(summary) + with open(os.path.join(self.results_dir, 'summary.txt'), 'w', buffering=BUFSIZE) as fd: + fd.write(summary) + fd.flush() + os.fsync(fd) + + if GLOBAL_TUNING_SPACE in best_space_records: + best_exp, best_metric_val, total_num_exps = best_space_records[GLOBAL_TUNING_SPACE] + if best_exp: + logger.info( + f"{best_exp['name']} is the optimal setup after tuning. The exp result is at {best_exp['result_dir']}." + ) + else: + logger.info(f"No optimal setup is found. Please check that experiments were run successfully.") + tuning_duration = datetime.timedelta(seconds=(time.time() - self.start_time)) + + logger.info(f"Tuning completed in {tuning_duration}") + with open(os.path.join(self.results_dir, 'summary.txt'), 'a') as f: + f.write( + f"\n\nTuning completed in {tuning_duration}. Total number of experiments: {self.rm.experiment_count - 1}." + ) + f.flush() + + def _get_user_config(self, user_args): + """Get DeepSpeed configuration from the user arguments passed to the launcher. + + Args: + user_args ([list]): user arguments passed to the DeepSpeed launcher + + Returns: + [dict]: DeepSpeed configuration dictionary + """ + user_config_file = None + if "--deepspeed_config" in user_args: + idx = user_args.index("--deepspeed_config") + assert ".json" in user_args[ + idx + 1], "DeepSpeed --deepspeed_config requires a json file to specify the configuration" + + user_config_file = user_args[idx + 1] + elif "--deepspeed" in user_args: + idx = user_args.index("--deepspeed") + if ".json" in user_args[idx + 1]: + user_config_file = user_args[idx + 1] + + logger.debug(f"user_config_file = {user_config_file}") + if user_config_file is not None: + assert os.path.isfile(user_config_file), "DeepSpeed configuration file: {} is not an existing file".format( + user_config_file) + if os.path.exists(user_config_file): + return json.load(open(user_config_file, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys) + + return None + + def _get_resource_manager(self, active_resources): + """Initialize and return a resource manager + + Args: + active_resources ([dict]): A dictionary of hostname and its slots (GPUs), e.g. {"worker-0": "0,1,2,3,4,5,6,7,8"} + + Raises: + RuntimeError: raises the error if no GPU is available + + Returns: + [ResourceManager]: A resource manager that schedules and runs autotuning experiments. + """ + logger.info(f"active_resources = {active_resources}") + + hosts = [] + ngpus_per_node = 100 + for hostname, slots in active_resources.items(): + hosts.append(hostname) + ngpus_per_node = min(len(slots), ngpus_per_node) + + assert ngpus_per_node > 0, "no gpu is available" + + return ResourceManager(args=self.args, + hosts=hosts, + num_gpus_per_node=ngpus_per_node, + results_dir=self.results_dir, + exps_dir=self.exps_dir, + arg_mappings=self.autotuning_config.arg_mappings) + + def _get_exp_resources(self, args): + """Get resource requirement for each autotuning experiment + + Args: + args (dict): user args + + Returns: + num_nodes, num_gpus: the number of gpus and number of nodes used in the autotuning experiments + """ + if args.num_nodes > 0: + num_nodes = args.num_nodes + else: + num_nodes = len(self.rm.nodes) + + if args.num_gpus > 0: + num_gpus = args.num_gpus + else: + num_gpus = self.rm.num_gpus_per_node + + return num_nodes, num_gpus + + def metric(self): + return self.autotuning_config.metric + + def fast_enabled(self): + return self.autotuning_config.fast + + def max_train_batch_size(self): + return self.autotuning_config.max_train_batch_size + + def mp_size(self): + return self.autotuning_config.mp_size + + def max_train_micro_batch_size_per_gpu(self): + if self.max_train_batch_size() and self.max_train_batch_size( + ) > 0: # if the user specifies a max_train_batch_size + max_train_micro_batch_size = self.max_train_batch_size() * self.mp_size() // ( + self.exp_num_gpus * self.exp_num_nodes) # gradient accumulation steps >=1 + return min(self.autotuning_config.max_train_micro_batch_size_per_gpu, max_train_micro_batch_size) + else: + return self.autotuning_config.max_train_micro_batch_size_per_gpu + + def min_train_micro_batch_size_per_gpu(self): + return self.autotuning_config.min_train_micro_batch_size_per_gpu + + def num_tuning_micro_batch_sizes(self): + return self.autotuning_config.num_tuning_micro_batch_sizes + + def fp16_enabled(self): + if FP16 in self.user_config.keys(): + return self.user_config[FP16].get(FP16_ENABLED, FP16_ENABLED_DEFAULT) + else: + return False + + def get_gpu_memory_info(self): + return get_accelerator().total_memory() + + def get_activation_memory_per_gpu(self): + if self.model_info and "activation_mem_per_gpu" in self.model_info: + return self.model_info["activation_mem_per_gpu"] + + def get_instantiation_memory_required_per_gpu(self, zero_stage): + num_params = self.get_model_num_params() + total_gpus = self.exp_num_nodes * self.exp_num_gpus + fp16_enabled = self.fp16_enabled() + + if not num_params: + return 0 + # assume the model uses Adam optimizer + # ZeroStageEnum.disabled: + params_mem = num_params * (2 if fp16_enabled else 4) + gradients_mem = num_params * (2 if fp16_enabled else 4) + optimizer_mem = num_params * (16 if fp16_enabled else 8) + + if zero_stage >= ZeroStageEnum.optimizer_states: + optimizer_mem = optimizer_mem / total_gpus + + if zero_stage >= ZeroStageEnum.gradients: + gradients_mem = gradients_mem / total_gpus + + if zero_stage >= ZeroStageEnum.weights: + params_mem = params_mem / total_gpus + + mem_per_gpu = (params_mem + gradients_mem + optimizer_mem) / self.mp_size() + + return mem_per_gpu + + def _generate_experiments(self, tuning_space, max_train_batch_size_per_gpu): + """Generates a list of autotuning experiments given a tuning_space. + The corresponding parameter values are replaced by user-defined values in the DeepSpeed configuration file. + Args: + tuning_space ([dict]): A DeepSpeed configuration dictionary where a value can be a list (called a tuning parameter). For example, + { + "zero_optimization": { + "stage": 1, + "reduce_bucket_size": [5e7, + 5e8, + 1e9], + "allgather_bucket_size": [5e7, + 5e8, + 1e9], + } + } + reduce_bucket_size and allgather_bucket_size are the tuning parameters in this tuning space. + Returns: + [list]: a list of experiments generated by taking combinations of values of the tuning space. The above tuning space generates 3*3 = 9 experiments if the user DeepSpeed configuration file does not overwrite the two tuning parameters or define more tuning parameters. + """ + exps = [] + + # each zero stage uses a different template configuration file + config_zero = tuning_space.get(ZERO_OPTIMIZATION, {}) + stage = config_zero.get(ZERO_OPTIMIZATION_STAGE, ZERO_OPTIMIZATION_STAGE_DEFAULT) + template_config = {} + if stage == 0: + template_path = DEFAULT_TEMPLATE_PATH_ZERO_0 + template_config = hjson.load(open(template_path, 'r')) + prefix = "z0_" + + elif stage == 1: + template_path = DEFAULT_TEMPLATE_PATH_ZERO_1 + template_config = hjson.load(open(template_path, 'r')) + prefix = "z1_" + + elif stage == 2: + template_path = DEFAULT_TEMPLATE_PATH_ZERO_2 + template_config = hjson.load(open(template_path, 'r')) + prefix = "z2_" + + elif stage == 3: + template_path = DEFAULT_TEMPLATE_PATH_ZERO_3 + template_config = hjson.load(open(template_path, 'r')) + model_info = self.model_info + if model_info and "hidden_size" in model_info: + hs = model_info["hidden_size"] + template_config[ZERO_OPTIMIZATION]['reduce_bucket_size'] = hs * hs + template_config[ZERO_OPTIMIZATION]['stage3_prefetch_bucket_size'] = 0.9 * hs * hs + template_config[ZERO_OPTIMIZATION]['stage3_param_persistence_threshold'] = 10 * hs + prefix = "z3_" + else: + return exps + + # replace the corresponding parameter values if the user specifies them in the DeepSpeed configuration file + replace_dict(tuning_space, self.user_config, [ZERO_OPTIMIZATION, TRAIN_MICRO_BATCH_SIZE_PER_GPU]) + + logger.debug(f"tuning_space = {json.dumps(tuning_space)}") + + all_configs = get_all_configs(tuning_space, ignore_keys=["optimizer"]) + + tuning_keys = get_tuning_keys(tuning_space) + + logger.debug(f"tuning_keys = {tuning_keys}") + + logger.debug(f"before pruning total configs = {len(all_configs)}") + + pruned_list = prune_configs(all_configs) + + logger.debug(f"after pruning total configs = {len(pruned_list)}") + + for config in pruned_list: + exp_config = copy.deepcopy(template_config) + # fill the template with the expr config + replace_dict(exp_config, config) + + # if the config does not use offloading, remove the offloading section + config_zero = config.get(ZERO_OPTIMIZATION, None) + if config_zero: + if OFFLOAD_OPTIMIZER not in config_zero and OFFLOAD_OPTIMIZER in exp_config[ZERO_OPTIMIZATION]: + del exp_config[ZERO_OPTIMIZATION][OFFLOAD_OPTIMIZER] + if OFFLOAD_PARAM not in config_zero and OFFLOAD_PARAM in exp_config[ZERO_OPTIMIZATION]: + del exp_config[ZERO_OPTIMIZATION][OFFLOAD_PARAM] + # set gradient accumulation steps according to max_train_batch_size_per_gpu + mbs = exp_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] + gas = max_train_batch_size_per_gpu // mbs + exp_config[GRADIENT_ACCUMULATION_STEPS] = gas + exp_config[TRAIN_BATCH_SIZE] = mbs * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp = {} + # generate the expr name + exp_name = canonical_name(exp_config, tuning_keys, prefix) + exp['name'] = exp_name + exp[DS_CONFIG] = exp_config + exp['num_gpus'] = self.exp_num_gpus + exp['num_nodes'] = self.exp_num_nodes + exps.append(exp) + + return exps + + def tune(self): + """ Tunes Zero stages, micro batch size per GPU, and other Zero configurations. Performance metrics of different tuning spaces are recorded in self.records. + """ + if has_mlflow: + self.mlflow_parent_id = os.environ['MLFLOW_RUN_ID'] + mlflow.start_run(run_id=self.mlflow_parent_id) + + self.start_time = time.time() + if self.fast_enabled(): + logger.info(f"Fast mode is enabled. Tuning micro batch size only.") + + # model info profile run with DEFAULT_MIN_MEM_CONFIG + model_info = self.model_info_profile_run() + if model_info: + self.model_info = model_info + else: + return + + logger.info(f"The model has {number_to_string(self.get_model_num_params())} parameters.") + + self.gpu_mem = self.get_gpu_memory_info() + logger.info(f"Memory per GPU in the system is {memory_to_string(self.gpu_mem, postfix='B')}.") + + self.activation_mem = self.get_activation_memory_per_gpu() + logger.info( + f"The model requires at least {memory_to_string(self.activation_mem, postfix='B')} activation memory for micro batch size 1." + ) + + stage = self.user_config.get(ZERO_OPTIMIZATION, {}).get(ZERO_OPTIMIZATION_STAGE, 0) + + user_zero_stages = [stage] if not isinstance(stage, list) else stage + logger.info(f"User-defined zero stages are {stage}.") + + mbs = 0 + max_mbs = 0 + metric_val = 0 + + required_gpu_mem = self.get_instantiation_memory_required_per_gpu(ZeroStageEnum.disabled) + self.activation_mem + if self.gpu_mem > required_gpu_mem: + if "all" in user_zero_stages or ZeroStageEnum.disabled in user_zero_stages: + logger.info( + f"The model might be runable with ZERO 0 (which requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory with mbs = 1), adding DEFAULT_TUNING_SPACE_ZERO_0 to the global tuning space" + ) + next_max_mbs, next_mbs, next_metric_val = self.tune_space(DEFAULT_TUNING_SPACE_ZERO_0) + if next_mbs > mbs: + mbs = next_mbs + max_mbs = next_max_mbs + metric_val = next_metric_val + if has_mlflow: + mlflow.log_metric(f"z0{self.metric()}", next_metric_val) + else: + logger.info( + f"The model is not runable with ZERO stage {ZeroStageEnum.disabled} (which requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory with mbs = 1)" + ) + + required_gpu_mem = self.get_instantiation_memory_required_per_gpu( + ZeroStageEnum.optimizer_states) + self.activation_mem + if self.gpu_mem > required_gpu_mem: + if "all" in user_zero_stages or ZeroStageEnum.optimizer_states in user_zero_stages: + logger.info( + f"The model might be runable with ZERO 1 (which requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory), adding DEFAULT_TUNING_SPACE_ZERO_1 to the global tuning space" + ) + next_max_mbs, next_mbs, next_metric_val = self.tune_space(DEFAULT_TUNING_SPACE_ZERO_1, + prev_max_mbs=max_mbs, + prev_best_mbs=mbs, + prev_best_metric_val=metric_val) + if next_mbs > mbs: + mbs = next_mbs + max_mbs = next_max_mbs + metric_val = next_metric_val + if has_mlflow: + mlflow.log_metric(f"z1{self.metric()}", next_metric_val) + else: + logger.info( + f"The model is not runable with ZERO stage {ZeroStageEnum.optimizer_states} (which requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory with mbs = 1)" + ) + + required_gpu_mem = self.get_instantiation_memory_required_per_gpu( + ZeroStageEnum.gradients) + self.activation_mem + if self.gpu_mem > required_gpu_mem: + if "all" in user_zero_stages or ZeroStageEnum.gradients in user_zero_stages: + logger.info( + f"The model might be runable with ZERO 2 (which requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory), adding DEFAULT_TUNING_SPACE_ZERO_2 to the global tuning space" + ) + next_max_mbs, next_mbs, next_metric_val = self.tune_space(DEFAULT_TUNING_SPACE_ZERO_2, + prev_max_mbs=max_mbs, + prev_best_mbs=mbs, + prev_best_metric_val=metric_val) + if next_mbs > mbs: + mbs = next_mbs + max_mbs = next_max_mbs + metric_val = next_metric_val + if has_mlflow: + mlflow.log_metric(f"z2{self.metric()}", next_metric_val) + else: + logger.info( + f"The model is not runable with ZERO stage {ZeroStageEnum.gradients} (which requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory with mbs = 1)" + ) + + required_gpu_mem = self.get_instantiation_memory_required_per_gpu(ZeroStageEnum.weights) + self.activation_mem + if self.gpu_mem > required_gpu_mem: + if "all" in user_zero_stages or ZeroStageEnum.weights in user_zero_stages: + logger.info( + f"The model might be runable with ZERO 3 (which requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory), adding DEFAULT_TUNING_SPACE_ZERO_3 to the global tuning space" + ) + _, _, next_metric_val = self.tune_space(DEFAULT_TUNING_SPACE_ZERO_3, + prev_max_mbs=max_mbs, + prev_best_mbs=mbs, + prev_best_metric_val=metric_val) + if has_mlflow: + mlflow.log_metric(f"z3{self.metric()}", next_metric_val) + else: + logger.info( + f"The model has {self.get_model_num_params()} parameters and requires at least {memory_to_string(required_gpu_mem, postfix='B')} memory per GPU with DeepSpeed Zero stage {ZeroStageEnum.weights} optimization. Memory per GPU in system is {memory_to_string(self.gpu_mem)}. No tuning is performed." + ) + return + if has_mlflow: + mlflow.end_run() + + def tune_space(self, tuning_space, prev_max_mbs=0, prev_best_mbs=0, prev_best_metric_val=0): + config_zero = tuning_space.get(ZERO_OPTIMIZATION, {}) + stage = config_zero.get(ZERO_OPTIMIZATION_STAGE, None) + tuning_space_name = TUNING_MICRO_BATCH_SIZE_PREFIX + str(stage) + tuning_micro_batch_sizes = [] + max_train_batch_size_per_gpu = 0 + tuning_micro_batch_sizes_overwritten = False + + # calculate max micro batch size using gpu memory, model instantiation memory and activation memory + # calculated_max_micro_batch_size = (memory_per_gpu - instantiation_memory) // activation_memory_micro_batch_size_1 + calculated_max_micro_batch_size = int( + self.gpu_mem - self.get_instantiation_memory_required_per_gpu(stage)) // self.activation_mem + logger.info( + f"Start tuning for space {tuning_space_name}, calculated_max_micro_batch_size = {calculated_max_micro_batch_size}" + ) + + if calculated_max_micro_batch_size < prev_max_mbs: + logger.info(f"No need to tune Zero stage {stage}. End tuning for space {tuning_space_name}") + return 0, 0, 0 + + if TRAIN_MICRO_BATCH_SIZE_PER_GPU in self.user_config and isinstance( + self.user_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU], list): + # user-specified micro batch size per gpu is a list which overwrites the default tuning behavior + tuning_micro_batch_sizes = [ + s for s in self.user_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] if isinstance(s, int) + ] + gas = self.get_gas_from_user_config() + min_micro_batch_size = min(tuning_micro_batch_sizes) + max_micro_batch_size = max(tuning_micro_batch_sizes) + max_train_batch_size_per_gpu = max_micro_batch_size * gas + tuning_micro_batch_sizes_overwritten = True + else: + # auto-detects the list of micro batch sizes to tune + min_micro_batch_size, max_micro_batch_size = self.get_min_max_micro_batch_size( + stage, prev_max_mbs, calculated_max_micro_batch_size) + + if max_micro_batch_size < prev_max_mbs: + logger.info(f"No need to tune Zero stage {stage}. End tuning for space {tuning_space_name}") + return 0, 0, 0 + + tuning_micro_batch_sizes, max_train_batch_size_per_gpu = self.get_tuning_micro_batch_size_list( + min_micro_batch_size, + max_micro_batch_size, + num_tuning_micro_batch_sizes=self.num_tuning_micro_batch_sizes()) + + logger.info( + f"tuning_micro_batch_sizes = {tuning_micro_batch_sizes}, max_train_batch_size_per_gpu = {max_train_batch_size_per_gpu}" + ) + + # return if the tuning_micro_batch_sizes list is empty + if not tuning_micro_batch_sizes: + logger.info(f"End tuning for space {tuning_space_name}") + return 0, 0, 0 + + # tune micro batch sizes and gradient accumulation steps given max_train_batch_size_per_gpu + tuning_micro_batch_sizes = self.run_tuning_micro_batch_sizes(tuning_micro_batch_sizes, + max_train_batch_size_per_gpu, + min_micro_batch_size, stage, + tuning_micro_batch_sizes_overwritten) + + fast_best_record = self.get_best_space_record(tuning_space_name) + fast_best_metric_val = fast_best_record[1] if fast_best_record else 0 + fast_best_mbs = fast_best_record[0][DS_CONFIG][TRAIN_MICRO_BATCH_SIZE_PER_GPU] if fast_best_record else 0 + logger.info(f"fast_best_mbs = {fast_best_mbs}, name = {fast_best_record[0]['name']}") + + if self.fast_enabled() or stage == 0: + logger.info(f"End tuning for space: {tuning_space_name}") + return max_micro_batch_size, fast_best_mbs, fast_best_metric_val + + # if the best metric or the micro batch size for that best metric in the current Zero stage after tuning micro batch size is less than the corresponding value in the previous Zero stage, return, do not tune other Zero configuration parameters + if stage > 0: + if fast_best_mbs <= prev_best_mbs or fast_best_metric_val < prev_best_metric_val: + logger.info( + f"End tuning for space: {tuning_space_name}. No need to tune other Zero configuration parameters.") + return max_micro_batch_size, fast_best_mbs, fast_best_metric_val + + tuning_space[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = tuning_micro_batch_sizes + tuning_space_name = canonical_name(tuning_space, + tuning_keys=get_tuning_keys(tuning_space), + prefix="z" + str(stage) + "_", + omit_val=True) + + logger.info(f'Tuning space is {tuning_space}') + logger.info(f'Tuning space name is {tuning_space_name}') + + exps = self._generate_experiments(tuning_space, max_train_batch_size_per_gpu) + + logger.info(f'Tuner type is {self.autotuning_config.tuner_type}') + if self.autotuning_config.tuner_type == AUTOTUNING_TUNER_MODELBASED: + t = ModelBasedTuner(exps, self.rm, self.metric(), tuning_space) + elif self.autotuning_config.tuner_type == AUTOTUNING_TUNER_RANDOM: + t = RandomTuner(exps, self.rm, self.metric()) + else: + t = GridSearchTuner(exps, self.rm, self.metric()) + + sample_size = len(self.rm.nodes) * self.rm.num_gpus_per_node // (self.exp_num_gpus * self.exp_num_nodes) + num_exps = t.tune(sample_size=sample_size, + n_trials=self.autotuning_config.tuner_num_trials, + early_stopping=self.autotuning_config.tuner_early_stopping) + exp = t.best_exp + metric_val = t.best_metric_val + if exp: + self.update_records(tuning_space_name, exp, metric_val, num_exps) + + full_best_record = self.get_best_space_record(tuning_space_name) + full_best_metric_val = full_best_record[1] if full_best_record else -1 + + if full_best_metric_val > fast_best_metric_val: + best_metric_val = full_best_metric_val + best_mbs = full_best_record[0][DS_CONFIG][TRAIN_MICRO_BATCH_SIZE_PER_GPU] if full_best_record else -1 + else: + best_metric_val = fast_best_metric_val + best_mbs = fast_best_mbs + + logger.info(f"End tuning for space: {tuning_space_name}") + return max_micro_batch_size, best_mbs, best_metric_val + + def get_plateau_mbs(self, tuning_space_name): + if tuning_space_name not in self.records: + return 0 + space_records = self.records[tuning_space_name] + sorted_space_records = sorted(space_records, key=lambda x: x[0][DS_CONFIG][TRAIN_MICRO_BATCH_SIZE_PER_GPU]) + prev_metric_val = None + prev_micro_batch_size = 0 + for (exp, metric_val, _) in sorted_space_records: + if prev_metric_val: + if metric_val < prev_metric_val: + break + if (metric_val >= prev_metric_val + and (metric_val - prev_metric_val) / prev_metric_val < METRIC_PERCENT_DIFF_CONST): + break + prev_metric_val = metric_val + prev_micro_batch_size = exp[DS_CONFIG][TRAIN_MICRO_BATCH_SIZE_PER_GPU] + plateau_mbs = prev_micro_batch_size + return plateau_mbs + + def get_model_num_params(self): + if self.model_info and "num_params" in self.model_info: + return self.model_info["num_params"] + + def model_info_profile_run(self): + """Does a model information profiling experiment that collects the number of model parameters and activation memory.\ + The experiment produces a "profile_model_info" folder under self.results_dir. + Returns: + [dict]: a model information dictionary, e.g., {"num_params": 335144976, "trainable_num_params": 335144976, "activation_mem_per_gpu": 324358144, "rank": 0} + """ + logger.info("Starting model info profile run.") + model_info = self.autotuning_config.model_info + if model_info and MODEL_INFO_NUM_PARAMS in model_info: + return model_info + + ds_config = copy.deepcopy(self.user_config) + replace_dict(ds_config, DEFAULT_MIN_MEM_CONFIG) + + model_info_path = os.path.join(self.results_dir, "profile_model_info", "model_info.json") + ds_config[AUTOTUNING] = {"enabled": True, "model_info_path": model_info_path, "model_info": {"profile": True}} + + exp_config = {} + exp_name = "profile_model_info" + exp_config['name'] = exp_name + exp_config[DS_CONFIG] = ds_config + exp_config['num_gpus'] = self.exp_num_gpus + exp_config['num_nodes'] = self.exp_num_nodes + exp_config['hostfile'] = self.args.hostfile + exp_path = os.path.join(self.exps_dir, f'{exp_name}.json') + + with open(exp_path, 'w', buffering=BUFSIZE) as fd: + json.dump(exp_config, fd) + fd.flush() + os.fsync(fd) + + self.rm.schedule_experiments([exp_path]) + self.rm.run() + + for exp_id, (exp_json, err) in self.rm.finished_experiments.items(): + self.rm.clear() + if err: + logger.error(f"The model is not runnable with DeepSpeed with error = {err}") + return None + + if os.path.exists(model_info_path): + with open(model_info_path, 'r') as f: + model_info = hjson.load(f) + return model_info + + def update_records(self, space_name, exp, metric_val, num_exps): + if space_name not in self.records: + self.records[space_name] = [(exp, metric_val, num_exps)] + else: + self.records[space_name].append((exp, metric_val, num_exps)) + + def get_best_space_record(self, space_name): + if space_name not in self.records: + return None + space_records = self.records[space_name] + best_space_record = None + space_num_exps = 0 + for (exp, metric_val, num_exps) in space_records: + space_num_exps += num_exps + if best_space_record is None or metric_val > best_space_record[1]: + best_space_record = (exp, metric_val) + if best_space_record: + best_space_record = best_space_record + (space_num_exps, ) + return best_space_record + + def get_best_space_records(self): + best_space_records = {} + global_best_record = None + for space_name, space_records in self.records.items(): + best_space_record = self.get_best_space_record(space_name) + if best_space_record: + best_space_records[space_name] = best_space_record + if not global_best_record or best_space_record[1] > global_best_record[1]: + global_best_record = best_space_record + if global_best_record: + best_space_records[GLOBAL_TUNING_SPACE] = global_best_record + return best_space_records + + def run_tuning_micro_batch_sizes(self, tuning_micro_batch_sizes, max_train_batch_size_per_gpu, + min_micro_batch_size, stage, tuning_micro_batch_sizes_overwritten): + assert tuning_micro_batch_sizes, "the tuning micro batch size list is empty" + tuning_micro_batch_sizes.sort() + max_micro_batch_size = tuning_micro_batch_sizes[-1] + max_micro_batch_size_metric_val = 0 + + ds_config = get_first_config(self.user_config) + ds_config[ZERO_OPTIMIZATION] = {ZERO_OPTIMIZATION_STAGE: stage} + tuning_space_name = TUNING_MICRO_BATCH_SIZE_PREFIX + str(stage) + + exp_paths = [] + for mbs in tuning_micro_batch_sizes: + ds_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = mbs + gas = max_train_batch_size_per_gpu // mbs + ds_config[GRADIENT_ACCUMULATION_STEPS] = gas + ds_config[TRAIN_BATCH_SIZE] = mbs * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp_name = tuning_space_name + "_gas" + str(gas) + "_tmbspg" + str(mbs) + exp_config = {} + exp_config['name'] = exp_name + exp_config[DS_CONFIG] = ds_config + exp_config['num_gpus'] = self.exp_num_gpus + exp_config['num_nodes'] = self.exp_num_nodes + exp_config['hostfile'] = self.args.hostfile + exp_path = os.path.join(self.exps_dir, f'{exp_name}.json') + + with open(exp_path, 'w', buffering=BUFSIZE) as fd: + json.dump(exp_config, fd) + fd.flush() + os.fsync(fd) + exp_paths.append(exp_path) + + self.rm.schedule_experiments(exp_paths) + self.rm.run() + + for exp_id, (exp, err) in self.rm.finished_experiments.items(): + if exp: + metric_file = exp[DS_CONFIG][AUTOTUNING][AUTOTUNING_METRIC_PATH] + if os.path.exists(metric_file): + + with open(metric_file, 'r') as f: + results = hjson.load(f) + metric_val = results[self.metric()] + self.update_records(tuning_space_name, exp, metric_val, 1) + if max_micro_batch_size == exp[DS_CONFIG][TRAIN_MICRO_BATCH_SIZE_PER_GPU]: + max_micro_batch_size_metric_val = metric_val + if has_mlflow: + os.environ.pop('MLFLOW_RUN_ID') + mlflow.start_run(nested=True, run_name=exp['name']) + for metric in results: + mlflow.log_metric(metric, results[metric]) + mlflow.end_run() + os.environ['MLFLOW_RUN_ID'] = self.mlflow_parent_id + else: + self.update_records(tuning_space_name, exp, 0, 1) + else: + mbs = exp[DS_CONFIG][TRAIN_MICRO_BATCH_SIZE_PER_GPU] + logger.info(f"micro batch size = {mbs} was not run successfully") + + self.rm.clear() + + if tuning_micro_batch_sizes_overwritten: + return tuning_micro_batch_sizes + + # in a auto-detected tuning_micro_batch_sizes list, max_micro_batch_size might not be performant as the memory consumption is close to max + # try smaller values while gas stays the same + # if finding a more performant mbs value, use it to replace max_micro_batch_size in the list + min_micro_batch_size_with_same_gas = (tuning_micro_batch_sizes[-2] + + 1) if len(tuning_micro_batch_sizes) > 1 else min_micro_batch_size + + prev_best_metric_val = max_micro_batch_size_metric_val + prev_best_mbs = max_micro_batch_size + + stride = (max_micro_batch_size - min_micro_batch_size_with_same_gas) // 3 + if stride == 0: + stride = 1 + for mbs in reversed(range(min_micro_batch_size_with_same_gas, max_micro_batch_size, stride)): + ds_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = mbs + gas = max_train_batch_size_per_gpu // mbs + ds_config[GRADIENT_ACCUMULATION_STEPS] = gas + ds_config[TRAIN_BATCH_SIZE] = mbs * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp_name = tuning_space_name + "_gas" + str(gas) + "_tmbspg" + str(mbs) + exp, metric_val = self.run_ds_config(ds_config, exp_name) + + if metric_val: + with open(metric_file, 'r') as f: + results = hjson.load(f) + metric_val = results[self.metric()] + if has_mlflow: + os.environ.pop('MLFLOW_RUN_ID') + mlflow.start_run(nested=True, run_name=exp_name) + for metric in results: + mlflow.log_metric(metric, results[metric]) + mlflow.end_run() + os.environ['MLFLOW_RUN_ID'] = self.mlflow_parent_id + self.update_records(tuning_space_name, exp, metric_val, 1) + if metric_val > prev_best_metric_val * (1 + METRIC_PERCENT_DIFF_CONST): + prev_best_metric_val = metric_val + prev_best_mbs = mbs + else: + break + else: + self.update_records(tuning_space_name, exp, 0, 1) + break + if prev_best_mbs != max_micro_batch_size: + tuning_micro_batch_sizes[-1] = prev_best_mbs + return tuning_micro_batch_sizes + + def get_min_max_micro_batch_size(self, stage, min_micro_batch_size, calculated_max_micro_batch_size): + # get min and max micro batch size with gradient accumulation steps = 1 + if min_micro_batch_size > calculated_max_micro_batch_size: + return -1, -1 + + used_micro_batch_sizes = [] + tuning_space_name = TUNING_MICRO_BATCH_SIZE_PREFIX + str(stage) + + ds_config = get_first_config(self.user_config) + ds_config[ZERO_OPTIMIZATION] = {ZERO_OPTIMIZATION_STAGE: stage} + gas = self.get_gas_from_user_config() + ds_config[GRADIENT_ACCUMULATION_STEPS] = gas + + # search for the min micro batch size + if min_micro_batch_size < 1: + if TRAIN_MICRO_BATCH_SIZE_PER_GPU in self.user_config and isinstance( + self.user_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU], int): + # user specifies train_micro_batch_size_per_gpu as an int + mbs = int(self.user_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU]) + else: + # user does not specify train_micro_batch_size_per_gpu or sets it to "auto" when using Hugging Face + val = self.get_val_from_user_args(TRAIN_MICRO_BATCH_SIZE_PER_GPU) + if val: + mbs = int(val) + else: + mbs = 1 + assert mbs > 0, "The micro batch size per GPU must be greater than 0." + ds_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = mbs + ds_config[GRADIENT_ACCUMULATION_STEPS] = gas + ds_config[TRAIN_BATCH_SIZE] = mbs * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp_name = tuning_space_name + "_gas" + str(gas) + "_tmbspg" + str(mbs) + exp, metric_val = self.run_ds_config(ds_config, exp_name) + if metric_val: + self.update_records(tuning_space_name, exp, metric_val, 1) + used_micro_batch_sizes.append(mbs) + min_micro_batch_size = mbs + else: + self.update_records(tuning_space_name, exp, 0, 1) + logger.info(f"User-specified micro batch size per GPU {mbs} does not run") + if self.min_train_micro_batch_size_per_gpu() == mbs: + return -1, -1 + mbs = self.min_train_micro_batch_size_per_gpu() + ds_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = mbs + ds_config[GRADIENT_ACCUMULATION_STEPS] = gas + ds_config[TRAIN_BATCH_SIZE] = mbs * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp_name = tuning_space_name + "_gas" + str(gas) + "_tmbspg" + str(mbs) + exp, metric_val = self.run_ds_config(ds_config, exp_name) + if not metric_val: + self.update_records(tuning_space_name, exp, 0, 1) + logger.info(f"min_train_micro_batch_size_per_gpu {mbs} is not runnable.") + return -1, -1 + self.update_records(tuning_space_name, exp, metric_val, 1) + min_micro_batch_size = mbs + used_micro_batch_sizes.append(mbs) + else: + ds_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = min_micro_batch_size + ds_config[GRADIENT_ACCUMULATION_STEPS] = gas + ds_config[TRAIN_BATCH_SIZE] = min_micro_batch_size * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp_name = tuning_space_name + "_gas" + str(gas) + "_tmbspg" + str(min_micro_batch_size) + exp, metric_val = self.run_ds_config(ds_config, exp_name) + if metric_val: + self.update_records(tuning_space_name, exp, metric_val, 1) + used_micro_batch_sizes.append(min_micro_batch_size) + else: + self.update_records(tuning_space_name, exp, 0, 1) + return -1, -1 + + # search for the max micro batch size + max_micro_batch_size = min(calculated_max_micro_batch_size, self.max_train_micro_batch_size_per_gpu()) + for mbs in [math.ceil(1.05 * max_micro_batch_size), max_micro_batch_size, int(0.95 * max_micro_batch_size)]: + if mbs > self.max_train_micro_batch_size_per_gpu(): + continue + if mbs in used_micro_batch_sizes: + return min_micro_batch_size, mbs + ds_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = mbs + ds_config[TRAIN_BATCH_SIZE] = mbs * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp_name = tuning_space_name + "_gas" + str(gas) + "_tmbspg" + str(mbs) + exp, metric_val = self.run_ds_config(ds_config, exp_name) + + if metric_val: + logger.info(f"mbs = {mbs} is found as max mbs") + self.update_records(tuning_space_name, exp, metric_val, 1) + used_micro_batch_sizes.append(mbs) + return min_micro_batch_size, mbs + else: + self.update_records(tuning_space_name, exp, 0, 1) + + space_records = self.records[tuning_space_name] if tuning_space_name in self.records else [] + if space_records: + prev_idx = min(range(len(space_records)), + key=lambda i: abs(space_records[i][0][DS_CONFIG][TRAIN_MICRO_BATCH_SIZE_PER_GPU] - + min_micro_batch_size)) + prev_metric_val = space_records[prev_idx][1] + else: + prev_metric_val = None + + low = min_micro_batch_size + high = max_micro_batch_size + # binary search until low is the smallest micro batch size that OOMs. + while low <= high: + mid = int((low + high) // 2) + logger.debug(f"trying mbs = {mid}, low = {low}, high = {high}") + if mid not in used_micro_batch_sizes: + ds_config[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = mid + ds_config[TRAIN_BATCH_SIZE] = mid * gas * \ + self.exp_num_gpus * self.exp_num_nodes // self.mp_size() + exp_name = tuning_space_name + "_gas" + str(gas) + "_tmbspg" + str(mid) + exp, metric_val = self.run_ds_config(ds_config, exp_name) + if metric_val: + low = mid + 1 + self.update_records(tuning_space_name, exp, metric_val, 1) + used_micro_batch_sizes.append(mid) + if prev_metric_val and ((metric_val - prev_metric_val) / + prev_metric_val) < METRIC_PERCENT_DIFF_CONST: + logger.info(f"performance plateaus at mbs = {low}") + break + prev_metric_val = metric_val + else: + self.update_records(tuning_space_name, exp, 0, 1) + high = mid - 1 + else: + low = mid + 1 + max_micro_batch_size = low - 1 + + logger.info(f"min_micro_batch_size = {min_micro_batch_size}, max_micro_batch_size = {max_micro_batch_size}.") + + return min_micro_batch_size, max_micro_batch_size + + def get_gas_from_user_config(self): + gas = 1 + if GRADIENT_ACCUMULATION_STEPS in self.user_config: + gas_in_config = self.user_config[GRADIENT_ACCUMULATION_STEPS] + if isinstance(gas_in_config, int): + gas = gas_in_config + elif gas_in_config == "auto": # GRADIENT_ACCUMULATION_STEPS: "auto" + val = self.get_val_from_user_args(GRADIENT_ACCUMULATION_STEPS) + if val: + gas = int(val) + elif isinstance(gas_in_config, list): + logger.info( + f"Specifying a list of {GRADIENT_ACCUMULATION_STEPS} to tune is not supported. 1 would be used.") + assert gas > 0, "Gradient accumulation steps must be positive." + return gas + + def get_val_from_user_args(self, ds_name): + arg_mappings = self.autotuning_config.arg_mappings + user_args = self.args.user_args + if arg_mappings and ds_name in arg_mappings: + arg_name = arg_mappings[ds_name] + if arg_name in user_args: + idx = user_args.index(arg_name) + if user_args[idx + 1].isnumeric(): + return (user_args[idx + 1]) + return None + + def get_tuning_micro_batch_size_list(self, min_micro_batch_size, max_micro_batch_size, + num_tuning_micro_batch_sizes): + """Get a list of micro batch sizes to tune based on min and max values, as well as the size of the list. + Args: + min_micro_batch_size ([int]): min micro batch size per GPU + max_micro_batch_size ([int]): max micro batch size per GPU + num_tuning_micro_batch_sizes (int): the number of items in the returned list + + Returns: + [list]: a list of micro batch sizes to tune. + """ + if min_micro_batch_size <= 0 or max_micro_batch_size <= 0: + logger.info( + f"min_micro_batch_size = {min_micro_batch_size}, max_micro_batch_size = {max_micro_batch_size}") + return [], 0 + + # NUM_GPUS=$(( ${NUM_WORKERS} * ${NUM_GPUS_PER_WORKER} )) + # DP_SIZE=$(( ${NUM_GPUS} / (${PP_SIZE} * ${MP_SIZE}) )) + # GRAD_ACC_STEPS=$(( ${TARGET_GLOBAL_BATCH_SIZE} / (${BATCH_SIZE} * ${DP_SIZE}) )) + if self.max_train_batch_size() and self.max_train_batch_size( + ) > 0: # if the user specifies a max_train_batch_size + max_train_batch_size_per_gpu = self.max_train_batch_size() * self.mp_size() // (self.exp_num_gpus * + self.exp_num_nodes) + else: + gas = self.get_gas_from_user_config() + max_train_batch_size_per_gpu = max_micro_batch_size * gas // self.mp_size() + logger.info(f"max_train_batch_size_per_gpu = {max_train_batch_size_per_gpu}") + if min_micro_batch_size < max_micro_batch_size // 2: + min_micro_batch_size = max_micro_batch_size // 2 + + # constant stride + stride = (max_micro_batch_size - min_micro_batch_size) // num_tuning_micro_batch_sizes + if stride == 0: + stride = 1 + ls = [] + min_gas = max_train_batch_size_per_gpu // max_micro_batch_size + # if gas is the same as min_gas, do not add mbs to the tuning list + for mbs in range(min_micro_batch_size, max_micro_batch_size, stride): + if max_train_batch_size_per_gpu // mbs != min_gas: + ls.append(mbs) + ls.append(max_micro_batch_size) + + return ls, max_train_batch_size_per_gpu + + def run_ds_config(self, ds_config, exp_name): + exp_config = {} + exp_config['name'] = exp_name + exp_config[DS_CONFIG] = ds_config + exp_config['num_gpus'] = self.exp_num_gpus + exp_config['num_nodes'] = self.exp_num_nodes + exp_config['hostfile'] = self.args.hostfile + exp_path = os.path.join(self.exps_dir, f'{exp_name}.json') + + logger.debug(f'run_ds_config exp_name = {exp_name}') + + with open(exp_path, 'w', buffering=BUFSIZE) as fd: + json.dump(exp_config, fd) + fd.flush() + os.fsync(fd) + self.rm.schedule_experiments([exp_path]) + self.rm.run() + exp, metric_val = self.rm.parse_results(self.metric()) + self.rm.clear() + return exp, metric_val + + def write_optimal_config(self): + best_space_records = self.get_best_space_records() + if GLOBAL_TUNING_SPACE not in best_space_records: + return + best_exp, best_metric_val, _ = best_space_records[GLOBAL_TUNING_SPACE] + if best_exp: + exp_dir = best_exp["result_dir"] + cmd = None + with open(os.path.join(exp_dir, "cmd.txt"), "r") as f: + cmd = [str(i) for i in f.read().split()] + + ds_config = hjson.load(open(os.path.join(exp_dir, "ds_config.json"), "r")) + ds_config.pop(AUTOTUNING) + + ds_config_path = os.path.join(self.results_dir, "ds_config_optimal.json") + json.dump(ds_config, open(ds_config_path, "w")) + + cmd_path = os.path.join(self.results_dir, "cmd_optimal.txt") + with open(cmd_path, "w") as fd: + fd.write(" ".join(cmd)) + fd.write("\n") + fd.flush() + self.optimal_cmd = cmd + self.optimal_ds_config = ds_config + logger.info( + f"Wrote the optimal DeepSpeed configuration found by autotuning to {ds_config_path}, and the corresponding DeepSpeed command to {cmd_path}" + ) + + def run_after_tuning(self): + """ Launches the training with the optimal DeepSpeed configuration found through the autotuning process. + "ds_config_optimal.json" describing the optimal DeepSpeed configuration as well the command used to launch training "cmd_optimal.txt" are saved to self.results_dir. + """ + if self.optimal_cmd: + result = subprocess.Popen(self.optimal_cmd) + result.wait() + + logger.info(f"Done running with the optimal DeepSpeed configuration using {self.optimal_cmd}") + else: + logger.info(f"No optimal DeepSpeed configuration found by autotuning.") diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/config.py b/lib/python3.12/site-packages/deepspeed/autotuning/config.py new file mode 100644 index 0000000000000000000000000000000000000000..6f58fb4e42965be30820809a1c8d6f424b209bd8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/config.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param, get_dict_param, DeepSpeedConfigObject +from deepspeed.autotuning.constants import * + + +class DeepSpeedAutotuningConfig(DeepSpeedConfigObject): + + def __init__(self, param_dict): + super(DeepSpeedAutotuningConfig, self).__init__() + + self.enabled = None + self.start_step = None + self.end_step = None + self.metric_path = None + self.arg_mappings = None + self.metric = None + self.model_info = None + self.results_dir = None + self.exps_dir = None + self.overwrite = None + + if param_dict and AUTOTUNING in param_dict.keys(): + autotuning_dict = param_dict[AUTOTUNING] + else: + autotuning_dict = {} + + self._initialize(autotuning_dict) + + def _initialize(self, autotuning_dict): + self.enabled = get_scalar_param(autotuning_dict, AUTOTUNING_ENABLED, AUTOTUNING_ENABLED_DEFAULT) + + self.fast = get_scalar_param(autotuning_dict, AUTOTUNING_FAST, AUTOTUNING_FAST_DEFAULT) + + self.results_dir = get_scalar_param(autotuning_dict, AUTOTUNING_RESULTS_DIR, AUTOTUNING_RESULTS_DIR_DEFAULT) + assert self.results_dir, "results_dir cannot be empty" + self.exps_dir = get_scalar_param(autotuning_dict, AUTOTUNING_EXPS_DIR, AUTOTUNING_EXPS_DIR_DEFAULT) + assert self.exps_dir, "exps_dir cannot be empty" + self.overwrite = get_scalar_param(autotuning_dict, AUTOTUNING_OVERWRITE, AUTOTUNING_OVERWRITE_DEFAULT) + + self.start_profile_step = get_scalar_param(autotuning_dict, AUTOTUNING_START_PROFILE_STEP, + AUTOTUNING_START_PROFILE_STEP_DEFAULT) + + self.end_profile_step = get_scalar_param(autotuning_dict, AUTOTUNING_END_PROFILE_STEP, + AUTOTUNING_END_PROFILE_STEP_DEFAULT) + + self.metric = get_scalar_param(autotuning_dict, AUTOTUNING_METRIC, AUTOTUNING_METRIC_DEFAULT) + + self.metric_path = get_scalar_param(autotuning_dict, AUTOTUNING_METRIC_PATH, AUTOTUNING_METRIC_PATH_DEFAULT) + + self.tuner_type = get_scalar_param(autotuning_dict, AUTOTUNING_TUNER_TYPE, AUTOTUNING_TUNER_TYPE_DEFAULT) + + self.tuner_early_stopping = get_scalar_param(autotuning_dict, AUTOTUNING_TUNER_EARLY_STOPPING, + AUTOTUNING_TUNER_EARLY_STOPPING_DEFAULT) + + self.tuner_num_trials = get_scalar_param(autotuning_dict, AUTOTUNING_TUNER_NUM_TRIALS, + AUTOTUNING_TUNER_NUM_TRIALS_DEFAULT) + + self.arg_mappings = get_dict_param(autotuning_dict, AUTOTUNING_ARG_MAPPINGS, AUTOTUNING_ARG_MAPPINGS_DEFAULT) + + self.model_info = get_model_info_config(autotuning_dict) + + self.model_info_path = get_scalar_param(autotuning_dict, AUTOTUNING_MODEL_INFO_PATH, + AUTOTUNING_MODEL_INFO_PATH_DEFAULT) + self.mp_size = get_scalar_param(autotuning_dict, AUTOTUNING_MP_SIZE, AUTOTUNING_MP_SIZE_DEFAULT) + + self.max_train_batch_size = get_dict_param(autotuning_dict, AUTOTUNING_MAX_TRAIN_BATCH_SIZE, + AUTOTUNING_MAX_TRAIN_BATCH_SIZE_DEFAULT) + + self.min_train_batch_size = get_dict_param(autotuning_dict, AUTOTUNING_MIN_TRAIN_BATCH_SIZE, + AUTOTUNING_MIN_TRAIN_BATCH_SIZE_DEFAULT) + + self.max_train_micro_batch_size_per_gpu = get_dict_param( + autotuning_dict, AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU, + AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT) + + self.min_train_micro_batch_size_per_gpu = get_dict_param( + autotuning_dict, AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU, + AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT) + + self.num_tuning_micro_batch_sizes = get_dict_param(autotuning_dict, AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES, + AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES_DEFAULT) + + +def get_model_info_config(param_dict): + if MODEL_INFO in param_dict and param_dict[MODEL_INFO] is not None: + model_info_config = {} + for key, default_value in MODEL_INFO_KEY_DEFAULT_DICT.items(): + model_info_config[key] = get_scalar_param(param_dict[MODEL_INFO], key, default_value) + return model_info_config + return None + + +def get_default_model_info_config(): + return MODEL_INFO_KEY_DEFAULT_DICT diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero0.json b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero0.json new file mode 100644 index 0000000000000000000000000000000000000000..b95c7da0948ef76ae2e593ffacc7b1c153c6671d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero0.json @@ -0,0 +1,5 @@ +{ + "zero_optimization": { + "stage": 0 + } +} diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero1.json b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero1.json new file mode 100644 index 0000000000000000000000000000000000000000..dc90f946f57436a74b540e60907b21283825fe31 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero1.json @@ -0,0 +1,7 @@ +{ + "zero_optimization": { + "stage": 1, + "reduce_bucket_size": 5e8, + "allgather_bucket_size": 5e8 + } +} diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero2.json b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero2.json new file mode 100644 index 0000000000000000000000000000000000000000..46f1817af7eead82a702822c5f2feacdf1a173e1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero2.json @@ -0,0 +1,11 @@ +{ + "zero_optimization": { + "stage": 2, + "allgather_partitions": true, + "allgather_bucket_size": 5e8, + "overlap_comm": false, + "reduce_scatter": true, + "reduce_bucket_size": 5e8, + "contiguous_gradients": false + } +} diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero3.json b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero3.json new file mode 100644 index 0000000000000000000000000000000000000000..620d7eb10e81102a9ec77cfaf3e9c1ab1ccb9a80 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero3.json @@ -0,0 +1,17 @@ +{ + "zero_optimization": { + "stage": 3, + "allgather_partitions": true, + "allgather_bucket_size": 5e8, + "overlap_comm": false, + "reduce_scatter": true, + "reduce_bucket_size": 5e8, + "contiguous_gradients": false, + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_prefetch_bucket_size": 5e8, + "stage3_param_persistence_threshold": 1e6, + "stage3_gather_16bit_weights_on_model_save": false, + "sub_group_size": 1e12 + } +} diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/constants.py b/lib/python3.12/site-packages/deepspeed/autotuning/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..908868a417bb0ab6fc00e9b277adce7201d21ec4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/constants.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######################################### +# autotuner implementation constants +######################################### + +import os + +DEFAULT_TEMPLATE_PATH_ZERO_0 = os.path.join(os.path.dirname(os.path.realpath(__file__)), "config_templates", + "template_zero0.json") +DEFAULT_TEMPLATE_PATH_ZERO_1 = os.path.join(os.path.dirname(os.path.realpath(__file__)), "config_templates", + "template_zero1.json") +DEFAULT_TEMPLATE_PATH_ZERO_2 = os.path.join(os.path.dirname(os.path.realpath(__file__)), "config_templates", + "template_zero2.json") +DEFAULT_TEMPLATE_PATH_ZERO_3 = os.path.join(os.path.dirname(os.path.realpath(__file__)), "config_templates", + "template_zero3.json") + +METRIC_PERCENT_DIFF_CONST = 0.05 +DS_CONFIG = "ds_config" +BUFSIZE = 1 # line buffer size for writing files + +######################################### +# autotuner configuration constants +######################################### +# Autotuner. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +AUTOTUNING_FORMAT = """ +autotuner should be enabled as: +"session_params": { + "autotuning": { + "enabled": true, + "start_step": 5, + "end_step": 15 + } +} +""" + +AUTOTUNING = "autotuning" + +AUTOTUNING_ENABLED = "enabled" +AUTOTUNING_ENABLED_DEFAULT = False + +AUTOTUNING_FAST = "fast" +AUTOTUNING_FAST_DEFAULT = True + +AUTOTUNING_RESULTS_DIR = "results_dir" +AUTOTUNING_RESULTS_DIR_DEFAULT = "autotuning_results" + +AUTOTUNING_EXPS_DIR = "exps_dir" +AUTOTUNING_EXPS_DIR_DEFAULT = "autotuning_exps" + +AUTOTUNING_OVERWRITE = "overwrite" +AUTOTUNING_OVERWRITE_DEFAULT = True + +AUTOTUNING_START_PROFILE_STEP = "start_profile_step" +AUTOTUNING_START_PROFILE_STEP_DEFAULT = 3 + +AUTOTUNING_END_PROFILE_STEP = "end_profile_step" +AUTOTUNING_END_PROFILE_STEP_DEFAULT = 5 +AUTOTUNING_METRIC_PATH = "metric_path" +AUTOTUNING_METRIC_PATH_DEFAULT = None + +AUTOTUNING_TUNER_TYPE = "tuner_type" +AUTOTUNING_TUNER_GRIDSEARCH = "gridsearch" +AUTOTUNING_TUNER_RANDOM = "random" +AUTOTUNING_TUNER_MODELBASED = "model_based" +AUTOTUNING_TUNER_TYPE_DEFAULT = AUTOTUNING_TUNER_GRIDSEARCH +AUTOTUNING_TUNER_EARLY_STOPPING = "tuner_early_stopping" +AUTOTUNING_TUNER_EARLY_STOPPING_DEFAULT = 5 +AUTOTUNING_TUNER_NUM_TRIALS = "tuner_num_trials" +AUTOTUNING_TUNER_NUM_TRIALS_DEFAULT = 50 + +AUTOTUNING_ARG_MAPPINGS = "arg_mappings" +AUTOTUNING_ARG_MAPPINGS_DEFAULT = None + +AUTOTUNING_MAX_TRAIN_BATCH_SIZE = "max_train_batch_size" +AUTOTUNING_MAX_TRAIN_BATCH_SIZE_DEFAULT = None +AUTOTUNING_MIN_TRAIN_BATCH_SIZE = "min_train_batch_size" +AUTOTUNING_MIN_TRAIN_BATCH_SIZE_DEFAULT = 1 +AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU = "max_train_micro_batch_size_per_gpu" +AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT = 1024 +AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU = "min_train_micro_batch_size_per_gpu" +AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT = 1 +AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES = "num_tuning_micro_batch_sizes" +AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES_DEFAULT = 3 + +AUTOTUNING_MP_SIZE = "mp_size" +AUTOTUNING_MP_SIZE_DEFAULT = 1 + +AUTOTUNING_METRIC = "metric" +AUTOTUNING_METRIC_LATENCY = "latency" +AUTOTUNING_METRIC_THROUGHPUT = "throughput" +AUTOTUNING_METRIC_FLOPS = "flops" +AUTOTUNING_METRIC_FORWARD = "forward" +AUTOTUNING_METRIC_BACKWRAD = "flops" +AUTOTUNING_METRIC_STEPS = "step" +AUTOTUNING_METRIC_DEFAULT = AUTOTUNING_METRIC_THROUGHPUT + +######################################### +# MODEL INFO +######################################### +AUTOTUNING_MODEL_INFO_PATH = "model_info_path" +AUTOTUNING_MODEL_INFO_PATH_DEFAULT = None + +MODEL_INFO_FORMAT = ''' +"model_info": { + "num_params": 1000000000, + "hidden_size": 10, + "num_layers": 12, +} +''' +MODEL_INFO = "model_info" +MODEL_INFO_PROFILE = "profile" +MODEL_INFO_PROFILE_DEFAULT = False +MODEL_INFO_NUM_PARAMS = "num_params" +MODEL_INFO_NUM_PARAMS_DEFAULT = None +MODEL_INFO_HIDDEN_SIZE = "hidden_size" +MODEL_INFO_HIDDEN_SIZE_DEFAULT = None +MODEL_INFO_NUM_LAYERS = "num_layers" +MODEL_INFO_NUM_LAYERS_DEFAULT = None + +MODEL_INFO_KEY_DEFAULT_DICT = { + MODEL_INFO_PROFILE: MODEL_INFO_PROFILE_DEFAULT, + MODEL_INFO_NUM_PARAMS: MODEL_INFO_NUM_PARAMS_DEFAULT, + MODEL_INFO_HIDDEN_SIZE: MODEL_INFO_HIDDEN_SIZE_DEFAULT, + MODEL_INFO_NUM_LAYERS: MODEL_INFO_NUM_LAYERS_DEFAULT +} + +######################################### +# autotuner search space constants +######################################### + +DEFAULT_HF_CONFIG = { + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "gradient_accumulation_steps": "auto", +} + +DEFAULT_MIN_MEM_CONFIG = { + "train_micro_batch_size_per_gpu": 1, + "zero_optimization": { + "stage": 3 + }, + "memory_break_down": False +} + +DEFAULT_TUNING_SPACE_ZERO_0 = {"zero_optimization": {"stage": 0}} + +DEFAULT_TUNING_SPACE_ZERO_1 = { + "zero_optimization": { + "stage": 1, + "reduce_bucket_size": [5e7, 5e8, 1e9], + "allgather_bucket_size": [5e7, 5e8, 1e9], + } +} + +DEFAULT_TUNING_SPACE_ZERO_2 = { + "zero_optimization": { + "stage": 2, + "overlap_comm": [True, False], + "reduce_scatter": [False, True], + "reduce_bucket_size": [5e7, 5e8, 1e9], + "allgather_bucket_size": [5e7, 5e8, 1e9], + "contiguous_gradients": [False, True] + }, +} + +DEFAULT_TUNING_SPACE_ZERO_3 = { + "zero_optimization": { + "stage": 3, + "overlap_comm": [True, False], + "reduce_scatter": [False, True], + "reduce_bucket_size": [5e7, 5e8, 1e9], + "allgather_partitions": [True, False], + "allgather_bucket_size": [5e7, 5e8, 1e9], + "contiguous_gradients": [False, True] + }, +} + +GLOBAL_TUNING_SPACE = 'global' +# TUNING_MICRO_BATCH_SIZE_PREFIX="tune_micro_batch_size_z" +TUNING_MICRO_BATCH_SIZE_PREFIX = "z" diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/scheduler.py b/lib/python3.12/site-packages/deepspeed/autotuning/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..35f709ead48a274591ee5aacbe2d269a20539f1a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/scheduler.py @@ -0,0 +1,432 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import copy + +import json +import subprocess +import sys +import threading +import time +import base64 + +import os +import hjson +from tqdm import tqdm + +from ..utils import logger +from .constants import AUTOTUNING, AUTOTUNING_METRIC_PATH, BUFSIZE +from .utils import get_val_by_key, search_error, was_interruptted +""" +thread-0: loop over experiment queue dispatching experiments if they become available +thread-N: start each experiment in its own thread +""" + +from deepspeed import comm as dist + +TIMEOUT = 5 + + +class ResourceManager: + + def __init__(self, args, hosts, num_gpus_per_node, results_dir, exps_dir, arg_mappings): + self.results_dir = results_dir + self.exps_dir = exps_dir + + self.nodes = [] + self.num_gpus_per_node = num_gpus_per_node + for host in hosts: + self.nodes.append(Node(host, num_gpus_per_node)) + + self.experiment_queue = [] + self.running_experiments = {} + self.finished_experiments = {} + self.experiment_count = 0 + self.exp_paths = set() + self.args = args + + self.arg_mappings = {} + if arg_mappings is not None: + for k, v in arg_mappings.items(): + k = k.strip() + v = v.strip() + if k not in self.arg_mappings: + self.arg_mappings[k] = v + + def schedule_experiments(self, exp_paths): + for exp_path in exp_paths: + if exp_path in self.exp_paths: + continue + else: + self.exp_paths.add(exp_path) + with open(exp_path, "r") as fd: + exp = hjson.load(fd) + exp["exp_id"] = self.experiment_count + self.experiment_count += 1 + + result_dir = exp["result_dir"] = os.path.join(self.results_dir, exp['name']) + if AUTOTUNING in exp["ds_config"]: + metric_file = os.path.join(result_dir, "metrics.json") + exp["ds_config"][AUTOTUNING][AUTOTUNING_METRIC_PATH] = metric_file + stderr_file = os.path.join(result_dir, "stderr.log") + model_info_file = os.path.join(result_dir, "model_info.json") + metric_file = os.path.join(result_dir, "metrics.json") + + # skip existing experiments (except for the ones that were interrupted) + if os.path.exists(result_dir) and os.path.exists(stderr_file): + if not was_interruptted(stderr_file): + err = search_error(stderr_file) + exp_id = exp["exp_id"] + self.finished_experiments[exp_id] = (exp, err) + if err or os.path.exists(metric_file) or os.path.exists(model_info_file): + logger.info(f"Skipping exp {exp['name']} whose result already exists") + continue + + self.experiment_queue.append(exp) + + def run_job(self, exp: dict, reservations): + exp_id = exp["exp_id"] + exp["master_port"] = self.args.master_port + exp_id + exp["result_dir"] = os.path.join(self.results_dir, exp['name']) + user_script = self.args.user_script + user_args = self.args.user_args + + # overwrite the user arg in the arg_mappings + for key, val in self.arg_mappings.items(): + nval = get_val_by_key(exp, key) + if nval and str(nval) != "auto": + if val in user_args: + idx = user_args.index(val) + user_args[idx + 1] = str(nval) + else: + user_args.append(val) + user_args.append(str(nval)) + + t = threading.Thread(target=run_experiment, args=(exp, reservations, user_script, user_args)) + t.start() + self.running_experiments[exp_id] = (t, exp, reservations, time.time()) + + def experiment_check(self, pbar): + finished_exps = [] + for exp_id, exp_data in self.running_experiments.items(): + thread, exp_json, reservations, start_time = exp_data + logger.debug(f"Checking exp_id = {exp_id}, alive = {thread.is_alive()}") + thread.join(timeout=TIMEOUT) + if not thread.is_alive(): + exp_dir = exp_json["result_dir"] + stderr_file = os.path.join(exp_dir, "stderr.log") + err = search_error(stderr_file) + finished_exps.append((exp_id, reservations)) + self.finished_experiments[exp_id] = (exp_json, err) + duration = time.time() - start_time + logger.debug(f"Finished exp_id = {exp_id}, duration={duration:.2f} sec") + pbar.update(len(finished_exps)) + for exp_id, reservations in finished_exps: + for reservation in reservations: + reservation.restore_slots() + self.running_experiments.pop(exp_id) + time.sleep(TIMEOUT) + + def resource_request(self, exp): + num_gpus, num_nodes = exp['num_gpus'], exp['num_nodes'] + slot_request = num_gpus + reservations = [] + for node in self.nodes: + if num_nodes == 0: + break + slots = node.reserve_slots(slot_request=slot_request) + if slots: + reservations.append(Reservation(node=node, slots=slots)) + num_nodes -= 1 + + if num_nodes == 0: + # request satisfied + return reservations + else: + # request not satisfied + for reservation in reservations: + reservation.restore_slots() + + def status(self): + status = "" + for node in self.nodes: + status += f"{node.host} ({len(node.idle_slots)} idle gpus), " + return status[:-1] + + def run(self): + pbar = tqdm(total=len(self.experiment_queue)) + + while len(self.experiment_queue) > 0: + exp = self.experiment_queue.pop(0) + logger.debug(f'Popped exp_id = {exp["exp_id"]} from the queue') + logger.debug(f'Resource status: {self.status()}') + reservations = self.resource_request(exp) + + if not reservations: + logger.debug(f'Unable to schedule exp_id = {exp["exp_id"]}') + self.experiment_queue.insert(0, exp) + logger.debug(f'Put exp_id = {exp["exp_id"]} back into the queue') + self.experiment_check(pbar) + else: + desc = "" + for reservation in reservations: + reservation.slots.sort() + slots = ",".join(map(str, reservation.slots)) + desc += f"{reservation.node.host}:{slots}@" + desc = desc[:-1] + logger.debug(f'Running exp_id = {exp["exp_id"]} on {desc}') + self.run_job(exp, reservations) + + # All pending experiments are scheduled, waiting for them to complete + while len(self.running_experiments) > 0: + self.experiment_check(pbar) + + def save_exp_results_to_database(self, message, ranks=None, path=None): + """Print message when one of following condition meets + + + not dist.is_initialized() + + dist.get_rank() in ranks if ranks is not None or ranks = [-1] + + Args: + message (str) + ranks (list) + path (str) + + """ + should_log = not dist.is_initialized() + ranks = ranks or [] + my_rank = dist.get_rank() if dist.is_initialized() else -1 + if ranks and not should_log: + should_log = ranks[0] == -1 + should_log = should_log or (my_rank in set(ranks)) + logger.debug(f"*** Should log: {should_log}") + if should_log: + message['rank'] = my_rank + with open(path, 'a') as outfile: + json.dump(message, outfile) + outfile.write('\n') + + def parse_results(self, metric): + """ Parses the metric file of the finished experiments to select the optimal DeepSpeed configuration. + + Args: + finished_experiments (dcit): a dictionary of experiment id and experiment description. + + Returns: + The path to the result folder of the experiment with the optimal configuration. + """ + max_throughput = sys.float_info.min + best_exp_id = -1 + for exp_id, (exp, err) in self.finished_experiments.items(): + if err: + logger.info( + f"The experiment exp_id = {exp_id}, exp_name = {exp['name']}, did not run successfully with error = {err}, thus a metrics.txt does not exist for it. Check the stderr.log in {exp['result_dir']}" + ) + continue + + metric_file = exp["ds_config"][AUTOTUNING][AUTOTUNING_METRIC_PATH] + + if os.path.exists(metric_file): + with open(metric_file, 'r') as f: + results = hjson.load(f) + curr_throughput = results[metric] + if curr_throughput > max_throughput: + max_throughput = curr_throughput + best_exp_id = exp_id + exp['results'] = results + + if best_exp_id != -1: + best_exp, _ = self.finished_experiments[best_exp_id] + return best_exp, max_throughput + + return exp, None + + def clear(self): + """Clear experiment queues, does not reset self.experiment_count + """ + self.experiment_queue = [] + # clean up the running experiments + for exp_id, exp_data in self.running_experiments.items(): + thread, exp_json, reservations, start_time = exp_data + clean_up(exp_json, reservations) + self.running_experiments = {} + self.finished_experiments = {} + self.exp_paths = set() + + +class Node: + + def __init__(self, host, max_slots): + self.host = host + self.max_slots = max_slots + self.idle_slots = list(range(max_slots)) + + def reserve_slots(self, slot_request: int) -> list: + if len(self.idle_slots) >= slot_request: + return [self.idle_slots.pop(0) for _ in range(slot_request)] + + def restore_slots(self, slots: list): + self.idle_slots += slots + + +class Reservation: + + def __init__(self, node, slots): + self.node = node + self.slots = slots + + def restore_slots(self): + self.node.restore_slots(self.slots) + + def desc(self): + slots = ",".join(map(str, self.slots)) + return f"{self.node.host}:{slots}@" + + +def get_job_id(): + # Infrastructure-specific job-id + infra_job_id = None + if "DLWS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLWS_JOB_ID"] + elif "DLTS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLTS_JOB_ID"] + else: + infra_job_id = "unknown-job-id" + + return infra_job_id + + +def get_user(): + user = None + if "USER" in os.environ: + user = os.environ["USER"] + else: + user = "unknown-user" + return user + + +def run_experiment(exp: dict, reservations, user_script, user_args): + include_str = "" + for reservation in reservations: + reservation.slots.sort() + slots = ",".join(map(str, reservation.slots)) + include_str += f"{reservation.node.host}:{slots}@" + include_str = include_str[:-1] + master_port = exp["master_port"] + hostfile = exp["hostfile"] + exp["launcher_args"] = [ + "--hostfile", + f"{hostfile}", + "--include", + f"{include_str}", + "--master_port", + str(master_port), + ] + logger.debug(f'launcher args={exp["launcher_args"]}') + + exp["user"] = get_user() + exp["job_id"] = get_job_id() + exp_dir = exp["result_dir"] + os.makedirs(exp_dir, exist_ok=True) + ds_config_path = os.path.join(exp_dir, "ds_config.json") + exp["ds_config_path"] = ds_config_path + + ds_config = copy.deepcopy(exp["ds_config"]) + ds_config_json = json.dumps(ds_config).encode('utf-8') + + exp["ds_config_base64"] = base64.urlsafe_b64encode(ds_config_json).decode('utf-8') + + with open(exp["ds_config_path"], "w", buffering=BUFSIZE) as fd: + json.dump(ds_config, fd) + fd.flush() + os.fsync(fd) + path = exp["ds_config_path"] + logger.info(f"Scheduler wrote ds_config to {path}, {os.path.abspath(path)}") + + with open(os.path.join(exp_dir, "exp.json"), "w", buffering=BUFSIZE) as fd: + json.dump(exp, fd) + fd.flush() + os.fsync(fd) + path = os.path.join(exp_dir, "exp.json") + logger.info(f"Scheduler wrote exp to {path}, {os.path.abspath(path)}") + + # remove "--deepspeed_config ds_config.json" from user_args + if user_args: + if "--deepspeed_config" in user_args: + idx = user_args.index("--deepspeed_config") + # "--deepspeed_config" is omitted in HF + elif "--deepspeed" in user_args: + idx = user_args.index("--deepspeed") + assert idx < len(user_args), "there is no ds_config file specified after --deepspeed_config or --deepspeed" + # user_args[idx + 1] = exp["ds_config_path"] + # pass base64 serialized ds_config to launcher + user_args[idx + 1] = exp["ds_config_base64"] + + exp["user_script"] = user_script + exp["user_args"] = user_args + + cmd = ["deepspeed"] + exp["launcher_args"] + [user_script] + user_args + + assert len(exp["launcher_args"]) > 0, "must provide launcher args" + + with open(os.path.join(exp_dir, "cmd.txt"), "w", buffering=BUFSIZE) as fd: + fd.write(" ".join(cmd)) + fd.write("\n") + fd.flush() + os.fsync(fd) + + logger.info( + f"Launching exp_id = {exp['exp_id']}, exp_name = {exp['name']}, with resource = {include_str}, and ds_config = {os.path.abspath(ds_config_path)}" + ) + + with open(os.path.join(exp_dir, "stdout.log"), "wb") as out, open(os.path.join(exp_dir, "stderr.log"), + "wb") as err: + result = subprocess.Popen(cmd, stdout=out, stderr=err) + result.wait() + out.flush() + err.flush() + os.fsync(out) + os.fsync(err) + + clean_up(exp, reservations) + + logger.info(f"Done running exp_id = {exp['exp_id']}, exp_name = {exp['name']}, with resource = {include_str}") + + +PDSH_MAX_FAN_OUT = 1024 + + +def clean_up(exp: dict, reservations): + env = os.environ.copy() + env['PDSH_RCMD_TYPE'] = 'ssh' + + nodes_str = "" + for reservation in reservations: + nodes_str += f"{reservation.node.host}," + nodes_str = nodes_str[:-1] + logger.debug(f"Cleaning up exp_id = {exp['exp_id']} on the following workers: {nodes_str}") + + # PDSH flags for max node fan out and specific hosts to launch on + # See https://linux.die.net/man/1/pdsh for flag details + pdsh_cmd = ['pdsh', '-f', str(PDSH_MAX_FAN_OUT), '-w', nodes_str] + + kill_cmd = [ + 'pkill', + '-f', + exp['name'], + ] + cmd = pdsh_cmd + kill_cmd + logger.debug("cmd = {}".format(' '.join(cmd))) + + result = subprocess.Popen(cmd, env=env) + result.wait() + + # In case of failure must propagate the error-condition back to the caller (usually shell). The + # actual error and traceback should have been printed in the subprocess, so in order to avoid + # unnecessary noise we just quietly exit here with the same code as the subprocess + if result.returncode > 0: + sys.exit(result.returncode) + + logger.info(f"Done cleaning up exp_id = {exp['exp_id']} on the following workers: {nodes_str}") diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__init__.py b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..676ae429e07745d3ad24051a2610e57ac42601f0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .index_based_tuner import RandomTuner, GridSearchTuner +# from .ga_tuner import GATuner +from .model_based_tuner import ModelBasedTuner diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..247257cce6d6898d8cae50f1979f10e6f5cb77f4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/base_tuner.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/base_tuner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..666ad71268d69e2748d534c9a418e0b27c5d585c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/base_tuner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/cost_model.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/cost_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aeec3f2bc6e4e9a8207ab3e6e31386ec5b68ed85 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/cost_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/index_based_tuner.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/index_based_tuner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c852ef0ed3be5f9a4b30bb5ae70bbffae8e44aba Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/index_based_tuner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/model_based_tuner.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/model_based_tuner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea8ed348145319bee27ca777e86b65f6f4fc4fb5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/model_based_tuner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a53e7d28a58d7d8f24ff952b633d94d04fb8fb0f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/base_tuner.py b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/base_tuner.py new file mode 100644 index 0000000000000000000000000000000000000000..b2da065e44f4074163dc22d1e2ef29d9d46a59f3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/base_tuner.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys + +from deepspeed.autotuning.constants import * +from deepspeed.autotuning.utils import write_experiments +from deepspeed.utils import logger + + +class BaseTuner: + + def __init__(self, exps, resource_manager, metric): + self.all_exps = exps + self.rm = resource_manager + self.best_iter = 0 + self.best_exp = None + self.best_metric_val = None + self.metric = metric if metric else AUTOTUNING_METRIC_DEFAULT + logger.info(f"total number of exps = {len(self.all_exps)}") + + def has_next(self): + """Whether there exists more configurations for evaluation""" + if len(self.all_exps) > 0: + return True + else: + return False + + def next_batch(self, sample_size): + """Select the next batch of configurations for evaluation""" + raise NotImplementedError + + def update(self): + """"Update the tuner with what configurations have been evaluated and their performance results""" + + def tune(self, sample_size=1, n_trials=1000, early_stopping=None): + i = 0 + try: + while i < n_trials and self.has_next(): + # Select the next batch of configuration for evaluation + sampled_exps = self.next_batch(sample_size) + # Generate experiments for measurement of performance + exp_paths = write_experiments(sampled_exps, self.rm.exps_dir) + self.rm.schedule_experiments(exp_paths) + self.rm.run() + exp, metric_val = self.rm.parse_results(self.metric) + if self.best_exp is None or self.best_metric_val is None or (metric_val + and metric_val > self.best_metric_val): + # logger.info(f"tuner finds better = {exp}") + self.best_exp = exp + self.best_metric_val = metric_val + self.best_iter = i + + i += len(sampled_exps) + + # Update the tuner with evaluated performance results + self.update() + + self.rm.clear() + + # Early stop if no more promising configurations are likely to be found + if early_stopping and i >= self.best_iter + early_stopping: + logger.info( + f"Tuner early stopped at iteration {i}. Best iteration is {self.best_iter}. Early stopping threshold is {early_stopping}" + ) + break + return i + except: + logger.info("Tuner Error:", sys.exc_info()[0]) + return i diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/cost_model.py b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/cost_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c12b10f743632c36a61711a411e8bb706041b762 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/cost_model.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .utils import * + +try: + import xgboost as xgb +except ImportError: + xgb = None + + +class XGBoostCostModel(): + + def __init__(self, loss_type, num_threads=None, log_interval=25, upper_model=None): + + assert xgb is not None, "missing requirements, please install deepspeed w. 'autotuning_ml' extra." + + self.loss_type = loss_type + + if loss_type == "reg": + self.xgb_params = { + "max_depth": 3, + "gamma": 0.0001, + "min_child_weight": 1, + "subsample": 1.0, + "eta": 0.3, + "lambda": 1.0, + "alpha": 0, + "objective": "reg:linear", + } + elif loss_type == "rank": + self.xgb_params = { + "max_depth": 3, + "gamma": 0.0001, + "min_child_weight": 1, + "subsample": 1.0, + "eta": 0.3, + "lambda": 1.0, + "alpha": 0, + "objective": "rank:pairwise", + } + else: + raise RuntimeError("Invalid loss type: " + loss_type) + + self.xgb_params["verbosity"] = 0 + if num_threads: + self.xgb_params["nthread"] = num_threads + + def fit(self, xs, ys): + x_train = np.array(xs, dtype=np.float32) + y_train = np.array(ys, dtype=np.float32) + y_max = np.max(y_train) + y_train = y_train / max(y_max, 1e-9) + + index = np.random.permutation(len(x_train)) + dtrain = xgb.DMatrix(x_train[index], y_train[index]) + + self.bst = xgb.train(self.xgb_params, dtrain) + + def predict(self, xs): + + features = xgb.DMatrix(xs) + + return self.bst.predict(features) diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/index_based_tuner.py b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/index_based_tuner.py new file mode 100644 index 0000000000000000000000000000000000000000..d3c822be0d35ff68a42355322797cc4b0c8c1429 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/index_based_tuner.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import random + +from .base_tuner import BaseTuner + + +class RandomTuner(BaseTuner): + """Explore the search space in random order""" + + def __init__(self, exps: list, resource_manager, metric): + super().__init__(exps, resource_manager, metric) + + def next_batch(self, sample_size=1): + if sample_size > len(self.all_exps): + sample_size = len(self.all_exps) + + sampled_batch = random.sample(self.all_exps, sample_size) + self.all_exps = [x for x in self.all_exps if x not in sampled_batch] + + return sampled_batch + + +class GridSearchTuner(BaseTuner): + """Explore the search space in sequential order""" + + def __init__(self, exps: list, resource_manager, metric): + super().__init__(exps, resource_manager, metric) + + def next_batch(self, sample_size=1): + if sample_size > len(self.all_exps): + sample_size = len(self.all_exps) + + sampled_batch = self.all_exps[0:sample_size] + self.all_exps = [x for x in self.all_exps if x not in sampled_batch] + + return sampled_batch diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/model_based_tuner.py b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/model_based_tuner.py new file mode 100644 index 0000000000000000000000000000000000000000..aec9264f9b7c8623923fd62cd0145871fbc74214 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/model_based_tuner.py @@ -0,0 +1,157 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import hjson + +from ..constants import AUTOTUNING, AUTOTUNING_METRIC_PATH +from .base_tuner import BaseTuner +from .cost_model import XGBoostCostModel +from .utils import * +from ..utils import * +import numbers +from ..constants import AUTOTUNING_METRIC_LATENCY + +INIT_NUM = 2 + + +class ModelBasedTuner(BaseTuner): + """Exploring the search space with a cost model""" + + def __init__(self, exps: list, resource_manager, metric, tuning_space): + super().__init__(exps, resource_manager, metric) + self.tuning_space = tuning_space + self.best_iter = 0 + + self.all_configs = [e['ds_config'] for e in exps] + self.num_all_configs = len(self.all_configs) + + self.dims = dict_to_dims(self.tuning_space) + + logger.info(f"Create config dim: {self.dims}, all configs: {self.num_all_configs}") + + self.visited = set([]) + + self.trials = [] + self.trial_pt = 0 + + init_num = min(INIT_NUM, self.num_all_configs) + + for _ in range(init_num): + exp_feature = np.random.randint(self.num_all_configs) + exp_feature = 0 + while exp_feature in self.visited: + exp_feature = np.random.randint(self.num_all_configs) + self.trials.append(exp_feature) + self.visited.add(exp_feature) + + self.cost_model = XGBoostCostModel("rank") + + self.evaluated_configs = [] + self.evaluated_perf = [] + + self.train_ct = 0 + + self.random_exploration_ratio = 0.2 # do random exploration + + def find_estimated_top_configs(self): + """Use the cost model to predict the estimated performance of configurations and find the top ones for the next round of evaluation""" + + configs = [] + + for c in self.all_configs: + flattened_ds_config = flatten(c) + feature_val = [] + for k, v in flattened_ds_config.items(): + if isinstance(v, numbers.Number): + feature_val.append(v) + configs.append(feature_val) + # print(configs) + # TODO the current implementation requires that all configs have the same shape. + configs = np.array(configs, dtype=np.float32) + estimates = self.cost_model.predict(configs) + + n = len(estimates) + top_idx = np.argsort(estimates) + top_idx_ret = top_idx if self.metric == AUTOTUNING_METRIC_LATENCY else top_idx[::-1][:n] + + # top_configs = [self.all_configs[i] for i in top_idx] + + return top_idx_ret + + def next_batch(self, sample_size): + sampled_batch = [] + + counter = 0 + while counter < sample_size: + + if len(self.visited) >= self.num_all_configs: + break + + while self.trial_pt < len(self.trials): + logger.debug(f"trials: {self.trials}") + # Select top promising trials + index = self.trials[self.trial_pt] + if index not in self.visited: + break + self.trial_pt += 1 + + # To avoid over-exploitation, randomly select one that has not been explored. + rand = np.random.rand() + if rand < self.random_exploration_ratio: + # Do normal selection + feature = np.random.choice(self.trials) + while index in self.visited: + index = np.random.randint(self.num_all_configs) + + # Need to track both the sampled configs and indices + + sampled_batch.append(self.all_exps[index]) + self.visited.add(index) + counter += 1 + + return sampled_batch + + def has_next(self): + return len(self.visited) < self.num_all_configs + + def update(self): + for exp_id, (exp, err) in self.rm.finished_experiments.items(): + feature_val = [] + if err: + logger.info( + f"Skipping exp_id = {exp_id}, exp_name = {exp['name']}, the experiment did not run successfully with error = {err}, thus a metrics.txt does not exist for it. Please check the stderr.log in {exp['result_dir']}" + ) + ds_config = exp["ds_config"] + flattened_ds_config = flatten(ds_config) + for k, v in flattened_ds_config.items(): + if isinstance(v, numbers.Number): + feature_val.append(v) + self.evaluated_configs.append(feature_val) + self.evaluated_perf.append(0.0) + continue + + p = exp["ds_config"][AUTOTUNING][AUTOTUNING_METRIC_PATH] + with open(p, 'r') as f: + results = hjson.load(f) + curr_iter = results[self.metric] + logger.debug(f"parsing the results for {exp_id}， Result is {curr_iter}") + + ds_config = exp["ds_config"] + flattened_ds_config = flatten(ds_config) + for k, v in flattened_ds_config.items(): + if isinstance(v, numbers.Number): + feature_val.append(v) + self.evaluated_configs.append(feature_val) + self.evaluated_perf.append(curr_iter) + + logger.debug(f"**Evaluated configs: {len(self.evaluated_configs)}, evaluated perf: {self.evaluated_perf}") + + self.cost_model.fit(self.evaluated_configs, self.evaluated_perf) + + estimated_top_configs = self.find_estimated_top_configs() + + self.trials = estimated_top_configs + self.trial_pt = 0 + self.train_ct += 1 diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/tuner/utils.py b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ada643f2c02ce9b0433c89295e932f9f49740eac --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/tuner/utils.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import numpy as np +import itertools +from ..utils import * +import collections.abc + + +def index_to_feature(p, dims): + """convert index form (single integer) to feature form (vector)""" + feature = [] + for dim in dims: + feature.append(p % dim) + p //= dim + return feature + + +def feature_to_index(feature, dims): + """convert feature form (vector) to index form (single integer)""" + p = 0 + for j, k in enumerate(feature): + print("j:", "k:", k, "dims", dims[:j]) + p += int(np.prod(dims[:j])) * k + return p + + +def dict_to_dims(tuning_space): + + dims = [] + + for key, val in tuning_space.items(): + if isinstance(val, dict): + dims.extend(dict_to_dims(val)) + elif isinstance(val, list): + dims.append(len(val)) + else: + dims.append(1) + + return dims + + +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)) + + +def flatten(d, parent_key='', sep='_'): + items = [] + for k, v in d.items(): + new_key = parent_key + sep + k if parent_key else k + if isinstance(v, collections.abc.MutableMapping): + items.extend(flatten(v, new_key, sep=sep).items()) + else: + items.append((new_key, v)) + return dict(items) + + +def dict_to_feature(feature_dict, keys, max_value=None): + """Extract values from dict""" + feature = [] + for key, val in feature_dict.items(): # First level + if key not in keys: + continue + if val is None or val == "auto" or key == "autotuning" or val == "": + continue + if isinstance(val, dict): + feature.append(dict_to_feature(val, max_value)) + else: + feature.append(float(val)) + + # normalization, should not matter in tree models + if max_value is not None: + norm_feature = [] + for f, mv in zip(feature, max_value): + norm_feature.append(f / mv) + feature = norm_feature + + return feature diff --git a/lib/python3.12/site-packages/deepspeed/autotuning/utils.py b/lib/python3.12/site-packages/deepspeed/autotuning/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b851353520fb87a5d9680644d0aa447cfa28d32f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/autotuning/utils.py @@ -0,0 +1,459 @@ +# 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) + " " diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__init__.py b/lib/python3.12/site-packages/deepspeed/checkpoint/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1f645a494ade557512ceaeb994f05663a910b43b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .reshape_meg_2d import reshape_meg_2d_parallel + +from .deepspeed_checkpoint import DeepSpeedCheckpoint + +from .utils import (get_layer_ckpt_name_for_rank, get_model_ckpt_name_for_rank, get_zero_ckpt_name_for_rank) + +from .reshape_utils import (merge_state) + +from .reshape_3d_utils import (model_3d_desc, get_model_3d_descriptor) + +from .zero_checkpoint import ZeROCheckpoint + +from .universal_checkpoint import enable_universal_checkpoint, SubparamShape + +from .constants import * diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd3f62065c383f2f563880e1563e403759c6c57d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9905111afd19ab4b2bc911c893419438f5a966d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/deepspeed_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/deepspeed_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59b256f442cd306b1fed3cf91f6bdbed95d288cd Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/deepspeed_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/ds_to_universal.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/ds_to_universal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f8d55fd54ea452c1ea468000d857b3cbbebf657 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/ds_to_universal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_3d_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_3d_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9281f2bc3504efedcf5b5271c92c2075ec07f32e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_3d_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_meg_2d.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_meg_2d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bd62088d1415a2172ccc75f0309bfa77d2c0a68 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_meg_2d.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37d61901adf0812e6a74fa9fcae4aaa7bd609c4a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/reshape_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/universal_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/universal_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f526d0ef9f5bd94cb00b1894272f9724d50b84d1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/universal_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94a7a566710a6d550123f81c7e35e36d8d6d9b73 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/zero_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/zero_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85e1092774a2a47f1d6c9e9b7c1d6f46ae716b2e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/checkpoint/__pycache__/zero_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/constants.py b/lib/python3.12/site-packages/deepspeed/checkpoint/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..046bc242002fc3702fc35191d38fe3f41d506fdb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/constants.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Various symbolic constants used for model checkpointing +""" + +######################################### +# Optimizer checkpoint keys +######################################### +OPTIMIZER_STATE_DICT = "optimizer_state_dict" +FP32_GROUPS = "fp32_groups" +FP32_FLAT_GROUPS = 'fp32_flat_groups' + +BASE_OPTIMIZER_STATE = 'base_optimizer_state' +BASE_OPTIMIZER_STATE_STEP = 'base_optimizer_state_step' +SINGLE_PARTITION_OF_FP32_GROUPS = "single_partition_of_fp32_groups" +PARAM_GROUPS = 'param_groups' +GROUP_PADDINGS = 'group_paddings' +PARTITION_COUNT = 'partition_count' +ZERO_STAGE = 'zero_stage' +CLIP_GRAD = 'clip_grad' +FP32_WEIGHT_KEY = "fp32" +LOSS_SCALER = 'loss_scaler' + +######################################### +# Module checkpoint keys +######################################### +PARAM = 'param' +PARAM_SHAPES = 'param_shapes' +BUFFER_NAMES = 'buffer_names' +FROZEN_PARAM_SHAPES = 'frozen_param_shapes' +FROZEN_PARAM_FRAGMENTS = 'frozen_param_fragments' + +######################################### +# Checkpoint naming constants +######################################### +MODEL_FILE_PREFIX = 'mp_rank_' +ZERO_FILE_PREFIX = 'zero_pp_rank_' +OPTIM_FILE_SUFFIX = '_optim_states.pt' +MODEL_FILE_SUFFIX = '_model_states.pt' +LAYER_FILE_PREFIX = 'layer_' +BF16_ZERO_FILE_PREFIX = 'bf16_' + ZERO_FILE_PREFIX +FP16_ZERO_FILE_PREFIX = 'fp16_' + ZERO_FILE_PREFIX + +######################################### +# Checkpoint utility keys +######################################### +DS_VERSION = 'ds_version' + +######################################### +# Universal Checkpoint keys +######################################### +UNIVERSAL_CHECKPOINT_INFO = 'universal_checkpoint_info' +UNIVERSAL_CHECKPOINT_VERSION_KEY = 'universal_checkpoint_version' +# Reserve version 0.1 for the hardcoded logic used in BLOOM-176B training +UNIVERSAL_CHECKPOINT_VERSION_VALUE = 0.2 + +# Vocabulary padding +VOCAB_TENSOR = 'vocab_tensor' +PADDED_VOCAB_SIZE = 'padded_vocab_size' +ORIGINAL_VOCAB_SIZE = 'original_vocab_size' + +# Parameter splitting/merging +PARAM_SLICE_MAPPINGS = 'param_slice_mappings' +CAT_DIM = "cat_dim" +# Following is a special case where a parameter effectively contains sub parameters. +# As an example, consider Megatron-DeepSpeed GPT SWIGLU implementation (mlp.h_to_4h). +# In this case, a single parameter ia allocated contiguously, but used as separate parameters. +# When using universal checkpoint, we have to normalize the representation of the full parameter. +# We normalize it by concatenating all slices of the sub params and then concatenating the sub params. +# All concat operations are done on CAT_DIM (currently, no support for different concat dims sub params and TP slicing). +# Similarly, load_hp_checkpoint_state has to take the needed actions when loading from universal. +PARAM_N_SUB_PARAMS = "param_n_sub_params" + +SUB_PARAM_SHAPE = "sub_param_shape" + +# Regex list of parameters that require special handling +VOCABULARY_PARAMETER_PATTERNS = 'vocabulary_parameter_patterns' +PIPELINE_REPLICATED_PARAMETER_PATTERNS = 'pipeline_replicated_parameter_patterns' +PARAMETER_TO_AVERAGE_PATTERNS = 'parameter_to_average_patterns' +PARAMETER_WITH_ROW_PARALLELISM_PATTERNS = 'parameter_with_row_parallelism_patterns' +TP_REPLICATED_PARAMETER_PATTERNS = 'tp_replicated_parameter_patterns' +PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0 = 'parameter_with_2_sub_params_cat_dim_0' +PARAMETER_WITH_SUB_PARAMS = 'parameter_with_sub_params' +SUB_PARAMS_SHAPE = 'sub_params_shape' diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/deepspeed_checkpoint.py b/lib/python3.12/site-packages/deepspeed/checkpoint/deepspeed_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..9a368b7a0a2543ba50658d11ebb76132bfe56d10 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/deepspeed_checkpoint.py @@ -0,0 +1,307 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import re +from typing import Dict +import torch + +from .reshape_3d_utils import model_3d_desc +from .reshape_utils import (basic_folder_validation, merge_state, partition_data, get_files, get_files_with_prefix) + +from .constants import (MODEL_FILE_PREFIX, LAYER_FILE_PREFIX) + +from .reshape_meg_2d import reshape_meg_2d_parallel, meg_2d_parallel_map +from .zero_checkpoint import ZeROCheckpoint +from .constants import * + +EMBEDDING_LAYER_INDEX = 0 +FINAL_LAYER_NORM_INDEX = -1 +ARGS_KEY = 'args' +CHECKPOINT_INFO_KEY = 'checkpoint_info' +ITERATION_KEY = 'iteration' +LAYER_FILE_PREFIX_PATTERN = r'layer_(\d+)-model_.*' + +SEQUENTIAL_LAYERS = [ + 'input_layernorm.weight', 'input_layernorm.bias', 'self_attention.dense.bias', 'post_attention_layernorm.weight', + 'post_attention_layernorm.bias', 'mlp.dense_4h_to_h.bias', 'position_embeddings.weight' +] + +LAYER_CONCAT_DIM = {'self_attention.dense.weight': 1, 'mlp.dense_4h_to_h.weight': 1} + + +class DeepSpeedCheckpoint(object): + + def __init__(self, + dir, + tp_degree=None, + pp_degree=None, + dp_degree=None, + final_layer_norm_idx=FINAL_LAYER_NORM_INDEX): + self.final_layer_norm_idx = final_layer_norm_idx + self.dir = dir + + pipeline_parallel = len(get_files_with_prefix(get_files(dir), LAYER_FILE_PREFIX)) > 0 + + self._validate_folder(dir, pipeline_parallel) + + self.zero_checkpoint = ZeROCheckpoint(dir) + + self.file_list = get_files(dir) + self.layer_files = get_files_with_prefix(self.file_list, LAYER_FILE_PREFIX) + self.mp_rank_files = get_files_with_prefix(self.file_list, MODEL_FILE_PREFIX) + + self.layer_keys = self._get_layer_keys() + self.layer_count = len(self.layer_keys) + + self.tp_degree = self.zero_checkpoint.get_src_tp_degree() if tp_degree is None else tp_degree + self.pp_degree = self.zero_checkpoint.get_src_pp_degree() if pp_degree is None else pp_degree + self.dp_degree = self.zero_checkpoint.get_src_dp_degree() if dp_degree is None else dp_degree + + self.original_world_size = self.zero_checkpoint.get_src_tp_degree() * self.zero_checkpoint.get_src_pp_degree( + ) * self.zero_checkpoint.get_src_dp_degree() + self.world_size = self.tp_degree * self.pp_degree * self.dp_degree + + self.old_2d_map = meg_2d_parallel_map(self.zero_checkpoint.get_src_pp_degree(), + self.zero_checkpoint.get_src_tp_degree()) + self.old_2d_map.simple_init() + self.new_2d_map = reshape_meg_2d_parallel(old_pp_degree=self.zero_checkpoint.get_src_pp_degree(), + old_tp_degree=self.zero_checkpoint.get_src_tp_degree(), + new_pp_degree=self.pp_degree, + new_tp_degree=self.tp_degree) + + if self.is_change_pp_degree() or self.is_change_tp_degree() or self.is_change_dp_degree(): + self.zero_checkpoint.reshape(model_3d_desc(self.pp_degree, self.tp_degree, self.dp_degree)) + + self.global_state = {} + + self._sanity_check() + self.pp_to_transformer_map = self._build_pp_transformer_map() + self.transformer_file_map = self._build_transformer_file_map() + self.tp_to_embedding_map = self._build_tp_other_layer_map(EMBEDDING_LAYER_INDEX) + self.tp_to_final_norm_map = self._build_tp_other_layer_map(self.final_layer_norm_idx) + self._build_global_state() + + def is_change_tp_degree(self): + return self.tp_degree != self.zero_checkpoint.get_src_tp_degree() + + def is_change_pp_degree(self): + return self.pp_degree != self.zero_checkpoint.get_src_pp_degree() + + def is_change_dp_degree(self): + return self.dp_degree != self.zero_checkpoint.get_src_dp_degree() + + def show_2d_mapping(self): + print(f'reshaped 2d map ---- begin') + + for i in range(self.pp_degree): + for j in range(self.tp_degree): + file_list = self.get_2d_parallel_files(pp_index=i, tp_index=j) + print(f'[{i}, {j}] = {file_list}') + + print(f'reshaped 2d map ---- end') + + def show_tp_embedding_map(self): + self._dump_mapping(self.tp_to_embedding_map, 'tp_to_embedding_layers') + + def show_tp_final_norm_map(self): + self._dump_mapping(self.tp_to_final_norm_map, 'tp_to_final_norm_layers') + + def show_pp_transformer_map(self): + self._dump_mapping(self.pp_to_transformer_map, 'pp_to_transformer_layers') + + def show_transformer_file_map(self): + self._dump_mapping(self.transformer_file_map, 'rank_to_transformer_files') + + def _build_global_state(self): + sd = torch.load(self.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False) + self.global_state[ITERATION_KEY] = sd.get(ITERATION_KEY, 0) + self.global_state[ARGS_KEY] = sd.get(ARGS_KEY, None) + + def get_zero_checkpoint_state(self, pp_index, tp_index, dp_index) -> dict: + return self.zero_checkpoint.get_state_for_rank(pp_index=pp_index, + tp_index=tp_index, + dp_index=dp_index, + keys_to_ignore=[PARAM_SHAPES]) + + def get_zero_files(self, pp_index, tp_index, dp_index) -> list: + return self.zero_checkpoint.get_files_for_rank(pp_index=pp_index, tp_index=tp_index, dp_index=dp_index) + + def get_embedding_layer_id(self): + return self.layer_keys[EMBEDDING_LAYER_INDEX] + + def get_final_norm_layer_id(self): + return self.layer_keys[self.final_layer_norm_idx] + + def get_iteration(self): + if not ITERATION_KEY in self.global_state: + sd = torch.load(self.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False) + self.global_state[ITERATION_KEY] = sd.get(ITERATION_KEY, 0) + + return self.global_state[ITERATION_KEY] + + def get_embedding_state(self, tp_index: int) -> Dict: + assert tp_index in self.tp_to_embedding_map.keys() + sd_list = [ + torch.load(fname, map_location=torch.device('cpu'), weights_only=False) + for fname in self.tp_to_embedding_map[tp_index] + ] + sd = self._merge_state_dicts(sd_list) + return sd + + def get_embedding_files(self, tp_index: int) -> list: + assert tp_index in self.tp_to_embedding_map.keys() + return self.tp_to_embedding_map[tp_index] + + def _get_checkpoint_value(self, key): + if not key in self.global_state: + sd = torch.load(self.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False) + self.global_state[key] = sd.get(key, None) + + return self.global_state[key] + + def get_args(self): + return self._get_checkpoint_value(ARGS_KEY) + + def get_checkpoint_info(self, info_key=CHECKPOINT_INFO_KEY): + return self._get_checkpoint_value(info_key) + + def get_2d_parallel_state(self, tp_index: int, pp_index: int) -> dict: + assert tp_index < self.tp_degree + assert pp_index < self.pp_degree + fname_list = self.get_2d_parallel_files(tp_index=tp_index, pp_index=pp_index) + sd_list = [torch.load(fname, map_location=torch.device('cpu'), weights_only=False) for fname in fname_list] + + merged_sd = None + for sd in sd_list: + if merged_sd is None: + merged_sd = sd + else: + merged_sd = merge_state(merged_sd, sd) + + return merged_sd + + def get_transformer_state(self, tp_index: int, pp_index: int) -> list: + assert tp_index < self.tp_degree + assert pp_index < self.pp_degree + t_list = [] + for fname_list in self.transformer_file_map[(tp_index, pp_index)]: + sd_list = [torch.load(fname, map_location=torch.device('cpu'), weights_only=False) for fname in fname_list] + sd = self._merge_state_dicts(sd_list) + t_list.append(sd) + return t_list + + def get_pp_transformer_map(self, pp_index: int) -> list: + assert pp_index < self.pp_degree + return self.pp_to_transformer_map[pp_index] + + def get_final_norm_state(self, tp_index: int) -> Dict: + assert tp_index in self.tp_to_final_norm_map.keys() + sd = torch.load(self.tp_to_final_norm_map[tp_index][0], map_location=torch.device('cpu'), weights_only=False) + return sd + + def get_final_norm_files(self, tp_index: int) -> list: + assert tp_index in self.tp_to_final_norm_map.keys() + return self.tp_to_final_norm_map[tp_index] + + def _build_tp_other_layer_map(self, layer_index: int): + data_map = {} + if len(self.layer_files) < 1: + return data_map + assert layer_index <= len(self.layer_files) + layer_files = get_files_with_prefix(self.layer_files, self.layer_keys[layer_index]) + layer_file_partitions = partition_data(layer_files, self.tp_degree) + data_map = {i: flist for i, flist in enumerate(layer_file_partitions)} + return data_map + + def get_2d_parallel_files(self, tp_index: int, pp_index: int) -> list: + assert tp_index < self.tp_degree + assert pp_index < self.pp_degree + file_indices = self.new_2d_map.get_data(pp_index=pp_index, tp_index=tp_index) + return [self.mp_rank_files[i] for i in file_indices] + + def _build_pp_transformer_map(self): + data_map = {} + if self.pp_degree > 0: + transformer_layers = self.layer_keys[1:self.final_layer_norm_idx] + layers_per_pp = len(transformer_layers) // self.pp_degree + data_map = { + i: transformer_layers[i * layers_per_pp:(i + 1) * layers_per_pp] + for i in range(0, self.pp_degree) + } + return data_map + + def _dump_mapping(self, data_map, map_tag=None): + if map_tag is not None: + print(f'Dump mapping: {map_tag}') + for k, v in data_map.items(): + print(f'{k} = {v}') + + def _build_transformer_file_map(self): + transformer_layer_keys = self.layer_keys[1:self.final_layer_norm_idx] + file_map = {} + # XXX: this is not guaranteed + layers_per_pp = 1 + if self.pp_degree > 0: + layers_per_pp = len(transformer_layer_keys) // self.pp_degree + #print(f"{transformer_layer_keys} {layers_per_pp}") + for key_index, layer_key in enumerate(transformer_layer_keys): + pp_index = key_index // layers_per_pp + layer_files = get_files_with_prefix(self.layer_files, layer_key + '-') + layer_file_partitions = partition_data(layer_files, self.tp_degree) + for tp_index in range(self.tp_degree): + map_key = (tp_index, pp_index) + if not map_key in file_map.keys(): + file_map[map_key] = [] + file_map[map_key].append(layer_file_partitions[tp_index]) + + return file_map + + def _sanity_check(self): + assert len(self.mp_rank_files) % self.tp_degree == 0 + assert self.zero_checkpoint.num_files % (self.pp_degree * self.tp_degree) == 0 + assert self.zero_checkpoint.num_files % (self.tp_degree) == 0 + # XXX: fix me - isn't always the case + # only true with --pp-partition-method 'type:transformer|embedding' \ + # assert (len(self.layer_keys) - 2) % self.pp_degree == 0 + + def validate_files(self): + for file in self.file_list: + if not os.path.isfile(file): + print(f'Error: {file} is not existent') + + def _get_layer_keys(self): + key_set = set() + for file_path in self.layer_files: + _, fname = os.path.split(file_path) + layer_id = re.search(LAYER_FILE_PREFIX_PATTERN, fname).group(1) + key_set.add(layer_id) + sorted_ids = sorted(list(key_set), key=int) + layer_keys = [LAYER_FILE_PREFIX + str(layer_id) for layer_id in sorted_ids] + return layer_keys + + def _merge_state_dicts(self, sd_list): + merged_sd = {} + for key in sd_list[0].keys(): + if not key in SEQUENTIAL_LAYERS: + cat_dim = LAYER_CONCAT_DIM.get(key, 0) + merged_sd[key] = torch.cat([sd[key] for sd in sd_list], dim=cat_dim) + else: + merged_sd[key] = sd_list[0][key] + + return merged_sd + + def _validate_folder(self, dir, pipeline_parallel): + basic_folder_validation(dir) + + file_list = get_files(dir) + file_prefix_list = [MODEL_FILE_PREFIX] + if pipeline_parallel: + file_prefix_list.extend([LAYER_FILE_PREFIX, f'{LAYER_FILE_PREFIX}01']) + for file_prefix in file_prefix_list: + ckpt_files = get_files_with_prefix(file_list, file_prefix) + assert len( + ckpt_files + ) > 0, f'{dir} seems a bogus DeepSpeed checkpoint folder: Cannot find {file_prefix}* files in there.' diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/ds_to_universal.py b/lib/python3.12/site-packages/deepspeed/checkpoint/ds_to_universal.py new file mode 100644 index 0000000000000000000000000000000000000000..f7b75eee66d06610a91d87019f37c153368b8323 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/ds_to_universal.py @@ -0,0 +1,549 @@ +#!/usr/bin/env python + +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from functools import partial +from itertools import chain +import argparse +import glob +import itertools +import math +from concurrent.futures import ProcessPoolExecutor +import os +import re +import shutil +import torch +import tqdm +#from pprint import pprint + +from deepspeed.checkpoint import DeepSpeedCheckpoint +from deepspeed.checkpoint import ( + OPTIMIZER_STATE_DICT, + ZERO_STAGE, + BASE_OPTIMIZER_STATE, + SINGLE_PARTITION_OF_FP32_GROUPS, + PARAM_GROUPS, + PARAM_SLICE_MAPPINGS, + PARAM_SHAPES, + PARAM, + CAT_DIM, + PARAM_N_SUB_PARAMS, + SUB_PARAM_SHAPE, + VOCAB_TENSOR, + UNIVERSAL_CHECKPOINT_INFO, + UNIVERSAL_CHECKPOINT_VERSION_KEY, + UNIVERSAL_CHECKPOINT_VERSION_VALUE, + VOCABULARY_PARAMETER_PATTERNS, + PIPELINE_REPLICATED_PARAMETER_PATTERNS, + TP_REPLICATED_PARAMETER_PATTERNS, + PARAMETER_TO_AVERAGE_PATTERNS, + PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, + PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0, + PARAMETER_WITH_SUB_PARAMS, + SubparamShape, +) + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--input_folder', type=str, required=True, help='Input DeepSpeed Checkpoint folder') + parser.add_argument('--output_folder', type=str, required=True, help='Output DeepSpeed checkpoint folder') + parser.add_argument('--num_extract_workers', + default=4, + type=int, + help='How many parallel processes to extract zero shards') + parser.add_argument( + '--num_merge_workers', + default=2, + type=int, + help= + 'How many parallel processes to merge tp slices (more memory intensive, use much fewer than --num_extract_workers))' + ) + parser.add_argument('--keep_temp_folder', + action='store_true', + help='Preserve temporary folder of intermediate checkpoint slice files. Useful for debugging.') + parser.add_argument('--no_strict', + dest='strict', + action='store_false', + help='Do not perform validity checks on converted checkpoint.') + parser.add_argument('--inject_missing_state', + action='store_true', + help='Inject missing checkpoint state into the checkpoint if it is absent.') + args = parser.parse_args() + print(f'args = {args}') + return args + + +def atoi(text): + return int(text) if text.isdigit() else text + + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + ''' + return [atoi(c) for c in re.split(r'(\d+)', text)] + + +def _create_checkpoint_paths(base_folder, iteration, tp_degree, pp_degree): + path_list = [] + iter_folder = f'iter_{iteration:07d}' + for i in range(0, tp_degree): + path_list.append([]) + for j in range(0, pp_degree): + rank_folder = f'mp_rank_{i:02d}' if pp_degree == 1 else f'mp_rank_{i:02d}_{j:03d}' + ckpt_path = os.path.join(rank_folder, 'model_optim_rng.pt') + path_list[i].append(os.path.join(base_folder, iter_folder, ckpt_path)) + + return path_list + + +def _save_checkpoint(file_path, chkpt_sd): + dir, _ = os.path.split(file_path) + os.makedirs(dir, exist_ok=True) + torch.save(chkpt_sd, file_path) + + +def extract_zero_shards(dir, ds_checkpoint, indices_3D): + pp_index, tp_index, dp_index = indices_3D + sd = ds_checkpoint.get_zero_checkpoint_state(pp_index=pp_index, tp_index=tp_index, dp_index=dp_index) + + # pprint(f"Processing {dp_index=} {pp_index=}, {tp_index=}") + + optim_sd = sd[OPTIMIZER_STATE_DICT] + param_slice_mappings = optim_sd[PARAM_SLICE_MAPPINGS] + universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO) + pipeline_replicated_params = universal_checkpoint_info.get(PIPELINE_REPLICATED_PARAMETER_PATTERNS, []) + # print(f'{pipeline_replicated_params=}') + + # dict + state_groups = optim_sd[BASE_OPTIMIZER_STATE]["state"] + # list + fp32_groups = optim_sd[SINGLE_PARTITION_OF_FP32_GROUPS] + param_groups_cnt = len(state_groups) + + for param_group_id in range(param_groups_cnt): + + flat_state = dict( + exp_avg=state_groups[param_group_id]["exp_avg"], + exp_avg_sq=state_groups[param_group_id]["exp_avg_sq"], + fp32=fp32_groups[param_group_id], + ) + + if "step" in state_groups[param_group_id]: + flat_state["step"] = state_groups[param_group_id]["step"] + + for name, fragment_mapping in param_slice_mappings[param_group_id].items(): + if pp_index > 0 and any(re.match(pattern, name) for pattern in pipeline_replicated_params): + # Skip tied weights that are replicated in first and last pp stages + continue + + # pprint(f"dpt{dp_index}{pp_index}{tp_index} {param_group_id} {name} => {fragment_mapping.start}:{fragment_mapping.numel}") + for state_key in flat_state.keys(): + dump_param_fragment(dir, tp_index, dp_index, state_key, flat_state[state_key], name, + fragment_mapping.start, fragment_mapping.numel) + + +def extract_zero_shards_stage3(optim_files, param_shapes, dp_degree, temp_dir, dp_index): + state_dict = torch.load(optim_files[dp_index], map_location='cpu', weights_only=False) + + flat_state = dict( + exp_avg=state_dict[OPTIMIZER_STATE_DICT]['optimizer_state_dict']['state'][0]["exp_avg"], + exp_avg_sq=state_dict[OPTIMIZER_STATE_DICT]['optimizer_state_dict']['state'][0]["exp_avg_sq"], + fp32=state_dict[OPTIMIZER_STATE_DICT]['fp32_flat_groups'][0], + ) + + offset = 0 + for name, shape in param_shapes.items(): + unpartitioned_numel = shape.numel() + partitioned_numel, _ = _zero_partitioned_param_info(unpartitioned_numel, dp_degree) + padding_free_numel = min(partitioned_numel, abs(unpartitioned_numel - dp_index * partitioned_numel)) + for state_key in flat_state.keys(): + dump_param_fragment(temp_dir, 0, dp_index, state_key, flat_state[state_key], name, offset, + padding_free_numel) + offset += partitioned_numel + + +cnt = 0 + + +def dp_index_to_str(dp_index): + return f"{dp_index:0>2d}" + + +def dump_param_fragment(dir, tp_index, dp_index, state_name, state_flat_tensor, param_name, offset, numel): + + global cnt # temp hack + + param_base_path = os.path.join(dir, param_name, str(tp_index)) + os.makedirs(param_base_path, exist_ok=True) + + cnt += 1 + + path = os.path.join(param_base_path, f"{state_name}.{dp_index_to_str(dp_index)}") + + #print(f"{param_name}: {offset}: {numel} => {path}") + + # State might be a python int or a tensor + if state_name != "step" and torch.is_tensor(state_flat_tensor): + state_flat_tensor = state_flat_tensor.narrow(0, offset, numel).clone() + _save_checkpoint(path, state_flat_tensor) + + +def _merge_zero_shards(param_base_path, state, tp_degree, slice_shape=None): + slices = [] + for tp_index in range(tp_degree): + prefix_path = os.path.join(param_base_path, str(tp_index), f"{state}") + paths = glob.glob(f"{prefix_path}.*") + + if len(paths) == 0: + continue + + pattern = re.compile(f"{prefix_path}\\.([0-9]+)") + dp_indices = set() + for p in paths: + m = pattern.match(p) + if m: + dp_indices.add(int(m.group(1))) + else: + raise ValueError(f"Cannot parse dp_rank from {p}") + + paths = [f"{prefix_path}.{dp_index_to_str(dp_index)}" for dp_index in sorted(list(dp_indices))] + shards = [torch.load(p, weights_only=False) for p in paths] + + if state == "step": + assert all(v == shards[0] for v in shards), "All shards must have the same step value" + slice = shards[0] + else: + if slice_shape is None: + slice = torch.cat(shards, dim=0) + else: + slice = torch.cat(shards, dim=0).reshape(slice_shape) + + slices.append(slice) + return slices + + +def merge_tp_slices(ds_checkpoint, dir, slice_dir, tp_degree, name_and_shape): + + name, shape = name_and_shape + slice_base_path = os.path.join(slice_dir, name) + param_base_path = os.path.join(dir, name) + + universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO) + replicated_parameters = universal_checkpoint_info.get(TP_REPLICATED_PARAMETER_PATTERNS, []) + parameters_to_average = universal_checkpoint_info.get(PARAMETER_TO_AVERAGE_PATTERNS, []) + parameters_with_row_parallelism = universal_checkpoint_info.get(PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, []) + vocabulary_parameters = universal_checkpoint_info.get(VOCABULARY_PARAMETER_PATTERNS, []) + parameters_with_2_sub_params_cat_dim_0 = universal_checkpoint_info.get(PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0, []) + parameter_with_sub_params = universal_checkpoint_info.get(PARAMETER_WITH_SUB_PARAMS, []) + + unmatched_patterns = set(replicated_parameters + parameters_to_average + parameters_with_row_parallelism + + vocabulary_parameters + parameters_with_2_sub_params_cat_dim_0) + unmatched_patterns.update(chain.from_iterable(SubparamShape(**s).patterns for s in parameter_with_sub_params)) + + def get_matched_pattern(patterns_, name_): + matched_ = [pattern_ for pattern_ in patterns_ if re.match(pattern_, name_)] + assert len(matched_) <= 1, f'Got more than one matching patterns={matched_} for {name_}' + if matched_: + pattern_ = matched_[0] + unmatched_patterns.discard(pattern_) + return pattern_ + return None + + def get_matched_sub_params_pattern(name_): + for subparam_shape_dict in parameter_with_sub_params: + subparam_shape = SubparamShape(**subparam_shape_dict) + for pattern_ in subparam_shape.patterns: + if re.match(pattern_, name_): + unmatched_patterns.discard(pattern_) + return subparam_shape + return None + + matched_sub_params_shape = get_matched_sub_params_pattern(name) + + step_merged = _merge_zero_shards(slice_base_path, "step", tp_degree, shape) + if step_merged: + _save_checkpoint(os.path.join(param_base_path, f"step.pt"), step_merged[0]) + + for state in ("fp32", "exp_avg", "exp_avg_sq"): + slices = _merge_zero_shards(slice_base_path, state, tp_degree, shape) + final_path = os.path.join(param_base_path, f"{state}.pt") + + #print(f"Expected shape: {shape}") + #print(f"Fragment sizes:", list(frag.shape for frag in slices)) + ckpt_dict = {} + if get_matched_pattern(replicated_parameters, name): + if len(slices) > 1: + assert all([slices[0].equal(other_slice) for other_slice in slices[1:]]) + param = slices[0] + # print(f'replicate {name} using first slice') + elif get_matched_pattern(parameters_to_average, name): + param = sum(slices) / len(slices) + # print(f'merge {name} using average') + elif get_matched_pattern(parameters_with_2_sub_params_cat_dim_0, name): + cat_dim = 0 + chunked_slices = [torch.chunk(s, 2, dim=cat_dim) for s in slices] + merged_chunks_0 = torch.cat([s[0] for s in chunked_slices], dim=cat_dim) + merged_chunks_1 = torch.cat([s[1] for s in chunked_slices], dim=cat_dim) + param = torch.cat([merged_chunks_0, merged_chunks_1], dim=cat_dim) + ckpt_dict[CAT_DIM] = cat_dim + ckpt_dict[PARAM_N_SUB_PARAMS] = 2 + elif matched_sub_params_shape: + merged_chunks = [] + partition_dim = matched_sub_params_shape.partition_dim + + sub_dim_sizes = matched_sub_params_shape.shape[partition_dim] + if not isinstance(sub_dim_sizes, tuple): + sub_dim_sizes = (sub_dim_sizes, ) + + partition_shape = [sum(d) if isinstance(d, tuple) else d for d in matched_sub_params_shape.shape] + partition_shape = [d // tp_degree if i == partition_dim else d for i, d in enumerate(partition_shape)] + slices = [s.view(partition_shape) for s in slices] + + offset = 0 + for sub_dim_size in sub_dim_sizes: + part_sub_dim_size = sub_dim_size // tp_degree + merged_chunks.append( + torch.cat([s.narrow(partition_dim, offset, part_sub_dim_size) for s in slices], dim=partition_dim)) + offset += part_sub_dim_size + param = torch.cat(merged_chunks, dim=partition_dim) + ckpt_dict[SUB_PARAM_SHAPE] = matched_sub_params_shape + else: + cat_dim = 1 if get_matched_pattern(parameters_with_row_parallelism, name) else 0 + # print(f"merge {name} with CAT DIM: {cat_dim}") + param = torch.cat(slices, dim=cat_dim) + ckpt_dict[CAT_DIM] = cat_dim + + if get_matched_pattern(vocabulary_parameters, name): + #print(f"Before {param.shape=}") + # strip padding + original_vocab_size = universal_checkpoint_info['original_vocab_size'] + param = param[:original_vocab_size, :] + ckpt_dict[VOCAB_TENSOR] = True + #print(f"After {param.shape=}") + + #print(f"Final shape: {param.shape}") + ckpt_dict[PARAM] = param + _save_checkpoint(final_path, ckpt_dict) + + return unmatched_patterns + + +def merge_zero3_slices(dp_degree, dir, slice_dir, name): + slice_base_path = os.path.join(slice_dir, name) + param_base_path = os.path.join(dir, name) + + for state in ("fp32", "exp_avg", "exp_avg_sq"): + slices = _merge_zero_shards(slice_base_path, state, 1) + final_path = os.path.join(param_base_path, f"{state}.pt") + _save_checkpoint(final_path, slices[0]) + + +def _do_parallel_work(do_work, work_chunks, num_workers): + results = [] + if num_workers > 1: + with ProcessPoolExecutor(max_workers=num_workers) as executor: + future_list = [executor.submit(do_work, work) for work in work_chunks] + for f in tqdm.tqdm(future_list): + results.append(f.result()) + else: + # No parallel pass for unit testing + # We can't create child processes in tests + for work in tqdm.tqdm(work_chunks): + results.append(do_work(work)) + return results + + +def _extract_zero_shard_files(args, ds_checkpoint, temp_dir): + _3d_range_list = list( + itertools.product(range(ds_checkpoint.pp_degree), range(ds_checkpoint.tp_degree), + range(ds_checkpoint.dp_degree))) + #pprint(f'{_3d_range_list=}') + + do_work = partial(extract_zero_shards, temp_dir, ds_checkpoint) + _do_parallel_work(do_work, _3d_range_list, args.num_extract_workers) + + +def _extract_zero_shard_files_stage3(args, optim_files, param_shapes, dp_degree, temp_dir): + do_work = partial(extract_zero_shards_stage3, optim_files, param_shapes, dp_degree, temp_dir) + _do_parallel_work(do_work, list(range(dp_degree)), args.num_extract_workers) + + +def _merge_tp_slice_files(args, ds_checkpoint, slice_shapes, temp_dir): + zero_output_folder = os.path.join(args.output_folder, "zero") + do_work = partial(merge_tp_slices, ds_checkpoint, zero_output_folder, temp_dir, ds_checkpoint.tp_degree) + unmatched_patterns_lists = _do_parallel_work(do_work, list(slice_shapes.items()), args.num_merge_workers) + + # verify that all patterns were used + # if a pattern was not used by any of the workers, then it was not used at all -> assert/alert + sets = [set(lst) for lst in unmatched_patterns_lists] + unmatched_patterns = list(set.intersection(*sets)) + if args.strict: + assert not unmatched_patterns, f'Unused patterns={unmatched_patterns} while merging tp slices' + elif unmatched_patterns: + print(f'Warning: Unused patterns={unmatched_patterns} while merging tp slices') + + +def _merge_zero3_slice_files(args, param_shapes, dp_degree, temp_dir): + zero_output_folder = os.path.join(args.output_folder, "zero") + do_work = partial(merge_zero3_slices, dp_degree, zero_output_folder, temp_dir) + _do_parallel_work(do_work, param_shapes.keys(), args.num_merge_workers) + + +def _zero_partitioned_param_info(unpartitioned_numel, world_size): + remainder = unpartitioned_numel % world_size + padding_numel = (world_size - remainder) if remainder else 0 + partitioned_numel = math.ceil(unpartitioned_numel / world_size) + return partitioned_numel, padding_numel + + +def _parse_model_states_stage3(files): + return torch.load(files[0], map_location=torch.device('cpu'), weights_only=False)[PARAM_SHAPES] + + +def _save_optimizer_state(args, ds_checkpoint): + sharded_states = [BASE_OPTIMIZER_STATE, PARAM_SLICE_MAPPINGS, SINGLE_PARTITION_OF_FP32_GROUPS] + sd = ds_checkpoint.get_zero_checkpoint_state(pp_index=0, tp_index=0, dp_index=0) + + optim_sd = sd[OPTIMIZER_STATE_DICT] + output_sd = {k: v for k, v in optim_sd.items() if k not in sharded_states} + output_sd[PARAM_GROUPS] = optim_sd[BASE_OPTIMIZER_STATE][PARAM_GROUPS] + zero_output_folder = os.path.join(args.output_folder, "zero") + output_file_path = os.path.join(zero_output_folder, f"optimizer_state.pt") + _save_checkpoint(output_file_path, output_sd) + + +def _save_optimizer_state_stage3(args, optim_files): + sd = torch.load(optim_files[0], map_location=torch.device('cpu'), weights_only=False) + output_sd = sd[OPTIMIZER_STATE_DICT] + output_sd[PARAM_GROUPS] = output_sd[OPTIMIZER_STATE_DICT][PARAM_GROUPS] + zero_output_folder = os.path.join(args.output_folder, "zero") + output_file_path = os.path.join(zero_output_folder, f"optimizer_state.pt") + _save_checkpoint(output_file_path, output_sd) + + +def _get_optim_files(checkpoint_dir): + return _get_checkpoint_files(checkpoint_dir, "*_optim_states.pt") + + +def _get_model_state_files(checkpoint_dir): + return _get_checkpoint_files(checkpoint_dir, "*_model_states.pt") + + +def _get_checkpoint_files(checkpoint_dir, glob_pattern): + ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys) + + if len(ckpt_files) == 0: + raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'") + + return ckpt_files + + +def _get_zero_stage(optim_files): + state_dict = torch.load(optim_files[0], map_location=torch.device('cpu'), weights_only=False) + optimizer_state = state_dict[OPTIMIZER_STATE_DICT] + zero_stage = optimizer_state.get(ZERO_STAGE, 1) + return zero_stage + + +def _inject_missing_state(ds_checkpoint): + if UNIVERSAL_CHECKPOINT_INFO not in ds_checkpoint.global_state: + sd = torch.load(ds_checkpoint.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False) + if UNIVERSAL_CHECKPOINT_INFO not in sd: + ds_checkpoint.global_state[UNIVERSAL_CHECKPOINT_INFO] = {} + ds_checkpoint.global_state[UNIVERSAL_CHECKPOINT_INFO][ + UNIVERSAL_CHECKPOINT_VERSION_KEY] = UNIVERSAL_CHECKPOINT_VERSION_VALUE + + +def _check_for_required_state(ds_checkpoint): + universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO) + assert universal_checkpoint_info is not None, f'Required {UNIVERSAL_CHECKPOINT_INFO} state is missing in checkpoint. Verify that client creates this state.' + + +def main(args): + print(f'Convert DeepSpeed Checkpoint to Universal Checkpoint') + + print(f'Converting DeepSpeed checkpoint in {args.input_folder} to Universal checkpoint in {args.output_folder}') + + optim_files = _get_optim_files(args.input_folder) + zero_stage = _get_zero_stage(optim_files) + + if zero_stage <= 2: + ds_checkpoint = DeepSpeedCheckpoint(args.input_folder) + if args.inject_missing_state: + _inject_missing_state(ds_checkpoint) + else: + _check_for_required_state(ds_checkpoint) + + iteration = ds_checkpoint.get_iteration() + #_create_latest_file(args.output_folder, iteration) + checkpoint_paths = _create_checkpoint_paths(args.output_folder, iteration, ds_checkpoint.tp_degree, + ds_checkpoint.pp_degree) + + slice_shapes = [] + for mp_rank_file in ds_checkpoint.mp_rank_files: + mp_sd = torch.load(mp_rank_file, map_location=torch.device('cpu'), weights_only=False) + slice_shapes += mp_sd[PARAM_SHAPES] + + # fix back to normal flat dict, merge duplicates for tp>1 + slice_shapes = dict((k, v) for d in slice_shapes for k, v in d.items()) + temp_dir = os.path.join(args.output_folder, 'tmp') + + print('*** 1. Extracting ZeRO fragments') + _extract_zero_shard_files(args, ds_checkpoint, temp_dir) + + print('*** 2. Merging slices .....') + _merge_tp_slice_files(args, ds_checkpoint, slice_shapes, temp_dir) + + print('*** 3. Saving common optimizer states') + _save_optimizer_state(args, ds_checkpoint) + + if not args.keep_temp_folder: + shutil.rmtree(temp_dir, ignore_errors=True) + + # Copy mp* files into output folder + for f in glob.glob(os.path.join(args.input_folder, 'mp*')): + shutil.copy2(f, args.output_folder) + + else: + model_files = _get_model_state_files(args.input_folder) + param_shapes = _parse_model_states_stage3(model_files) + param_shapes = {k: v for d in param_shapes for k, v in d.items()} + dp_degree = len(model_files) + + temp_dir = os.path.join(args.output_folder, 'tmp') + + print('*** 1. Extracting ZeRO fragments') + _extract_zero_shard_files_stage3(args, optim_files, param_shapes, dp_degree, temp_dir) + + print('*** 2. Merging slices .....') + _merge_zero3_slice_files(args, param_shapes, dp_degree, temp_dir) + + print('*** 3. Saving common optimizer states') + _save_optimizer_state_stage3(args, optim_files) + + if not args.keep_temp_folder: + shutil.rmtree(temp_dir, ignore_errors=True) + + # Copy *model_states files into output folder + for f in glob.glob(os.path.join(args.input_folder, '*model_states.pt')): + shutil.copy2(f, args.output_folder) + + # Update latest to output folder + checkpoint_root_folder, step_folder = os.path.split(args.output_folder) + latest_file = os.path.join(checkpoint_root_folder, 'latest_universal') + with open(latest_file, "w") as f: + f.write(step_folder) + + print('*** Done!') + + +if __name__ == "__main__": + args = parse_arguments() + main(args) diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_3d_utils.py b/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_3d_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..02b3947624a12153aeb83967e515b18002890a40 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_3d_utils.py @@ -0,0 +1,111 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .reshape_utils import (get_files, get_files_with_prefix, partition_data, get_zero_files) + +from .constants import (MODEL_FILE_PREFIX, LAYER_FILE_PREFIX) + +from .reshape_meg_2d import (reshape_meg_2d_parallel, meg_2d_parallel_map) + +PP_DIM = 'PP' +TP_DIM = 'TP' +DP_DIM = 'DP' + + +class model_3d_desc(object): + + def __init__(self, pp_degree=1, tp_degree=1, dp_degree=1): + self.pp_degree = pp_degree + self.tp_degree = tp_degree + self.dp_degree = dp_degree + + def reshape(self, target_3d_desc, verbose=False): + valid_reshape, reshape_errors = self.can_reshape(target_3d_desc) + assert valid_reshape, ','.join(reshape_errors) + tgt_2d_map = reshape_meg_2d_parallel(old_pp_degree=self.pp_degree, + old_tp_degree=self.tp_degree, + new_pp_degree=target_3d_desc.pp_degree, + new_tp_degree=target_3d_desc.tp_degree, + verbose=verbose) + + flat_3d_map = flatten_dp_dimension(meg_2d_map=tgt_2d_map, + src_2d_size=self.pp_degree * self.tp_degree, + dp_degree=self.dp_degree) + + return unflatten_dp_dimension(meg_2d_map=flat_3d_map, dp_degree=target_3d_desc.dp_degree) + + def get_desc(self): + return f'{PP_DIM},{TP_DIM},{DP_DIM} = ({self.pp_degree}, {self.tp_degree}, {self.dp_degree})' + + def world_size(self): + return self.pp_degree * self.tp_degree * self.dp_degree + + def is_valid(self, pp_index, tp_index, dp_index): + err_msg = [] + valid = True + for index, degree, dim_name in [(pp_index, self.pp_degree, PP_DIM), (tp_index, self.tp_degree, TP_DIM), + (dp_index, self.dp_degree, DP_DIM)]: + if index >= degree: + valid = False + err_msg.append(f'{dim_name} indexing error: index {index} >= degree {degree}') + + return valid, err_msg + + def can_reshape(self, target_3d_desc): + err_msg = [] + if target_3d_desc.pp_degree > self.pp_degree: + err_msg.append( + f'Expansion reshape not supported - {PP_DIM}: {self.pp_degree} ---> {target_3d_desc.pp_degree}') + + if target_3d_desc.tp_degree > self.tp_degree: + err_msg.append( + f'Expansion reshape not supported - {TP_DIM}: {self.tp_degree} ---> {target_3d_desc.tp_degree}') + + if target_3d_desc.dp_degree > self.dp_degree: + err_msg.append( + f'Expansion reshape not supported - {DP_DIM}: {self.dp_degree} ---> {target_3d_desc.dp_degree}') + + return len(err_msg) == 0, err_msg + + +def get_model_3d_descriptor(dir): + file_list = get_files(dir) + zero_file_list = get_zero_files(dir) + num_pp0_files = len(get_files_with_prefix(file_list, f'{LAYER_FILE_PREFIX}01')) + if num_pp0_files > 0: + tp_degree = num_pp0_files + pp_degree = len(get_files_with_prefix(file_list, MODEL_FILE_PREFIX)) // tp_degree + dp_degree = max(1, len(zero_file_list) // (pp_degree * tp_degree)) + else: + tp_degree = len(get_files_with_prefix(file_list, MODEL_FILE_PREFIX)) + dp_degree = max(1, len(zero_file_list) // tp_degree) + pp_degree = 1 + + return model_3d_desc(pp_degree, tp_degree, dp_degree) + + +def flatten_dp_dimension(meg_2d_map, src_2d_size, dp_degree): + new_meg_2d_map = meg_2d_parallel_map(meg_2d_map.pp_degree, meg_2d_map.tp_degree) + for pp_index in range(meg_2d_map.pp_degree): + for tp_index in range(meg_2d_map.tp_degree): + dp0_indices = meg_2d_map.get_data(pp_index, tp_index) + for idx in dp0_indices: + dpX_indices = [idx + (i * src_2d_size) for i in range(dp_degree)] + new_meg_2d_map.add_data(pp_index, tp_index, dpX_indices) + return new_meg_2d_map + + +def unflatten_dp_dimension(meg_2d_map, dp_degree): + pp_degree = meg_2d_map.pp_degree + tp_degree = meg_2d_map.tp_degree + meg_2d_map_list = [meg_2d_parallel_map(pp_degree=pp_degree, tp_degree=tp_degree) for _ in range(dp_degree)] + for pp_index in range(pp_degree): + for tp_index in range(tp_degree): + flat_dp_indices = meg_2d_map.get_data(pp_index, tp_index) + partitioned_dp_indices = partition_data(flat_dp_indices, dp_degree) + for dp_indices, _2d_map in zip(partitioned_dp_indices, meg_2d_map_list): + _2d_map.add_data(pp_index, tp_index, dp_indices) + + return meg_2d_map_list diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_meg_2d.py b/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_meg_2d.py new file mode 100644 index 0000000000000000000000000000000000000000..3bff87f4344f4a1e74b348ae77f9b9a7d9212c6e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_meg_2d.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .reshape_utils import partition_data + + +class meg_2d_parallel_map(object): + + def __init__(self, pp_degree, tp_degree): + self.pp_degree = pp_degree + self.tp_degree = tp_degree + self.map = {} + + def simple_init(self): + self.map = { + self._make_key(i // self.tp_degree, i % self.tp_degree): [i] + for i in range(self.pp_degree * self.tp_degree) + } + + def add_data(self, pp_index, tp_index, data): + self._validate_indices(pp_index, tp_index) + assert type(data) is list + + key = self._make_key(pp_index, tp_index) + if not key in self.map.keys(): + self.map[key] = [] + self.map[key] += data + + def get_data(self, pp_index=None, tp_index=None): + self._validate_indices(pp_index, tp_index) + pp_indices = list(range(self.pp_degree)) if pp_index is None else [pp_index] + tp_indices = list(range(self.tp_degree)) if tp_index is None else [tp_index] + + result = [] + for i in pp_indices: + for j in tp_indices: + result += self.map[self._make_key(i, j)] + + return result + + def print_data(self, tag): + print(f'{tag}') + for key, value in self.map.items(): + print(f'{key} = {value}') + + def _validate_indices(self, pp_index, tp_index): + assert pp_index is None or pp_index < self.pp_degree + assert tp_index is None or tp_index < self.tp_degree + + def _make_key(self, i, j): + return f'{i},{j}' + + +def _reshape_tp_dimension(old_2d_map, new_tp_degree): + old_pp_degree = old_2d_map.pp_degree + new_2d_map = meg_2d_parallel_map(old_pp_degree, new_tp_degree) + for i in range(old_pp_degree): + ranks_for_pp_index = old_2d_map.get_data(pp_index=i, tp_index=None) + split_ranks = partition_data(ranks_for_pp_index, new_tp_degree) + for j in range(new_tp_degree): + new_2d_map.add_data(i, j, split_ranks[j]) + + return new_2d_map + + +def _reshape_pp_dimension(old_2d_map, new_pp_degree): + old_tp_degree = old_2d_map.tp_degree + new_2d_map = meg_2d_parallel_map(new_pp_degree, old_tp_degree) + for i in range(old_tp_degree): + ranks_for_tp_index = old_2d_map.get_data(pp_index=None, tp_index=i) + split_ranks = partition_data(ranks_for_tp_index, new_pp_degree) + for j in range(new_pp_degree): + new_2d_map.add_data(j, i, split_ranks[j]) + + return new_2d_map + + +def reshape_meg_2d_parallel(old_pp_degree, old_tp_degree, new_pp_degree, new_tp_degree, verbose=False): + assert new_pp_degree <= old_pp_degree + assert new_tp_degree <= old_tp_degree + + old_2d_map = meg_2d_parallel_map(old_pp_degree, old_tp_degree) + old_2d_map.simple_init() + if verbose: + old_2d_map.print_data(f'original_2d_map:') + + if old_tp_degree != new_tp_degree: + new_tp_map = _reshape_tp_dimension(old_2d_map, new_tp_degree) + else: + new_tp_map = old_2d_map + if verbose: + new_tp_map.print_data(f'after_tp_reshape:') + + if old_pp_degree != new_pp_degree: + final_map = _reshape_pp_dimension(new_tp_map, new_pp_degree) + else: + final_map = new_tp_map + + if verbose: + final_map.print_data(f'final_2d_map:') + + return final_map + + +def get_mpu_ranks(tp_size=1, pp_size=1, dp_size=1, virtual_pp_size=None): + """ + Initialize model data parallel groups. + + Arguments: + tp_size: number of GPUs used to parallelize model tensor. + pp_size: number of GPUs used to parallelize model pipeline. + dp_size: number of GPUs used to parallelize model data. + + Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we + use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize + the model pipeline. The present function will + create 8 tensor model-parallel groups, 4 pipeline model-parallel groups + and 8 data-parallel groups as: + 8 data_parallel groups: + [g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15] + 8 tensor model-parallel groups: + [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15] + 4 pipeline model-parallel groups: + [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15] + Note that for efficiency, the caller should make sure adjacent ranks + are on the same DGX box. For example if we are using 2 DGX-1 boxes + with a total of 16 GPUs, rank 0 to 7 belong to the first box and + ranks 8 to 15 belong to the second box. + """ + + world_size = tp_size * pp_size * dp_size + + print(f"\n\n*** tp={tp_size}, pp={pp_size}, dp={dp_size}, world={world_size}") + + tensor_model_parallel_size = min(tp_size, world_size) + pipeline_model_parallel_size = min(pp_size, world_size) + data_parallel_size = world_size // (tensor_model_parallel_size * pipeline_model_parallel_size) + + num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size + num_pipeline_model_parallel_groups = world_size // pipeline_model_parallel_size + num_data_parallel_groups = world_size // data_parallel_size + + # Build the data-parallel groups. + all_dp_group_ranks = [] + for i in range(pipeline_model_parallel_size): + start_rank = i * num_pipeline_model_parallel_groups + end_rank = (i + 1) * num_pipeline_model_parallel_groups + for j in range(tensor_model_parallel_size): + ranks = range(start_rank + j, end_rank, tensor_model_parallel_size) + all_dp_group_ranks.append(list(ranks)) + + print("DP", all_dp_group_ranks) + + # Build the model-parallel groups. + all_pp_group_ranks = [] + for i in range(data_parallel_size): + ranks = [data_parallel_group_ranks[i] for data_parallel_group_ranks in all_dp_group_ranks] + all_pp_group_ranks.append(list(ranks)) + + print(f"PP", all_pp_group_ranks) + + # Build the tensor model-parallel groups. + all_tp_group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) + all_tp_group_ranks.append(list(ranks)) + + print(f"TP", all_tp_group_ranks) + + return all_tp_group_ranks, all_pp_group_ranks, all_dp_group_ranks + + # # Build the pipeline model-parallel groups and embedding groups + # # (first and last rank in each pipeline model-parallel group). + # for i in range(num_pipeline_model_parallel_groups): + # ranks = range(i, world_size, + # num_pipeline_model_parallel_groups) + # print(f"EMB{i}", list(ranks)) + + +def reshape(src, tgt): + """ + reshape([tp_size_src, pp_size_src, dp_size_src], + [tp_size_tgt, pp_size_tgt, dp_size_tgt]) + """ + + print(f"\n\n*** Reshaping: {src} => {tgt}") + + tp_size_src, pp_size_src, dp_size_src = src + tp_size_tgt, pp_size_tgt, dp_size_tgt = tgt + + tp_ranks1, pp_ranks1, dp_ranks1 = get_mpu_ranks(tp_size=tp_size_src, pp_size=pp_size_src, dp_size=dp_size_src) + tp_ranks2, pp_ranks2, dp_ranks2 = get_mpu_ranks(tp_size=tp_size_tgt, pp_size=pp_size_src, dp_size=dp_size_src) + tp_ranks3, pp_ranks3, dp_ranks3 = get_mpu_ranks(tp_size=tp_size_tgt, pp_size=pp_size_tgt, dp_size=dp_size_src) + + # handle tp contraction first + print("\n*** TP contraction:") + + for i, r in enumerate(tp_ranks1): + print(f'{tp_ranks1[i]} => {tp_ranks2[i]}') + + # handle pp contraction next + + print("\n*** PP contraction:") + + for i, r in enumerate(pp_ranks1): + print(f'{pp_ranks2[i]} => {pp_ranks3[i]}') + + +# easy +#reshape([2,2,1],[1,1,1]) + +# probably need more logic to suggest how to pack +#reshape([4,4,1],[2,2,1]) + +#reshape([2,4,2], [8,32,1]) + +# get_mpu_ranks(2,2,2) +# get_mpu_ranks(4,2,1) +# get_mpu_ranks(2,4,1) +# get_mpu_ranks(1,1,8) diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_utils.py b/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..137607721ebf91804f40105daf9ea3bf5c2600aa --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/reshape_utils.py @@ -0,0 +1,113 @@ +# 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 diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/universal_checkpoint.py b/lib/python3.12/site-packages/deepspeed/checkpoint/universal_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..266d5a0635951f828f687d5eca144e375d89b8ae --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/universal_checkpoint.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import re +import torch +import types +from typing import List, Tuple, Union +from dataclasses import dataclass +from .constants import (FP32_WEIGHT_KEY, PARAM, VOCAB_TENSOR, CAT_DIM, PARAM_N_SUB_PARAMS, SUB_PARAM_SHAPE) + + +@dataclass +class SubparamShape: + patterns: List[str] + shape: Tuple[Union[Tuple[int], int]] + partition_dim: int + + +def load_hp_checkpoint_state(self, folder, tp_rank, tp_world_size): + hp_mapping = self._hp_mapping + hp_mapping.optim_fragment = {} + + hp_keys = [] + for file in os.listdir(folder): + # We expect files named something like "exp_avg.pt", "exp_avg_sq.pt", "fp32.pt" + pattern = r'(.+).pt' + match = re.search(pattern, file) + if match: + hp_keys.append(match.group(1)) + + step = None + for key in hp_keys: + ckpt_file = os.path.join(folder, f"{key}.pt") + ckpt_dict = torch.load(ckpt_file, weights_only=False) + + if key == "step": + step = ckpt_dict + continue + + full_hp_param = ckpt_dict[PARAM] + + # need to deal with slices that were averaged. + # the opposite of averaging here becomes an exact copy of the first slice + # I thought of 2 ways: + # implementation a. find a way for a client to pass a dict with patterns + # if any(re.search(pattern, folder) for pattern in WEIGHTS_TO_AVERAGE_PATTERNS): + # tp_rank = 0 + # tp_world_size = 1 + # the other approach is to assume that the saved data is correct and if full_hp_param.shape == + # self.shape that means we automatically copy? + # implementation b. + # this version requires no additional data passed from the client + # if the shapes already match it must be slices that were averaged - so we just hack around those + if full_hp_param.shape == self.shape: + tp_rank = 0 + tp_world_size = 1 + + # special case for word_embeddings weights which get padded differently depending on TP degree. + # the converter to universal currently strips the original padding completely so the saved + # weight is padding-free and we just need to add new padding depending on the target TP + # degree + is_vocab_tensor = ckpt_dict.get(VOCAB_TENSOR, False) + if is_vocab_tensor: + # In the absence of data passed from the user wrt new padded vocab specific to tp degree + # we can again derive that data by reverse engineering the target shapes like so: + padded_target_vocab_size = self.shape[0] * tp_world_size + assert padded_target_vocab_size >= full_hp_param.shape[0], \ + f'Vocab tensor padded size {padded_target_vocab_size} < loaded universal size {full_hp_param.shape[0]}' + if padded_target_vocab_size > full_hp_param.shape[0]: + padding_size = padded_target_vocab_size - full_hp_param.shape[0] + full_hp_param = torch.nn.functional.pad(full_hp_param, (0, 0, 0, padding_size), "constant", 0) + + full_param_numel = full_hp_param.numel() + tp_slice_numel = self.numel() + # if key == FP32_WEIGHT_KEY and 'word_embeddings.weight' in folder: + # print_rank_0(f'{full_hp_param[:10]=}', force=True) + + + assert full_param_numel == tp_world_size * tp_slice_numel, \ + f'Loading {ckpt_file} full param numel {full_param_numel} != tensor slice numel {tp_slice_numel} * tp_world_size {tp_world_size}' + + # print(f"{full_hp_param.shape=} {full_param_numel=} {folder=}") + # print(f"{dst_tensor.shape=} {dst_tensor.numel()=}{folder=}") + + sub_param_shape = ckpt_dict.get(SUB_PARAM_SHAPE, None) + # since when we do many to 1 on tp we cat sometimes on dim=0 and other times on dim=1 we have to do exactly the same in reverse + # special case is when a single parameter is effectively a container for multiple sub parameters + # (more details at PARAM_N_SUB_PARAMS definition) + chunk_dim = ckpt_dict.get(CAT_DIM, 0) + n_sub_params = ckpt_dict.get(PARAM_N_SUB_PARAMS, 1) + if sub_param_shape: + partition_dim = sub_param_shape.partition_dim + sub_dim_sizes = sub_param_shape.shape[partition_dim] + if not isinstance(sub_dim_sizes, tuple): + sub_dim_sizes = (sub_dim_sizes, ) + + partition_shape = [sum(d) if isinstance(d, tuple) else d for d in sub_param_shape.shape] + full_hp_param = full_hp_param.view(partition_shape) + + offset = 0 + merged_chunks = [] + for sub_dim_size in sub_dim_sizes: + sub_params_tp_slice = full_hp_param.narrow(partition_dim, + offset, sub_dim_size).chunk(tp_world_size, + dim=partition_dim)[tp_rank] + merged_chunks.append(sub_params_tp_slice) + offset += sub_dim_size + tp_hp_slice = torch.cat(merged_chunks, dim=partition_dim) + + elif n_sub_params > 1: + sub_params = full_hp_param.chunk(n_sub_params, dim=chunk_dim) + sub_params_tp_slice = [p.chunk(tp_world_size, dim=chunk_dim)[tp_rank] for p in sub_params] + tp_hp_slice = torch.cat(sub_params_tp_slice, dim=chunk_dim) + else: + # this performs the opposite of cat when merging TP slices + tp_hp_slice = full_hp_param.chunk(tp_world_size, chunk_dim)[tp_rank] + + tp_hp_slice = tp_hp_slice.flatten() + + lp_frag_address = hp_mapping.lp_fragment_address + tp_hp_fragment = tp_hp_slice.narrow(0, lp_frag_address.start, lp_frag_address.numel) + + # print(f"{key} SHAPE: {tp_hp_slice.shape=}") + # print(f"{key} SHAPE: {dst_tensor.shape=}") + # print(f"{key} SHAPE: {tp_hp_fragment.shape=}") + + if key == FP32_WEIGHT_KEY: + dst_tensor = hp_mapping.get_hp_fragment() + assert dst_tensor.numel() == lp_frag_address.numel, \ + f'Load checkpoint {key} dst numel {dst_tensor.numel()} != src numel {lp_frag_address.numel}' + dst_tensor.data.copy_(tp_hp_fragment.data) + else: + assert tp_hp_fragment.numel() == lp_frag_address.numel, \ + f'Load checkpoint {key} dst numel {tp_hp_fragment.numel()} != src numel {lp_frag_address.numel}' + + hp_mapping.optim_fragment[key] = tp_hp_fragment.clone().detach() + + return step + + +def enable_universal_checkpoint(param_list): + for param in param_list: + param.load_hp_checkpoint_state = types.MethodType(load_hp_checkpoint_state, param) diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/utils.py b/lib/python3.12/site-packages/deepspeed/checkpoint/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5964da00728e4ee0ba74a179744aeba9ba1ae329 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/utils.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch +from .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX) + + +def get_model_ckpt_name_for_rank(base_folder, mp_rank_str): + ckpt_name = os.path.join( + base_folder, + MODEL_FILE_PREFIX + mp_rank_str + MODEL_FILE_SUFFIX, + ) + return ckpt_name + + +def get_zero_ckpt_name_for_rank(base_folder, dp_rank, mp_rank): + zero_prefix = f'{ZERO_FILE_PREFIX}{dp_rank}' + mp_rank_string = f'_{MODEL_FILE_PREFIX}{mp_rank:02d}' + zero_ckpt_name = os.path.join( + base_folder, + zero_prefix + mp_rank_string + OPTIM_FILE_SUFFIX, + ) + return zero_ckpt_name + + +def get_layer_ckpt_name_for_rank(base_folder, layer_id, tp_rank): + ckpt_file = f'{layer_id}-model_{tp_rank:02d}{MODEL_FILE_SUFFIX}' + ckpt_path = os.path.join(base_folder, ckpt_file) + return ckpt_path + + +# We pass cloned tensors to torch.save() to avoid checkpoint bloat that occurs when torch.save() +# saves the underlying storage rather than the slice of the storage corresponding to individual tensors. +# This is a problem in DeepSpeed because we often allocate tensors using slices of large flattened buffers. +# Tensor cloning helps to avoid this problem because the storage of cloned tensors are closer to the true size. +# It is expected that the garbage collector will reclaim the cloned tensor storage to avoid memory bloat. +# See https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing +def clone_tensors_for_torch_save(item, device=torch.device('cpu')): + """ + Returns a copy of ``item`` with all enclosed tensors replaced by clones on a specified device. + Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. + + Parameters: + - ``item``: tensor to clone or (possibly nested) container of tensors to clone. + - ``device``: target device (defaults to 'cpu') + + Returns: + - copy of ``item`` with cloned tensors on target device + """ + if torch.is_tensor(item): + if type(device) is str: + device = torch.device(device) + if device == item.device: + return item.detach().clone() + else: + return item.detach().to(device) + elif isinstance(item, list): + return [clone_tensors_for_torch_save(v, device) for v in item] + elif isinstance(item, tuple): + return tuple([clone_tensors_for_torch_save(v, device) for v in item]) + elif isinstance(item, dict): + return type(item)({k: clone_tensors_for_torch_save(v, device) for k, v in item.items()}) + else: + return item diff --git a/lib/python3.12/site-packages/deepspeed/checkpoint/zero_checkpoint.py b/lib/python3.12/site-packages/deepspeed/checkpoint/zero_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..c85f0241005d15e002fc1a55b0be33cca4b69204 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/checkpoint/zero_checkpoint.py @@ -0,0 +1,142 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from .constants import (BASE_OPTIMIZER_STATE, GROUP_PADDINGS, OPTIMIZER_STATE_DICT, PARTITION_COUNT) + +from .reshape_utils import (basic_folder_validation, get_zero_files, merge_state) + +from .reshape_3d_utils import (model_3d_desc, get_model_3d_descriptor) + +GROUP_STATE_KEY = 'state' + + +class ZeROCheckpoint(object): + + def __init__(self, dir): + basic_folder_validation(dir) + self.dir = dir + self.file_list = get_zero_files(dir) + self.num_files = len(self.file_list) + assert self.num_files > 0, f'No ZeRO files found in {dir}' + + self.src_3d = get_model_3d_descriptor(dir) + self.target_3d = model_3d_desc(pp_degree=self.src_3d.pp_degree, + tp_degree=self.src_3d.tp_degree, + dp_degree=self.src_3d.dp_degree) + self._3d_file_map = self.src_3d.reshape(self.target_3d) + + def get_src_world_size(self): + return self.src_3d.world_size() + + def get_src_tp_degree(self): + return self.src_3d.tp_degree + + def get_src_pp_degree(self): + return self.src_3d.pp_degree + + def get_src_dp_degree(self): + return self.src_3d.dp_degree + + def get_file_indices_for_rank(self, pp_index, tp_index, dp_index): + assert dp_index < len(self._3d_file_map), f'DP index {dp_index} >= DP degree {len(self._3d_file_map)}' + dp_2d_map = self._3d_file_map[dp_index] + return dp_2d_map.get_data(pp_index, tp_index) + + def get_files_for_rank(self, pp_index, tp_index, dp_index): + file_idx_list = self.get_file_indices_for_rank(pp_index, tp_index, dp_index) + return [self.file_list[idx] for idx in file_idx_list] + + def get_state_for_rank(self, pp_index, tp_index, dp_index, keys_to_ignore=[], strip_tensor_paddings=True): + state_file_list = self.get_files_for_rank(pp_index, tp_index, dp_index) + merged_sd = None + for state_file in state_file_list: + sd = torch.load(state_file, map_location=torch.device('cpu'), weights_only=False) + for key in keys_to_ignore: + sd.pop(key, None) + + if strip_tensor_paddings: + self._strip_tensor_paddings(sd) + + if merged_sd is None: + merged_sd = sd + else: + merged_sd = merge_state(merged_sd, sd) + + self._update_partition_count(merged_sd) + if strip_tensor_paddings: + self._clear_group_paddings(merged_sd) + + return merged_sd + + def print_3d_index_map(self, tag=None): + if tag: + print(f'3D index map: {tag}') + for dp_index, _2d_map in enumerate(self._3d_file_map): + _2d_map.print_data(f'dp = {dp_index}') + + def print_3d_file_map(self, tag=None): + if tag: + print(f'3D file map: {tag}') + for dp_index, _2d_map in enumerate(self._3d_file_map): + for pp_index in _2d_map.pp_degree: + for tp_index in _2d_map.tp_degree: + file_index_list = _2d_map.get_data(pp_index, tp_index) + file_list = [self.file_list[idx] for idx in file_index_list] + print(f'{pp_index}, {tp_index}, {dp_index} => {file_list}') + + def reshape(self, target_3d_desc: model_3d_desc): + self.target_3d = target_3d_desc + self._3d_file_map = self.src_3d.reshape(self.target_3d) + + def _strip_tensor_paddings(self, sd): + param_group_states = self._get_param_group_states(sd) + if param_group_states is None: + return + + group_paddings = self._get_optimizer_state(sd, GROUP_PADDINGS) + if group_paddings is None: + return + + for key, group_state in param_group_states.items(): + if group_paddings[key] == 0: + continue + for state_name, state_value in group_state.items(): + if state_name != "step" and torch.is_tensor(state_value): + raw_length = state_value.numel() - group_paddings[key] + group_state[state_name] = torch.narrow(state_value, 0, 0, raw_length).clone() + else: + group_state[state_name] = state_value + + def _clear_group_paddings(self, sd): + group_paddings = self._get_optimizer_state(sd, GROUP_PADDINGS) + if group_paddings: + num_groups = len(group_paddings) + sd[OPTIMIZER_STATE_DICT][GROUP_PADDINGS] = [0] * num_groups + + def _get_optimizer_state(self, sd, state_key): + optimizer_state = sd.get(OPTIMIZER_STATE_DICT, None) + if optimizer_state is None: + return None + + return optimizer_state.get(state_key, None) + + def _get_param_group_states(self, sd): + optimizer_state = sd.get(OPTIMIZER_STATE_DICT, None) + if optimizer_state is None: + return None + + base_optimizer_state = optimizer_state.get(BASE_OPTIMIZER_STATE, None) + if base_optimizer_state is None: + return None + + return base_optimizer_state.get(GROUP_STATE_KEY, None) + + def _update_partition_count(self, sd): + partition_counts = self._get_optimizer_state(sd, PARTITION_COUNT) + if partition_counts: + num_groups = len(partition_counts) + sd[OPTIMIZER_STATE_DICT][PARTITION_COUNT] = [self.target_3d.dp_degree] * num_groups diff --git a/lib/python3.12/site-packages/deepspeed/comm/__init__.py b/lib/python3.12/site-packages/deepspeed/comm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2c26ef4913085748dc10d720bb723d784ee9335e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .utils import * +from .comm import * diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e0885fefcc003123eac467794c4612a2bc427ae Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/backend.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52111eaceed84faba2df731b236fb9db1ce29b4c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/ccl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/ccl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..beaa578311c04f40491af1da6dfaf5e4e53d2369 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/ccl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/comm.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/comm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d2c638f9addd0343c43fedc96c2ef683b03112e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/comm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1fa798fc3866d571e1ffe403fb8e0c131ac7104b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca3fa9c7545320c13db650e8865699da3e316629 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/reduce_op.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/reduce_op.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e99abeb931366395e7372253b9a620e7652aba0a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/reduce_op.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/torch.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/torch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea0ecb3da117c88c4f6c80fa567c44a6a2ecb86a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/torch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2cad3e0950b6014e63b7a98ac0fe11590067ab8 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/comm/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/comm/backend.py b/lib/python3.12/site-packages/deepspeed/comm/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..15dc96c40521222128310740110420ac38668819 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/backend.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +DeepSpeed Communication Backend. + +In the future, directly use NCCL/MPI/Gloo/etc without requiring torch.distributed. Simply wrap torch.distributed for now. + +# Custom DS Backends -- Direct C/Ops + - NCCL -- [EXPERIMENTAL] + - MPI -- [EXPERIMENTAL] + - RCCL -- [EXPERIMENTAL] + - GLOO -- [EXPERIMENTAL] +# DS backend wrapper for torch.distributed [DEFAULT] + - T-NCCL -- [DEFAULT] + - T-GLOO + - T-MPI + +# Backend is the base class + -- NcclBackend, MpiBackend, and TorchBackend are the main subclasses. TorchBackend is the only officially supported backend for now. +""" + + +class Backend(object): + + def __init__(self, name='backend', rank=0, size=1): + self.name = name + # The world size and rank of the world process group + self.world_group = None + self.world_size = size + self.world_rank = rank + # Single process group (pg) implementation for now but keep a list for future + self.process_groups = [] + self.initialized = False + + def is_initialized(self): + return self.initialized + + def new_group(self): + # create a new pg and add it to pg list + pass + + def init_process_group(self): + # subclasses will initialize them fully + # - initialize a default world process group and add it to pg list + self.initialized = True diff --git a/lib/python3.12/site-packages/deepspeed/comm/ccl.py b/lib/python3.12/site-packages/deepspeed/comm/ccl.py new file mode 100644 index 0000000000000000000000000000000000000000..8753cf9f70ed5106ccfea931abf4ae5c65febefb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/ccl.py @@ -0,0 +1,194 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +''' +Copyright 2021 The Microsoft DeepSpeed Team +''' + +import torch +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import NotImplementedBuilder +from .reduce_op import ReduceOp +from .torch import TorchBackend + + +def build_ccl_op(): + builder = get_accelerator().create_op_builder("CCLCommBuilder") + if builder is None or isinstance(builder, NotImplementedBuilder): + return None + ccl_cpp_module = builder.load() + print(f'DeepSpeed {builder.absolute_name()} built successfully') + return ccl_cpp_module + + +class CCLHandler(): + + def __init__(self, ccl_comm_op=None): + self.ccl_comm_op = ccl_comm_op + + def wait(self): + # backend covered it + pass + + +class CCLBackend(TorchBackend): + + def __init__(self, name='ccl', rank=-1, world_size=-1, mpu=None, timeout=None, init_method=None): + self.ccl_comm_op = build_ccl_op() + if self.ccl_comm_op is None: + # set CCLBackend to uninitialized state if CCLCommBuilder cannot be loaded + self.initialized = False + return + super(CCLBackend, self).__init__(backend='ccl', + name='torch', + rank=rank, + world_size=world_size, + timeout=timeout, + init_method=init_method) + self.name = 'ccl' + size = self.get_world_size() + rank = self.get_rank() + main_kvs = self.ccl_comm_op.get_kvs_addr(rank) + main_kvs = torch.tensor(main_kvs).to(torch.uint8).to(get_accelerator().current_device_name()) + super(CCLBackend, self).broadcast(main_kvs, 0) + self.ccl_comm_op.initialize(size, rank, main_kvs) + self.initialized = True + self.groups = [tuple(range(self.get_world_size()))] + self.available_coll = self.ccl_comm_op.get_available_coll() + + def is_initialized(self): + return self.initialized + + def run_collective(self, name, **kwargs): + if name in self.available_coll: + if 'group' in kwargs: + kwargs['group'] = self.get_all_ranks_from_group(kwargs['group']) + if 'dst' in kwargs: + kwargs['dst'] = kwargs['group'].index(kwargs['dst']) + if 'src' in kwargs: + kwargs['src'] = kwargs['group'].index(kwargs['src']) + func = "self.ccl_comm_op." + name + eval(func)(*(kwargs.values())) + return CCLHandler(self.ccl_comm_op) + else: + func = "super(CCLBackend, self)." + name + eval(func)(*(kwargs.values())) + return CCLHandler(self.ccl_comm_op) + + def all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=False): + use_caching = False + if use_caching: + match_id = f"{tensor.size()}-{op}" + name = "all_reduce_caching" + if name in self.available_coll: + group = self.get_all_ranks_from_group(group) + return self.ccl_comm_op.all_reduce_caching(tensor, op, match_id, group, async_op) + else: + return self.run_collective(name=name, + tensor=tensor, + op=op, + match_id=match_id, + group=group, + async_op=async_op) + else: + name = "all_reduce" + if name in self.available_coll: + group = self.get_all_ranks_from_group(group) + return self.ccl_comm_op.all_reduce(tensor, op, group, async_op) + else: + return self.run_collective(name=name, tensor=tensor, op=op, group=group, async_op=async_op) + + def inference_all_reduce(self, tensor, op=ReduceOp.SUM, group=None): + name = "inference_all_reduce" + if name in self.available_coll: + return self.ccl_comm_op.inference_all_reduce(tensor, op) + else: + return self.run_collective(name=name, tensor=tensor, op=op, group=None, async_op=False) + + def broadcast(self, tensor, src, group=None, async_op=False): + return self.run_collective(name="broadcast", tensor=tensor, src=src, group=group, async_op=async_op) + + def all_gather(self, tensor_list, tensor, group=None, async_op=False): + return self.run_collective(name="all_gather", + tensor_list=tensor_list, + tensor=tensor, + group=group, + async_op=async_op) + + def reduce_scatter_tensor(self, output_tensor, input_tensor, op, group=None, async_op=False): + return self.run_collective(name="reduce_scatter_tensor", + output_tensor=output_tensor, + input_tensor=input_tensor, + op=op, + group=group) + + def all_gather_into_tensor(self, output_tensor, input_tensor, group=None, async_op=False): + return self.run_collective(name="all_gather_into_tensor", + output_tensor=output_tensor, + input_tensor=input_tensor, + group=group) + + def all_to_all_single(self, output, input, output_split_sizes, input_split_sizes, group=None, async_op=False): + return self.run_collective(name="all_to_all_single", + output=output, + input=input, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group) + + def send(self, tensor, dst, group=None, tag=0): + return self.run_collective(name="send", tensor=tensor, dst=dst, group=group, tag=tag) + + def recv(self, tensor, src, group=None, tag=0): + return self.run_collective(name="recv", tensor=tensor, src=src, group=group, tag=tag) + + def gather(self, tensor, gather_list, dst, group=None, async_op=False): + return self.run_collective(name="gather", tensor=tensor, gather_list=gather_list, dst=dst, group=group) + + def scatter(self, tensor, gather_list, dst, group=None, async_op=False): + return self.run_collective(name="scatter", tensor=tensor, gather_list=gather_list, dst=dst, group=group) + + def barrier(self, group=None, async_op=False): + return self.run_collective(name="barrier", group=group, async_op=async_op) + + def monitored_barrier(self, group=None, timeout=None, wait_all_ranks=False): + return self.run_collective(name="monitored_barrier", group=group) + + def reduce_scatter(self, output, input_list, op=ReduceOp.SUM, group=None, async_op=False): + return self.run_collective(name="reduce_scatter", + output=output, + input_list=input_list, + op=op, + group=group, + async_op=async_op) + + def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False): + return self.run_collective(name="reduce", tensor=tensor, dst=dst, op=op, group=group, async_op=async_op) + + def new_group(self, ranks): + return super(CCLBackend, self).new_group(ranks) + + def _new_group(self, ranks, group): + size = len(ranks) + rank = self.get_rank() + sub_main_kvs = self.ccl_comm_op.get_sub_kvs_addr(rank == ranks[0]) + sub_main_kvs = torch.tensor(sub_main_kvs).to(torch.uint8).to(get_accelerator().current_device_name()) + super(CCLBackend, self).broadcast(sub_main_kvs, ranks[0], group) + self.ccl_comm_op.initialize_sub_comm(size, ranks.index(rank), sub_main_kvs, ranks) + self.groups.append(tuple(ranks)) + + def get_all_ranks_from_group(self, group): + if group is None: + return list(range(self.get_world_size())) + rank = 0 + results = [] + try: + while True: + results.append(super(CCLBackend, self).get_global_rank(group, rank)) + rank += 1 + except (ValueError, RuntimeError): + pass + if tuple(results) not in self.groups: + self._new_group(results, group) + return results diff --git a/lib/python3.12/site-packages/deepspeed/comm/comm.py b/lib/python3.12/site-packages/deepspeed/comm/comm.py new file mode 100644 index 0000000000000000000000000000000000000000..d0e5bf225b9bc291c610afb0ebd29a34a770df87 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/comm.py @@ -0,0 +1,814 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" + DeepSpeed Communication Package: deepspeed.comm + deepspeed.comm + -- import and use deepspeed.ops.comm + -- use torch.distributed directly if both this package and torch.distributed use the same NCCL version + -- use custom collectives + -- can either use torch.dist or ds.ops.comm? + Note: the old 1-bit compressed allreduce variants that resided in deepspeed.runtime.comm will be moved here as well. + deepspeed.comm API + -- must be kept fully compatible (same signatures) as torch.dist API to ensure backward/cross-framework compatibility. + -- e.g. if a client code used + from deepspeed import comm as dist + instead of + import torch.distributed as dist + The code should work without breaking any of the public torch.distributed functionality + Future: + -- deepspeed groups API should be brought into ds.comm +""" + +import torch +import os + +from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT, default_pg_timeout +from .constants import * +from .reduce_op import * +from deepspeed.accelerator import get_accelerator +from deepspeed.comm.ccl import CCLBackend + +from deepspeed.utils.comms_logging import CommsLogger +from deepspeed.utils import timer, get_caller_func +from deepspeed.comm.torch import TorchBackend + +from deepspeed import utils +from datetime import timedelta + +# Current deepspeed.comm backend (cdb) global object for simple access by client code +cdb = None + +# Create global timer for ops +timers = timer.SynchronizedWallClockTimer() +timer_summary = {} + +comms_logger = CommsLogger() + +# Maintain objects of all initialized ds backends and assign them using the API functions in this file +nccl_backend = None +mpi_backend = None +ccl_backend = None +hccl_backend = None + +# This should be set here so all rank/size information from the launcher can be propagated +from deepspeed.comm.utils import * + + +class ProcessGroup(): + + def __init__(self, comm_id, ranks=[]): + self.ranks = ranks + self.comm_id = comm_id + self.size = len(ranks) + + +def _configure_using_config_file(config): + if config.comms_logger_enabled: + comms_logger.configure(config) + + +def configure( + deepspeed_config=None, + enabled=None, + prof_all=None, + prof_ops=None, + verbose=None, + debug=None, +): + + if deepspeed_config is not None: + _configure_using_config_file(deepspeed_config.comms_config) + + if enabled is not None: + comms_logger.enabled = enabled + + if prof_all is not None: + comms_logger.prof_all = prof_all + + if prof_ops is not None: + comms_logger.prof_ops = prof_ops + + if verbose is not None: + comms_logger.verbose = verbose + + if debug is not None: + comms_logger.debug = debug + + +# Logging wrapper for timing ops +def timed_op(func): + + def log_wrapper(*args, **kwargs): + # Add enabled flag so that overhead to each comm op is two if conditions at most + if comms_logger.enabled: + if ('prof' in kwargs + and kwargs['prof']) or comms_logger.prof_all or ('log_name' in kwargs + and kwargs['log_name'] in comms_logger.prof_ops): + # Need func args for their defaults + func_args = get_default_args(func) + func_args.update(kwargs) + msg_size = get_msg_size_from_args(func, *args, **kwargs) + log_name = get_debug_log_name(func_args, comms_logger.debug) + timers(log_name).start() + # Return the op, then stop the op's timer + try: + return func(*args, **kwargs) + finally: + if comms_logger.enabled: + # Need to make op blocking for accurate logging + get_accelerator().synchronize() + # If we're using MPI, we can't simply sync the stream + if cdb.using_mpi: + cdb.barrier() + if ('prof' in kwargs and kwargs['prof']) or comms_logger.prof_all or ( + 'log_name' in kwargs and kwargs['log_name'] in comms_logger.prof_ops): + log_name = get_debug_log_name(func_args, comms_logger.debug) + raw_name = func.__name__ + timers(log_name).stop() + # need temp var since 'elapsed' resets events + time_elapsed = timers(log_name).elapsed(reset=False) + comms_logger.append(raw_name, log_name, time_elapsed, msg_size) + + return log_wrapper + + +# For compatibility with torch distributed's init_process_group, we shall retain the signature from PyTorch code. +# DeepSpeed NCCL/MPI backend may not need all these params as we will have our own implementation. +# Please read full torch.distributed API docs from https://pytorch.org/docs/stable/distributed.html + + +# UNUSED: Future helper function to initialize DS backends +def init_deepspeed_backend(ds_backend, timeout, init_method): + global cdb + global nccl_backend + global mpi_backend + global ccl_backend + global hccl_backend + + rank = int(os.getenv('RANK', '-1')) + size = int(os.getenv('WORLD_SIZE', '-1')) + + if ds_backend == NCCL_BACKEND: + utils.logger.debug("NCCL backend in DeepSpeed not yet implemented") + elif ds_backend == MPI_BACKEND: + utils.logger.debug("MPI backend in DeepSpeed not yet implemented") + elif ds_backend == GLOO_BACKEND: + utils.logger.debug("Gloo backend in DeepSpeed not yet implemented") + elif ds_backend == CCL_BACKEND: + ccl_backend = CCLBackend(rank=rank, world_size=size, timeout=timeout, init_method=init_method) + utils.logger.info(f"Initialize {ds_backend} backend") + elif ds_backend == HCCL_BACKEND: + utils.logger.debug("HCCL backend in DeepSpeed not yet implemented") + else: + utils.logger.debug(f"DeepSpeed does not support {ds_backend} backend") + + +def is_initialized(): + #assert cdb is not None, 'DeepSpeed backend not set, please initialize it using init_process_group()' + if cdb is None: + return False + else: + return cdb.is_initialized() + + +def destroy_process_group(group=None): + global cdb + return cdb.destroy_process_group(group=group) + + +def new_group(ranks): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.new_group(ranks) + + +def is_available() -> bool: + + # Returns ``True`` if the deepspeed comm package is available. + + # TODO: load other ops. Clients including deepspeed itself should use deepspeed.comm to import + # any communication related primitives from this package. + # use hasattr(deepspeed.csrc.ops, "_comm") or something + return True + + +def set_backend(): + global cdb + global nccl_backend + global mpi_backend + global ccl_backend + global hccl_backend + + backend_name = get_accelerator().communication_backend_name() + + if backend_name == NCCL_BACKEND: + if nccl_backend is not None and nccl_backend.is_initialized(): + cdb = nccl_backend + elif backend_name == MPI_BACKEND: + if mpi_backend is not None and mpi_backend.is_initialized(): + cdb = mpi_backend + elif backend_name == CCL_BACKEND: + if ccl_backend is not None and ccl_backend.is_initialized(): + cdb = ccl_backend + elif backend_name == HCCL_BACKEND: + if hccl_backend is not None and hccl_backend.is_initialized(): + cdb = hccl_backend + + +@timed_op +def broadcast(tensor, src, group=None, async_op=False, prof=False, log_name='broadcast', debug=get_caller_func()): + global cdb + return cdb.broadcast(tensor=tensor, src=src, group=group, async_op=async_op) + + +@timed_op +def broadcast_object_list(object_list, src, group=None, device=None): + global cdb + return cdb.broadcast_object_list(object_list=object_list, src=src, group=group, device=device) + + +@timed_op +def all_gather(tensor_list, + tensor, + group=None, + async_op=False, + prof=False, + log_name='all_gather', + debug=get_caller_func()): + global cdb + return cdb.all_gather(tensor_list=tensor_list, tensor=tensor, group=group, async_op=async_op) + + +@timed_op +def all_gather_object(object_list, obj, group=None, prof=False, log_name='all_gather_object', debug=get_caller_func()): + global cdb + return cdb.all_gather_object(object_list=object_list, obj=obj, group=group) + + +def has_reduce_scatter_tensor(): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.has_reduce_scatter_tensor() + + +def reduce_scatter_fn(output_tensor, + tensor, + op=ReduceOp.SUM, + group=None, + async_op=False, + prof=False, + debug=get_caller_func()): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + if cdb.has_reduce_scatter_tensor(): + return reduce_scatter_tensor(output_tensor, + tensor, + op=op, + group=group, + async_op=async_op, + prof=prof, + debug=debug) + else: + if get_rank() == 0: + utils.logger.warning_once("unable to find torch.distributed.reduce_scatter_tensor. will fall back to " + "torch.distributed.reduce_scatter which will result in suboptimal performance. " + "please consider upgrading your pytorch installation.") + input_tensor_lst = list(torch.chunk(tensor, cdb.get_world_size(group))) + return reduce_scatter(output_tensor, + input_tensor_lst, + op=op, + group=group, + async_op=async_op, + prof=prof, + debug=debug) + + +@timed_op +def reduce_scatter_tensor(output_tensor, + tensor, + op=ReduceOp.SUM, + group=None, + async_op=False, + prof=False, + log_name='reduce_scatter_tensor', + debug=get_caller_func()): + global cdb + return cdb.reduce_scatter_tensor(output_tensor=output_tensor, + input_tensor=tensor, + op=op, + group=group, + async_op=async_op) + + +@timed_op +def all_gather_into_tensor(output_tensor, + tensor, + group=None, + async_op=False, + prof=False, + log_name='all_gather_into_tensor', + debug=get_caller_func()): + global cdb + return cdb.all_gather_into_tensor(output_tensor=output_tensor, input_tensor=tensor, group=group, async_op=async_op) + + +def has_all_gather_into_tensor(): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.has_all_gather_into_tensor() + + +def allgather_fn(output_tensor, input_tensor, group=None, async_op=False, debug=get_caller_func()): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + if cdb.has_all_gather_into_tensor(): + return all_gather_into_tensor(output_tensor, input_tensor, group=group, async_op=async_op, debug=debug) + else: + if get_rank() == 0: + utils.logger.warning_once("unable to find torch.distributed.all_gather_into_tensor. will fall back to " + "torch.distributed.all_gather which will result in suboptimal performance. " + "please consider upgrading your pytorch installation.") + output_tensors = list(torch.chunk(output_tensor, cdb.get_world_size(group))) + return all_gather(output_tensors, input_tensor, group=group, async_op=async_op, debug=debug) + + +@timed_op +def all_to_all_single(output, + tensor, + output_split_sizes=None, + input_split_sizes=None, + group=None, + async_op=False, + prof=False, + log_name='all_to_all_single', + debug=get_caller_func()): + global cdb + return cdb.all_to_all_single(output=output, + input=tensor, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op) + + +@timed_op +def all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False): + global cdb + return cdb.all_to_all(output_tensor_list, input_tensor_list, group=group, async_op=async_op) + + +@timed_op +def send(tensor, dst, group=None, tag=0, prof=False, log_name='send', debug=get_caller_func()): + global cdb + return cdb.send(tensor=tensor, dst=dst, group=group, tag=tag) + + +@timed_op +def recv(tensor, src=None, group=None, tag=0, prof=False, log_name='recv', debug=get_caller_func()): + global cdb + return cdb.recv(tensor=tensor, src=src, group=group, tag=tag) + + +@timed_op +def isend(tensor, dst, group=None, tag=0, prof=False, log_name='isend', debug=get_caller_func()): + global cdb + return cdb.send(tensor=tensor, dst=dst, group=group, tag=tag) + + +@timed_op +def irecv(tensor, src=None, group=None, tag=0, prof=False, log_name='irecv', debug=get_caller_func()): + global cdb + return cdb.recv(tensor=tensor, src=src, group=group, tag=tag) + + +@timed_op +def gather(tensor, + gather_list=None, + dst=0, + group=None, + async_op=False, + prof=False, + log_name='gather', + debug=get_caller_func()): + global cdb + return cdb.gather(tensor=tensor, gather_list=gather_list, dst=dst, group=group, async_op=async_op) + + +@timed_op +def scatter(tensor, + scatter_list=None, + src=0, + group=None, + async_op=False, + prof=False, + log_name='scatter', + debug=get_caller_func()): + global cdb + return cdb.scatter(tensor=tensor, scatter_list=scatter_list, src=src, group=group, async_op=async_op) + + +@timed_op +def barrier(group=None, async_op=False, device_ids=None, prof=False, log_name='barrier', debug=get_caller_func()): + global cdb + return cdb.barrier(group=group, async_op=async_op) + + +@timed_op +def monitored_barrier(group=None, + timeout=None, + wait_all_ranks=False, + prof=False, + log_name='monitored_barrier', + debug=get_caller_func()): + global cdb + return cdb.monitored_barrier(group=group, timeout=timeout, wait_all_ranks=wait_all_ranks) + + +def log_summary(show_straggler=False): + global cdb + barrier(log_name='log_summary_barrier') + if cdb.get_rank() == 0: + comms_logger.log_all(print_log=True, show_straggler=show_straggler) + else: + comms_logger.log_all(print_log=False, show_straggler=show_straggler) + barrier(log_name='log_summary_barrier') + + +@timed_op +def reduce(tensor, + dst, + op=ReduceOp.SUM, + group=None, + async_op=False, + prof=False, + log_name='reduce', + debug=get_caller_func()): + global cdb + return cdb.reduce(tensor=tensor, dst=dst, op=op, group=group, async_op=async_op) + + +@timed_op +def reduce_scatter(output, + input_list, + op=ReduceOp.SUM, + group=None, + async_op=False, + prof=False, + log_name='reduce_scatter', + debug=get_caller_func()): + global cdb + return cdb.reduce_scatter(output=output, input_list=input_list, op=op, group=group, async_op=async_op) + + +def has_all_reduce_coalesced(): + """""" + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + assert cdb.has_all_reduce_coalesced is not None, 'has_all_reduce_coalesced is not yet defined' + return cdb.has_all_reduce_coalesced + + +def has_coalescing_manager(): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + assert cdb.has_coalescing_manager is not None, 'has_coalescing_manager is not yet defined' + return cdb.has_coalescing_manager + + +def all_gather_coalesced(output_tensors, input_tensors, group=None, async_op=False): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.all_gather_coalesced(output_tensors, input_tensors, group=group, async_op=async_op) + + +@timed_op +def all_reduce(tensor, + op=ReduceOp.SUM, + group=None, + async_op=False, + prof=False, + log_name='all_reduce', + debug=get_caller_func()): + #if profile_comm: + # context of the timers? + # timers.start() + # TensorBoard logging for comm calls.? + global cdb + #print(f'op = {op}, cdb= {cdb.name}') + return cdb.all_reduce(tensor, op, group, async_op) + + +@timed_op +def inference_all_reduce(tensor, + op=ReduceOp.SUM, + group=None, + async_op=False, + prof=False, + log_name='all_reduce', + debug=get_caller_func()): + global cdb + return cdb.inference_all_reduce(tensor, op, group) + + +@timed_op +def all_reduce_coalesced(tensors, + op=ReduceOp.SUM, + group=None, + async_op=False, + prof=False, + log_name='all_reduce', + debug=get_caller_func()): + global cdb + return cdb.all_reduce_coalesced(tensors, op, group, async_op) + + +def get_world_group(): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.get_world_group() + + +def get_world_size(group=None) -> int: + """ + Returns the number of processes in the current process group + Args: + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + Returns: + The world size of the process group + -1, if not part of the group + """ + global cdb + + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.get_world_size(group) + + +def get_rank(group=None): + """ + Returns the rank of the current process in the provided ``group`` or the + default group if none was provided. + Rank is a unique identifier assigned to each process within a distributed + process group. They are always consecutive integers ranging from 0 to + ``world_size``. + Args: + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + Returns: + The rank of the process group + -1, if not part of the group + """ + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.get_rank(group) + + +def get_local_rank(): + """ + Helper function to get local rank after a backend has been set and initialized + Args: + None + Returns: + local rank (= GPU device ID) + """ + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return get_local_rank_from_launcher() + + +def get_global_rank(group=None, group_rank=0): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + return cdb.get_global_rank(group, group_rank) + + +def get_all_ranks_from_group(group=None): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + rank = 0 + group_ranks = [] + try: + while True: + group_ranks.append(cdb.get_global_rank(group, rank)) + rank += 1 + except (RuntimeError, ValueError): + pass + return group_ranks + + +def initialize_mesh_device(mesh_shape, mesh_dim_names): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + mesh_device = None + if hasattr(cdb, 'init_device_mesh'): + utils.logger.info(f"Initializing mesh device with backend {cdb.name} \ + with shape {mesh_shape} and dim names {mesh_dim_names}") + mesh_device = cdb.init_device_mesh(mesh_shape, mesh_dim_names) + else: + if get_rank() == 0: + utils.logger.warning_once(f"Backend {cdb.name} does not support mesh device initialization") + return mesh_device + + +def enable_symm_mem_for_group(group_name: str): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + + if hasattr(cdb, 'enable_symm_mem_for_group'): + cdb.enable_symm_mem_for_group(group_name) + else: + raise RuntimeError(f"Backend {cdb.name} does not support symmetric memory initialization") + + +# Main DeepSpeed Comms. public API. +def init_distributed(dist_backend=None, + auto_mpi_discovery=True, + distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, + verbose=True, + timeout=default_pg_timeout, + init_method=None, + dist_init_required=None, + config=None, + rank=-1, + world_size=-1): + ''' Initialize dist backend, potentially performing MPI discovery if needed + + Arguments: + dist_backend: Optional (str). torch distributed backend, e.g., nccl, mpi, gloo, hccl + auto_mpi_discovery Optional (bool). if distributed environment variables are not set, attempt to discover them from MPI + distributed_port: Optional (int). torch distributed backend port + verbose: Optional (bool). verbose logging + timeout: Optional (timedelta). Timeout for operations executed against the process group. The default value of 30 minutes can be overridden by the environment variable `DEEPSPEED_TIMEOUT`. + init_method: Optional (string). Torch distributed, URL specifying how to initialize the process group. Default is “env://” if no init_method or store is specified. + config: Optional (dict). DeepSpeed configuration for setting up comms options (e.g. Comms profiling) + rank: Optional (int). The current manually specified rank. Some init_method like “tcp://” need the rank and world_size as well (see: https://pytorch.org/docs/stable/distributed.html#tcp-initialization) + world_size: Optional (int). Desired world_size for the TCP or Shared file-system initialization. + ''' + global cdb + + configure(deepspeed_config=config) + + if dist_init_required is None: + dist_init_required = cdb is None or not cdb.is_initialized() + + if cdb is None: + init_deepspeed_backend(get_accelerator().communication_backend_name(), timeout, init_method) + set_backend() + utils.logger.info(f'cdb={cdb}') + if cdb is None and torch.distributed.is_initialized(): + # The user initialized torch.dist themselves, create cdb and short-circuit + cdb = TorchBackend(dist_backend, timeout, init_method) + return + + if dist_init_required is False: + assert ( + cdb is not None and cdb.is_initialized() is True + ), "Distributed backend is not initialized. Please set dist_init_required to True or initialize before calling deepspeed.initialize()" + else: + # Initialize torch distributed if needed + required_env = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"] + if auto_mpi_discovery and not all(map(lambda v: v in os.environ, required_env)): + if verbose: + utils.logger.info("Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...") + if in_aml() and not in_dlts(): + patch_aml_env_for_torch_nccl_backend(verbose=verbose) + elif in_aws_sm(): + patch_aws_sm_env_for_torch_nccl_backend(verbose=verbose) + else: + mpi_discovery(distributed_port=distributed_port, verbose=verbose) + + if cdb is not None and cdb.is_initialized(): + if int(os.getenv('RANK', '0')) == 0: + utils.logger.info('Distributed backend already initialized') + else: + assert isinstance(timeout, timedelta) + if dist_backend is None: + dist_backend = get_accelerator().communication_backend_name() + if int(os.getenv('RANK', '0')) == 0: + utils.logger.info('Initializing TorchBackend in DeepSpeed with backend {}'.format(dist_backend)) + # Create a torch backend object, initialize torch distributed, and assign to cdb + cdb = TorchBackend(dist_backend, timeout, init_method, rank, world_size) + + +def mpi_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True): + ''' + Discovery MPI environment via mpi4py and map to relevant dist state + ''' + from mpi4py import MPI + import subprocess + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + master_addr = None + if rank == 0: + import shlex + try: + hostname_cmd = shlex.split("hostname -I") + result = subprocess.check_output(hostname_cmd) + master_addr = result.decode('utf-8').split()[0] + except subprocess.CalledProcessError: # hostname -I not available (e.g. on macOS) + import socket + master_addr = socket.gethostbyname(socket.gethostname()) + master_addr = comm.bcast(master_addr, root=0) + + # Determine local rank by assuming hostnames are unique + proc_name = MPI.Get_processor_name() + all_procs = comm.allgather(proc_name) + local_rank = sum([i == proc_name for i in all_procs[:rank]]) + + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = str(world_size) + os.environ['LOCAL_RANK'] = str(local_rank) + os.environ['MASTER_ADDR'] = master_addr + os.environ['MASTER_PORT'] = str(distributed_port) + + if verbose: + utils.logger.info( + "Discovered MPI settings of world_rank={}, local_rank={}, world_size={}, master_addr={}, master_port={}". + format(os.environ['RANK'], os.environ['LOCAL_RANK'], os.environ['WORLD_SIZE'], os.environ['MASTER_ADDR'], + os.environ['MASTER_PORT'])) + + if cdb is not None and cdb.is_initialized(): + assert cdb.get_rank() == rank, "MPI rank {} does not match torch rank {}".format(rank, cdb.get_rank()) + assert cdb.get_world_size() == world_size, "MPI world size {} does not match torch world size {}".format( + world_size, cdb.get_world_size()) + + +def in_aml(): + # Are we running inside an Azure Machine Learning (AML) environment? + return 'AZUREML_EXPERIMENT_ID' in os.environ + + +def in_aws_sm(): + # Are we running inside an AWS SageMaker environment? + return 'SM_TRAINING_ENV' in os.environ + + +def in_dlts(): + # Are we running on a DLTS cluster? + return 'DLTS_JOB_ID' in os.environ + + +def patch_aml_env_for_torch_nccl_backend(master_port=6105, verbose=True): + """Helper routine to get and set environment variables. + This is adapted from Azure ML's documentation available from: + https://azure.github.io/azureml-web/docs/cheatsheet/distributed-training/#environment-variables-from-openmpi + """ + os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] + os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"] + single_node = int(os.environ["OMPI_COMM_WORLD_LOCAL_SIZE"]) == int(os.environ["WORLD_SIZE"]) + + if not single_node: + master_node_params = os.environ["AZ_BATCH_MASTER_NODE"].split(":") + os.environ["MASTER_ADDR"] = master_node_params[0] + # Do not overwrite master port with that defined in AZ_BATCH_MASTER_NODE + if "MASTER_PORT" not in os.environ: + os.environ["MASTER_PORT"] = str(master_port) + else: + os.environ["MASTER_ADDR"] = os.environ["AZ_BATCHAI_MPI_MASTER_NODE"] + os.environ["MASTER_PORT"] = DEFAULT_AML_MASTER_PORT + + if verbose: + utils.logger.info("NCCL_SOCKET_IFNAME original value = {}".format(os.environ["NCCL_SOCKET_IFNAME"])) + + os.environ["NCCL_SOCKET_IFNAME"] = DEFAULT_AML_NCCL_SOCKET_IFNAME + os.environ['LOCAL_RANK'] = os.environ["OMPI_COMM_WORLD_LOCAL_RANK"] + + if verbose: + utils.logger.info( + "Discovered AzureML settings of world_rank={}, local_rank={}, world_size={}, master_addr={}, master_port={}" + .format(os.environ['RANK'], os.environ['LOCAL_RANK'], os.environ['WORLD_SIZE'], os.environ['MASTER_ADDR'], + os.environ['MASTER_PORT'])) + + +def patch_aws_sm_env_for_torch_nccl_backend(verbose=True): + """Helper routine to get and set environment variables when running inside an AWS SageMaker environment. + """ + os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] + os.environ['LOCAL_RANK'] = os.environ["OMPI_COMM_WORLD_LOCAL_RANK"] + os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"] + + if verbose: + utils.logger.info( + "Discovered AWS SageMaker settings of world_rank={}, local_rank={}, world_size={}, master_addr={}, master_port={}" + .format(os.environ['RANK'], os.environ['LOCAL_RANK'], os.environ['WORLD_SIZE'], os.environ['MASTER_ADDR'], + os.environ['MASTER_PORT'])) diff --git a/lib/python3.12/site-packages/deepspeed/comm/config.py b/lib/python3.12/site-packages/deepspeed/comm/config.py new file mode 100644 index 0000000000000000000000000000000000000000..57501c9dd237c5ebae50e43c627891568f1cb6b5 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/config.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import DeepSpeedConfigModel + +from .constants import * + + +class CommsLoggerConfig(DeepSpeedConfigModel): + enabled: bool = COMMS_LOGGER_ENABLED_DEFAULT + prof_all: bool = COMMS_LOGGER_PROF_ALL_DEFAULT + prof_ops: list = COMMS_LOGGER_PROF_OPS_DEFAULT + verbose: bool = COMMS_LOGGER_VERBOSE_DEFAULT + debug: bool = COMMS_LOGGER_DEBUG_DEFAULT + + +class DeepSpeedCommsConfig: + + def __init__(self, ds_config): + self.comms_logger_enabled = 'comms_logger' in ds_config + + if self.comms_logger_enabled: + self.comms_logger = CommsLoggerConfig(**ds_config['comms_logger']) diff --git a/lib/python3.12/site-packages/deepspeed/comm/constants.py b/lib/python3.12/site-packages/deepspeed/comm/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..50d234c93fa0df954d5aec18cda209aa3909a32e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/constants.py @@ -0,0 +1,51 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +NCCL_BACKEND = 'nccl' +CCL_BACKEND = 'ccl' +MPI_BACKEND = 'mpi' +GLOO_BACKEND = 'gloo' +SCCL_BACKEND = 'sccl' +HCCL_BACKEND = 'hccl' + +DEFAULT_AML_MASTER_PORT = "54965" +DEFAULT_AML_NCCL_SOCKET_IFNAME = "^docker0,lo" + +######################################### +# Comms Logger +######################################### +# Comms Logger. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +COMMS_LOGGER_FORMAT = ''' +The Comms Logger can be specified as: +"comms_logger": { + "enabled": true, + "verbose": false, + "prof_all": true, + "debug": false, + "prof_ops": ["all_reduce", "custom_all_reduce_name"] +} +''' +COMMS_LOGGER = "comms_logger" + +# Comms logger enable signal +COMMS_LOGGER_ENABLED = "enabled" +COMMS_LOGGER_ENABLED_DEFAULT = False + +# Comms logger verbose signal +COMMS_LOGGER_VERBOSE = "verbose" +COMMS_LOGGER_VERBOSE_DEFAULT = False + +# comms logger profile all ops signal +COMMS_LOGGER_PROF_ALL = "prof_all" +COMMS_LOGGER_PROF_ALL_DEFAULT = True + +# comms logger show all ops signal +COMMS_LOGGER_DEBUG = "debug" +COMMS_LOGGER_DEBUG_DEFAULT = False + +# comms logger profile specific ops in list +COMMS_LOGGER_PROF_OPS = "prof_ops" +COMMS_LOGGER_PROF_OPS_DEFAULT = [] diff --git a/lib/python3.12/site-packages/deepspeed/comm/reduce_op.py b/lib/python3.12/site-packages/deepspeed/comm/reduce_op.py new file mode 100644 index 0000000000000000000000000000000000000000..49393f0c7227d667508359e227c6761dd6b35505 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/reduce_op.py @@ -0,0 +1,18 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import Enum + + +class ReduceOp(Enum): + SUM = 0 + PRODUCT = 1 + MIN = 2 + MAX = 3 + BAND = 4 + BOR = 5 + BXOR = 6 + AVG = 7 + UNUSED = 8 diff --git a/lib/python3.12/site-packages/deepspeed/comm/torch.py b/lib/python3.12/site-packages/deepspeed/comm/torch.py new file mode 100644 index 0000000000000000000000000000000000000000..af906359900dfc4d2efbcc74bc8185b7db1433e0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/torch.py @@ -0,0 +1,427 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import deepspeed +from deepspeed import utils + +from .utils import * +from .backend import * +from .comm import * +from ..runtime import compiler +from deepspeed.utils.torch import required_torch_version +import os + +DS_COMM_ALL_GATHER_OFF = False +DS_COMM_REDUCE_SCATTER_OFF = False +DS_COMM_BROADCAST_OFF = False +DS_COMM_ALL_REDUCE_OFF = False +DS_COMM_REDUCE_OFF = False + + +def disable_compiler_collective(func): + if required_torch_version(min_version=2.3): + return func + return compiler.disable(func) + + +def build_shm_op(): + builder = get_accelerator().create_op_builder("ShareMemCommBuilder") + if builder is None or not deepspeed.ops.__compatible_ops__[builder.NAME]: + return None + shm_cpp_module = builder.load() + print(f'DeepSpeed {builder.absolute_name()} built successfully') + return shm_cpp_module + + +def has_coalescing_manager(): + has_c10d = hasattr(torch.distributed, 'distributed_c10d') + return has_c10d and hasattr(torch.distributed.distributed_c10d, '_coalescing_manager') + + +def has_all_reduce_coalesced(): + return hasattr(torch.distributed, "all_reduce_coalesced") and required_torch_version(min_version=1.13) + + +def get_coalescing_manager(group, device, reqs, async_op): + if required_torch_version(min_version=2.0, max_version=2.0): + return torch.distributed.distributed_c10d._coalescing_manager(group, device=device, reqs=reqs) + elif required_torch_version(min_version=2.1): + return torch.distributed.distributed_c10d._coalescing_manager(group, device=device, async_ops=async_op) + else: + return torch.distributed.distributed_c10d._coalescing_manager(group, reqs) + + +##Utilities to turn comm off +##TODO: move to base comm (wrapper) +def all_gather_comm_off(flag=False): + global DS_COMM_ALL_GATHER_OFF + DS_COMM_ALL_GATHER_OFF = flag + + +def reduce_scatter_comm_off(flag=False): + global DS_COMM_REDUCE_SCATTER_OFF + DS_COMM_REDUCE_SCATTER_OFF = flag + + +def broadcast_comm_off(flag=False): + global DS_COMM_BROADCAST_OFF + DS_COMM_BROADCAST_OFF = flag + + +def all_reduce_comm_off(flag=False): + global DS_COMM_ALL_REDUCE_OFF + DS_COMM_ALL_REDUCE_OFF = flag + + +def reduce_comm_off(flag=False): + global DS_COMM_REDUCE_OFF + DS_COMM_REDUCE_OFF = flag + + +#assumption: all_gather and reduce scatter +## are what we care about +def backward_comm_off(flag=False): + all_gather_comm_off(flag) + reduce_scatter_comm_off(flag) + + +class Noop: + + def wait(self): + return None + + +class TorchBackend(Backend): + """ + A light-weight wrapper class for torch.distributed API. + Only a subset of functions are wrapped. Once the init_process_group + is initialized, standard torch.distributed.* can be used directly + so no need to wrap all the functions. We can keep adding wrappers as + needed. + """ + + def __init__(self, backend, timeout, init_method, rank=-1, world_size=-1, name='torch'): + super(TorchBackend, self).__init__() + self.shm_comm_op = build_shm_op() + self.has_all_reduce_coalesced = has_all_reduce_coalesced() + self.has_coalescing_manager = has_coalescing_manager() + self.all_gather_function = self.get_all_gather_function() + self.reduce_scatter_function = self.get_reduce_scatter_function() + self.initialized = True + self.name = name + # Future functionality to support ds.initialize() on a single GPU + # The idea is to fake that dist backend is initialized even when + # it is not so we can run on a single GPU without doing any init_process_group + self.single_gpu_mode = True + self.init_process_group(backend, timeout, init_method, rank, world_size) + if self.shm_comm_op != None: + self.shm_comm_op.initialize(self.get_world_size(), self.get_rank()) + + @classmethod + @disable_compiler_collective + def get_all_gather_function(self): + if hasattr(torch.distributed, "all_gather_into_tensor"): + return torch.distributed.all_gather_into_tensor + elif hasattr(torch.distributed, "_all_gather_base"): + return torch.distributed._all_gather_base + return None + + @classmethod + @disable_compiler_collective + def get_reduce_scatter_function(self): + if hasattr(torch.distributed, "reduce_scatter_tensor"): + return torch.distributed.reduce_scatter_tensor + elif hasattr(torch.distributed, "_reduce_scatter_base"): + return torch.distributed._reduce_scatter_base + return None + + def has_all_gather_into_tensor(self): + return self.all_gather_function is not None + + def has_reduce_scatter_tensor(self): + return self.reduce_scatter_function is not None + + def init_process_group(self, backend, timeout, init_method, rank, world_size): + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend, + timeout=timeout, + init_method=init_method, + rank=rank, + world_size=world_size) + self.using_mpi = torch.distributed.get_backend() == 'mpi' + + @disable_compiler_collective + def all_reduce(self, tensor, op=torch.distributed.ReduceOp.SUM, group=None, async_op=False): + op = self._reduce_op(op) + return torch.distributed.all_reduce(tensor=tensor, op=op, group=group, async_op=async_op) + + def inference_all_reduce(self, tensor, op, group=None): + if not hasattr(torch.ops, 'deepspeed') or not hasattr(torch.ops.deepspeed, 'inference_all_reduce_'): + op = self._reduce_op(op) + return torch.distributed.all_reduce(tensor=tensor, op=op, group=group, async_op=False) + else: + return torch.ops.deepspeed.inference_all_reduce_(tensor) + + @disable_compiler_collective + def all_reduce_coalesced(self, tensors, op=torch.distributed.ReduceOp.SUM, group=None, async_op=False): + """ proxy func to torch.distributed.all_reduce_coalesced, + which is included in PyTorch 1.13 and above + """ + if not self.has_all_reduce_coalesced: + raise RuntimeError(f"Current torch version does not have all_reduce_coalesced " + f"api (torch.__version__: {torch.__version__})") + op = self._reduce_op(op) + return torch.distributed.all_reduce_coalesced(tensors=tensors, op=op, group=group, async_op=async_op) + + @disable_compiler_collective + def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False): + if DS_COMM_REDUCE_OFF: + if int(os.getenv('RANK', '0')) == 0: + utils.logger.warning("REDUCE is OFF") + return Noop() + return torch.distributed.reduce(tensor=tensor, dst=dst, op=self._reduce_op(op), group=group, async_op=async_op) + + @disable_compiler_collective + def reduce_scatter(self, output, input_list, op=ReduceOp.SUM, group=None, async_op=False): + if DS_COMM_REDUCE_SCATTER_OFF: + if int(os.getenv('RANK', '0')) == 0: + utils.logger.warning("REDUCE SCATTER is OFF") + return Noop() + else: + return torch.distributed.reduce_scatter(output=output, + input_list=input_list, + op=self._reduce_op(op), + group=group, + async_op=async_op) + + @disable_compiler_collective + def broadcast(self, tensor, src, group=None, async_op=False): + if DS_COMM_BROADCAST_OFF: + if int(os.getenv('RANK', '0')) == 0: + utils.logger.warning("BROADCAST is OFF") + return Noop() + else: + return torch.distributed.broadcast(tensor=tensor, src=src, group=group, async_op=async_op) + + @disable_compiler_collective + def broadcast_object_list(self, object_list, src, group=None, device=None): + return torch.distributed.broadcast_object_list(object_list=object_list, src=src, group=group, device=device) + + @disable_compiler_collective + def all_gather(self, tensor_list, tensor, group=None, async_op=False): + if DS_COMM_ALL_GATHER_OFF: + if int(os.getenv('RANK', '0')) == 0: + utils.logger.warning("All Gather is OFF") + return Noop() + else: + return torch.distributed.all_gather(tensor_list=tensor_list, tensor=tensor, group=group, async_op=async_op) + + @disable_compiler_collective + def all_gather_into_tensor(self, output_tensor, input_tensor, group=None, async_op=False): + if self.has_all_gather_into_tensor(): + return self.all_gather_function(output_tensor=output_tensor, + input_tensor=input_tensor, + group=group, + async_op=async_op) + + @disable_compiler_collective + def all_gather_base(self, output_tensor, input_tensor, group=None, async_op=False): + if DS_COMM_ALL_GATHER_OFF: + if int(os.getenv('RANK', '0')) == 0: + utils.logger.warning("All Gather is OFF") + return Noop() + else: + if self.has_allgather_base: + return torch.distributed.distributed_c10d._all_gather_base(output_tensor=output_tensor, + input_tensor=input_tensor, + group=group, + async_op=async_op) + else: + utils.logger.warning("unable to find torch.distributed._all_gather_base. will fall back to " + "torch.distributed.reduce_scatter which will result in suboptimal performance. " + "please consider upgrading your pytorch installation.") + pass + + @disable_compiler_collective + def all_gather_coalesced(self, output_tensors, input_tensors, group=None, async_op=False): + """""" + assert len(output_tensors) == len(input_tensors), "" + if hasattr(torch.distributed.distributed_c10d, '_all_gather_base_coalesced'): + # customized PyTorch + return torch.distributed.distributed_c10d._all_gather_base_coalesced(output_tensors, + input_tensors, + group=group, + async_op=async_op) + elif has_coalescing_manager(): + reqs = [] + with get_coalescing_manager(group, input_tensors[0].device, reqs, async_op): + for output, input in zip(output_tensors, input_tensors): + handle = torch.distributed.distributed_c10d.all_gather_into_tensor(output, + input, + group=group, + async_op=True) + reqs.append(handle) + if async_op: + return reqs[-1] + else: + reqs[-1].wait() + + @disable_compiler_collective + def all_gather_object(self, object_list, obj, group=None): + return torch.distributed.all_gather_object(object_list=object_list, obj=obj, group=group) + + @disable_compiler_collective + def reduce_scatter_tensor(self, output_tensor, input_tensor, op=ReduceOp.SUM, group=None, async_op=False): + if self.has_reduce_scatter_tensor(): + return self.reduce_scatter_function(output_tensor, + input_tensor, + op=self._reduce_op(op), + group=group, + async_op=async_op) + else: + utils.logger.warning("unable to find torch.distributed.reduce_scatter_tensor. will fall back to " + "torch.distributed.reduce_scatter which will result in suboptimal performance. " + "please consider upgrading your pytorch installation.") + pass + + @disable_compiler_collective + def all_to_all_single(self, + output, + input, + output_split_sizes=None, + input_split_sizes=None, + group=None, + async_op=False): + return torch.distributed.all_to_all_single(output=output, + input=input, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op) + + @disable_compiler_collective + def all_to_all(self, output_tensor_list, input_tensor_list, group=None, async_op=False): + return torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=group, async_op=async_op) + + @disable_compiler_collective + def send(self, tensor, dst, group=None, tag=0): + return torch.distributed.send(tensor=tensor, dst=dst, group=group, tag=tag) + + @disable_compiler_collective + def recv(self, tensor, src=None, group=None, tag=0): + return torch.distributed.recv(tensor=tensor, src=src, group=group, tag=tag) + + @disable_compiler_collective + def isend(self, tensor, dst, group=None, tag=0): + return torch.distributed.isend(tensor=tensor, dst=dst, group=group, tag=tag) + + @disable_compiler_collective + def irecv(self, tensor, src=None, group=None, tag=0): + return torch.distributed.irecv(tensor=tensor, src=src, group=group, tag=tag) + + @disable_compiler_collective + def gather(self, tensor, gather_list=None, dst=0, group=None, async_op=False): + return torch.distributed.gather(tensor=tensor, + gather_list=gather_list, + dst=dst, + group=group, + async_op=async_op) + + @disable_compiler_collective + def scatter(self, tensor, scatter_list=None, src=0, group=None, async_op=False): + return torch.distributed.scatter(tensor=tensor, + scatter_list=scatter_list, + src=src, + group=group, + async_op=async_op) + + @disable_compiler_collective + def barrier(self, group=torch.distributed.GroupMember.WORLD, async_op=False, device_ids=None): + if group is None: + group = torch.distributed.GroupMember.WORLD + return torch.distributed.barrier(group=group, async_op=async_op, device_ids=device_ids) + + @disable_compiler_collective + def monitored_barrier(self, group=torch.distributed.GroupMember.WORLD, timeout=None, wait_all_ranks=False): + if group is None: + group = torch.distributed.GroupMember.WORLD + return torch.distributed.monitored_barrier(group=group, timeout=timeout, wait_all_ranks=wait_all_ranks) + + def get_rank(self, group=None): + return torch.distributed.get_rank(group=group) + + def get_world_size(self, group=None): + return torch.distributed.get_world_size(group=group) + + def is_initialized(self): + return torch.distributed.is_initialized() + + def get_backend(self, group=None): + return torch.distributed.get_backend(group=group) + + def new_group(self, ranks): + return torch.distributed.new_group(ranks) + + def get_global_rank(self, group, group_rank): + if hasattr(torch.distributed.distributed_c10d, "get_global_rank"): + from torch.distributed.distributed_c10d import get_global_rank as _get_global_rank + else: + from torch.distributed.distributed_c10d import _get_global_rank + return _get_global_rank(group, group_rank) + + def get_world_group(self): + return torch.distributed.group.WORLD + + def destroy_process_group(self, group=None): + return torch.distributed.destroy_process_group(group=group) + + def _reduce_op(self, op): + ''' + Helper function. If the op provided is not a torch.dist.ReduceOp, convert it and return + ''' + if not isinstance(op, torch.distributed.ReduceOp): + if op == ReduceOp.SUM: + op = torch.distributed.ReduceOp.SUM + elif op == ReduceOp.PRODUCT: + op = torch.distributed.ReduceOp.PRODUCT + elif op == ReduceOp.AVG: + op = torch.distributed.ReduceOp.AVG + elif op == ReduceOp.MIN: + op = torch.distributed.ReduceOp.MIN + elif op == ReduceOp.MAX: + op = torch.distributed.ReduceOp.MAX + elif op == ReduceOp.BAND: + op = torch.distributed.ReduceOp.BAND + elif op == ReduceOp.BOR: + op = torch.distributed.ReduceOp.BOR + elif op == ReduceOp.BXOR: + op = torch.distributed.ReduceOp.BXOR + return op + + def init_device_mesh(self, mesh_shape, mesh_dim_names): + if not required_torch_version(min_version=2.2): + raise RuntimeError(f"Current torch version does not have device mesh" + f"api (torch.__version__: {torch.__version__})") + if not required_torch_version(max_version=2.4): + return torch.distributed.device_mesh.init_device_mesh(get_accelerator().device_name(), + mesh_shape, + mesh_dim_names=mesh_dim_names) + else: + return torch.distributed.device_mesh.init_device_mesh(get_accelerator().current_device_name(), + mesh_shape, + mesh_dim_names=mesh_dim_names) + + def enable_symm_mem_for_group(self, group_name): + if not required_torch_version(min_version=2.5): + raise RuntimeError(f"Torch version must be 2.5 or higher to use symmetric memory. " + f"Current version: {torch.__version__}") + from torch.distributed._symmetric_memory import enable_symm_mem_for_group + return enable_symm_mem_for_group(group_name) + + +# This will become a light-weight wrapper around torch.distributed functions +# TODO: create some example to show how this wrapper can help profile communication +# TODO: make sure there is no performance regression with this approach +# TODO: explore monkey-patching if this does not work diff --git a/lib/python3.12/site-packages/deepspeed/comm/utils.py b/lib/python3.12/site-packages/deepspeed/comm/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..78bfdf4462e4f57458a1137cb24a271e09fc1b6f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/comm/utils.py @@ -0,0 +1,133 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import inspect +from deepspeed.utils import get_caller_func + + +def get_local_rank_from_launcher(): + + # DeepSpeed launcher will set it so get from there + rank = os.environ.get('LOCAL_RANK') + + if rank is None: + rank = os.environ.get('OMPI_COMM_WORLD_LOCAL_RANK') + + # Make it a single process job and set rank to 0 + if rank is None: + rank = 0 + + return int(rank) + + +def get_world_rank_from_launcher(): + + # DeepSpeed launcher will set it so get from there + rank = os.environ.get('RANK') + + if rank is None: + rank = os.environ.get('OMPI_COMM_WORLD_RANK') + + # Make it a single process job and set rank to 0 + if rank is None: + rank = 0 + + return int(rank) + + +def get_world_size_from_launcher(): + # DeepSpeed launcher will set it so get from there + size = os.environ.get('WORLD_SIZE') + rank = os.environ.get('RANK') + + if size is None: + size = os.environ.get('OMPI_COMM_WORLD_SIZE') + + # Make it a single process job and set size to 1 + if size is None: + size = 1 + + if rank == 0: + print(f"set world size to {size}") + + return int(size) + + +def get_default_args(func): + signature = inspect.signature(func) + return {k: v.default for k, v in signature.parameters.items() if v.default is not inspect.Parameter.empty} + + +# We need this hacky function since torch doesn't consistently name or place the input tensor args +def get_tensor_position(func): + sig_params = inspect.signature(func).parameters + arg = None + # most colls + if 'tensor' in sig_params: + arg = 'tensor' + # all_reduce_coalesced coll + elif 'tensors' in sig_params: + arg = 'tensors' + # reduce scatter coll + elif 'input_list' in sig_params: + arg = 'input_list' + # all_to_all and torch multiGPU colls + elif 'input_tensor_list' in sig_params: + arg = 'input_tensor_list' + if arg is None: + return -1 + else: + return list(sig_params).index(arg) + + +def get_tensor_kwarg(func, kwargs): + func_args = get_default_args(func) + func_args.update(kwargs) + arg = None + + if 'tensor' in func_args: + arg = func_args['tensor'] + elif 'tensors' in func_args: + arg = func_args['tensors'] + elif 'input_list' in func_args: + arg = func_args['input_list'] + elif 'input_tensor_list' in func_args: + arg = func_args['input_tensor_list'] + return arg + + +def get_msg_size_from_args(func, *args, **kwargs): + # 3 cases: + # - tensor arg is in args + # - tensor arg is in kwargs + # - tensor arg is not present (e.g. barrier) + tensor_arg_position = -1 + tensor_arg = None + # check if tensor arg is in args + if len(args) > 0: + tensor_arg_position = get_tensor_position(func) + if tensor_arg_position > -1: + tensor_arg = args[get_tensor_position(func)] + # check if tensor arg is in kwargs + if tensor_arg is None and len(kwargs) > 0: + tensor_arg = get_tensor_kwarg(func, kwargs) + # if tensor arg is not present, no data is being transmitted + if tensor_arg is None: + return 0 + else: + # Sum of tensor sizes for list colls such as torch's all_to_all + # NOTE: msg_size for list colls will not be the actual size transmitted by a given MPI/NCCL call within the coll op. Instead, it's the total amount of data transmitted. + if type(tensor_arg) is list: + return sum(x.element_size() * x.nelement() for x in tensor_arg) + else: + return tensor_arg.element_size() * tensor_arg.nelement() + + +def get_debug_log_name(func_args, debug): + if debug: + return func_args['log_name'] + ' | [Caller Func: ' + get_caller_func() + ']' + else: + return func_args['log_name'] diff --git a/lib/python3.12/site-packages/deepspeed/compile/__init__.py b/lib/python3.12/site-packages/deepspeed/compile/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..208299fb8c50f73468d293b6fa5dca71649d62e7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be5e0e1064c2d03db5db59df051d3e57f1a2f169 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/backend.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0fda96d8e2780f90183a9e862ee7e875ece9ed0c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b069b9133a22592556d8c5b3992080ac98bd5ba Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/fx.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/fx.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35744646595c0727440e599f553c19ade2c03e96 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/fx.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/graph_param.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/graph_param.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2c3be16a721bd6ae8c2d8d3020ce72432e3000b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/graph_param.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/inductor.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/inductor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4fbf707f822f1ef38cbd03f536488a2bfd76c86f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/inductor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/init_z1.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/init_z1.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a470d6bc973569c8a4e50e6df1fca9e85684834b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/init_z1.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/init_z3.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/init_z3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f84f0bfde3a6fde768be56020f1c57e63ac8530d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/init_z3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/list_schedule.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/list_schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c5b35f8cd605a7166c0ea7c75a4faa0903c42d9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/list_schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/partitioner.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/partitioner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7def70889fb9dafa8cdeed2fc0fc0cdabf8715ec Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/partitioner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/patch_compiled_func.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/patch_compiled_func.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d20eace82c0cc023a4c509afeb415144cbf2824c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/patch_compiled_func.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/patch_fake_tensor.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/patch_fake_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7889f4314c3e1228a59db7465782873521bcb351 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/patch_fake_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b978e17ac53379317994b454b1e5071d2fd5595f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/backend.py b/lib/python3.12/site-packages/deepspeed/compile/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..ee33447aaf4a46e063394fbdca4386e44bad6b54 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/backend.py @@ -0,0 +1,279 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Dict, List, Callable +import time +import gc + +import torch +from torch.fx import Graph, GraphModule + +try: + import torch.utils._pytree as pytree + import torch._dynamo + import torch._inductor.scheduler + from functorch.compile import make_boxed_func + from torch._functorch.aot_autograd import aot_module_simplified + from torch._subclasses.fake_tensor import unset_fake_temporarily +except ImportError: + pass + +from deepspeed.accelerator import get_accelerator + +from .fx import add_free_activations +from .graph_param import DSGraphParamManager +from .profilers import ProfilingResult +from .profilers.graph_profile import MemoryProfilingInterpreter +from .patch_compiled_func import patch_compiled_func, unpatch_compiled_func, get_backward_inputs +from .util import get_input_nodes, get_activation_node_names, get_index_by_graph_id, get_deepcompile_handle, log_rank0 +from .partitioner import get_wrapped_partitioner +from .inductor import register_custom_ops, patch_create_aot_dispatcher_function + +remaining_schedule = None +next_pass_step = -1 +next_passes = None +current_passes = None + +param_manager: Dict[int, DSGraphParamManager] = {} +graph_order = [] +profiling_results: Dict[int, ProfilingResult] = {} +opt_pass_times = [] + +opt_passes = {} + +fwd_real_inputs = [] +remaining_bwd_compile_count = 0 + + +def register_compile_pass(name: str, opt_pass_fn): + opt_passes[name] = opt_pass_fn + + +def init_schedule(schedule): + + assert isinstance(schedule, list), f"schedule should be a list, but got {type(schedule)}" + + for step, passes in schedule: + assert isinstance(step, int), f"Each step in schedule should be an integer, but got {type(step)}" + assert isinstance(passes, list), f"Passes at a certain step should be a list, but got {type(passes)}" + + global remaining_schedule + remaining_schedule = schedule + + +def launch_compile_passes(global_steps: int): + global next_pass_step, next_passes + + if len(remaining_schedule) > 0 and global_steps == remaining_schedule[0][0]: + _, next_passes = remaining_schedule.pop(0) + log_rank0(f"Launching compile passes: global_steps={global_steps} passes={next_passes}", True) + + torch._dynamo.reset() + get_deepcompile_handle().reset() + patch_compiled_func() + graph_order.clear() + profiling_results.clear() + param_manager.clear() + + +def set_time_and_tensor_size(graph_id, graph: Graph, mem, bwd, profiling_results): + node_time = [] + tensor_sizes = [] + + for n in graph.nodes: + node_time.append((n.name, n.meta["device_time"] if "device_time" in n.meta else 0.0, + n.meta["wall_time"] if "wall_time" in n.meta else 0.0)) + tensor_sizes.append((n.name, n.meta["tensor_size"] if "tensor_size" in n.meta else 0)) + + if bwd: + profiling_results[graph_id].bwd_graph = graph + profiling_results[graph_id].bwd_time = node_time + profiling_results[graph_id].bwd_tensor_sizes = tensor_sizes + profiling_results[graph_id].bwd_mem = mem + else: + profiling_results[graph_id].fwd_graph = graph + profiling_results[graph_id].fwd_time = node_time + profiling_results[graph_id].fwd_tensor_sizes = tensor_sizes + profiling_results[graph_id].fwd_mem = mem + + +def run_opt_passes(opt_passes: List[Callable], + gm: GraphModule, + graph_id: int, + graph_order: List[int], + profiling_results, + create_inputs_fn, + mem_budget: float, + param_manager, + bwd: bool, + debug_log=False) -> None: + + with unset_fake_temporarily(): + get_accelerator().synchronize() + gc.collect() + get_accelerator().empty_cache() + + for i, opt_pass_fn in enumerate(opt_passes): + log_rank0(f"Running opt pass {i} for graph {graph_id}. bwd={bwd}", enable=debug_log) + + gm_new = opt_pass_fn(gm, graph_id, graph_order, profiling_results, create_inputs_fn, mem_budget, param_manager, + bwd) + if gm_new is not None: + gm = gm_new + gm.graph.lint() + gm.recompile() + + mem_prof = MemoryProfilingInterpreter(gm, debug_log=debug_log) + mem_prof.run(*create_inputs_fn()) + mem = [(name, current_alloc, delta, peak) for name, current_alloc, delta, peak in mem_prof.mem_record] + + set_time_and_tensor_size(graph_id, gm.graph, mem, bwd, profiling_results) + + with unset_fake_temporarily(): + get_accelerator().synchronize() + gc.collect() + get_accelerator().empty_cache() + + +def make_backend(backend, compile_kwargs={}, free_activation=False, debug_log=False): + + register_custom_ops() + + def backend_fn(gm: GraphModule, real_inputs): + graph_id = id(gm.graph) + needs_backward = pytree.tree_any(lambda x: x.requires_grad if torch.is_tensor(x) else False, real_inputs) + + global graph_order + graph_order.append((graph_id, needs_backward)) + + z3_partition = any(hasattr(v, "ds_id") for v in real_inputs) + if z3_partition: + param_indices = [(i, input_val.ds_id, input_val.ds_shape) for i, input_val in enumerate(real_inputs) + if isinstance(input_val, torch.nn.Parameter)] + else: + assert all(hasattr(v, "param_id") for v in real_inputs + if isinstance(v, torch.nn.Parameter)), "All param inputs should have param_id" + param_indices = [(i, input_val.param_id, input_val.shape) for i, input_val in enumerate(real_inputs) + if isinstance(input_val, torch.nn.Parameter)] + + global fwd_real_inputs + fwd_real_inputs.append(real_inputs) + + global profiling_results + if graph_id not in profiling_results: + profiling_results[graph_id] = ProfilingResult() + profiling_results[graph_id].param_indices = param_indices + profiling_results[graph_id].needs_backward = needs_backward + + def make_fw_graph(gm, sample_inputs): + time_start = time.time() + graph_index = len(graph_order) - 1 + real_inputs = fwd_real_inputs.pop(0) + + param_manager[graph_id] = DSGraphParamManager(gm.graph, real_inputs, param_indices) + + real_inputs_with_rng = real_inputs + sample_inputs[len(real_inputs):] + run_opt_passes( + opt_passes=next_passes, + gm=gm, + graph_id=graph_id, + graph_order=graph_order, + profiling_results=profiling_results, + create_inputs_fn=lambda: real_inputs_with_rng, + mem_budget=.0, # unused + param_manager=param_manager, + bwd=False, + debug_log=debug_log) + + if needs_backward: + global remaining_bwd_compile_count + remaining_bwd_compile_count += 1 + + opt_pass_times.append(("fwd", graph_index, graph_id, time.time() - time_start)) + + log_rank0( + f"Fwd end {graph_index} graph_id={graph_id} alloc_mem={get_accelerator().memory_allocated()} graph={gm.graph}", + enable=debug_log) + + return gm.graph + + def make_bw_graph(gm, sample_inputs): + time_start = time.time() + + graph_index = get_index_by_graph_id(graph_order, graph_id) + log_rank0( + f"Bwd start {graph_index} graph_id={graph_id} alloc_mem={get_accelerator().memory_allocated()} graph={gm.graph}", + enable=debug_log) + + bwd_inputs_stack = get_backward_inputs() + + if len(bwd_inputs_stack) == 0: + # dynamo calls bw compiler ahead of time when symints are saved for backward. See the details for aot_dispatch_autograd in jit_compile_runtime_wrappers. + # As we currently use actually bwd input values in bw compiler, we return None to skip the compilation there. + # This would need be handled properly in the future. + return None + + bwd_real_inputs = bwd_inputs_stack.pop() + run_opt_passes( + opt_passes=next_passes, + gm=gm, + graph_id=graph_id, + graph_order=graph_order, + profiling_results=profiling_results, + create_inputs_fn=lambda: tuple(bwd_real_inputs), + mem_budget=.0, # unused + param_manager=param_manager, + bwd=True, + debug_log=debug_log) + + # assert graph_id in param_manager, f"Graph {graph_id} not found in param_manager" + + if free_activation: + param_nodes_bw, _ = param_manager[graph_id].get_bwd_mapping(gm.graph) + param_names = [n.name for n in param_nodes_bw] + non_param_input_names = [n.name for n in get_input_nodes(gm.graph) if n.name not in param_names] + add_free_activations(graph_id, gm.graph, + get_activation_node_names(gm.graph, param_nodes_bw, non_param_input_names)) + + global remaining_bwd_compile_count + remaining_bwd_compile_count -= 1 + if remaining_bwd_compile_count == 0: + unpatch_compiled_func() + + log_rank0( + f"Bwd end {graph_index} graph_id={graph_id} alloc_mem={get_accelerator().memory_allocated()} graph={gm.graph}", + enable=debug_log) + + gm.recompile() + opt_pass_times.append(("bwd", graph_index, graph_id, time.time() - time_start)) + + return gm.graph + + if backend == "eager": + + def make_compiler_fn(make_graph_fn): + + def compiler_fn(gm, sample_inputs): + return None if make_graph_fn(gm, sample_inputs) is None else make_boxed_func(gm.forward) + + return compiler_fn + + aot_mod = aot_module_simplified(gm, + real_inputs, + fw_compiler=make_compiler_fn(make_fw_graph), + bw_compiler=make_compiler_fn(make_bw_graph), + partition_fn=get_wrapped_partitioner(param_indices)) + return torch._dynamo.optimize(**compile_kwargs)(aot_mod) + elif backend == "inductor": + patch_create_aot_dispatcher_function(graph_id, z3_partition, make_fw_graph, make_bw_graph, real_inputs, + param_indices, param_manager) + from .partitioner import get_wrapped_choose_saved_values_set + torch._functorch.partitioners.choose_saved_values_set = get_wrapped_choose_saved_values_set(param_indices) + + return torch._inductor.compile(gm, real_inputs) + + raise ValueError(f"Unsupported backend {backend}") + + return backend_fn diff --git a/lib/python3.12/site-packages/deepspeed/compile/config.py b/lib/python3.12/site-packages/deepspeed/compile/config.py new file mode 100644 index 0000000000000000000000000000000000000000..d88458fc594e4b86dc6d1dcf481060a857ed4065 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/config.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import DeepSpeedConfigModel + + +class CompileConfig(DeepSpeedConfigModel): + """ Configure compile settings """ + + deepcompile: bool = False + """ Turn on/off the DeepCompile mode """ + + free_activation: bool = False + """ Turn on/off the free activation mode """ + + offload_activation: bool = False + """ Turn on/off the activation offloading """ + + offload_opt_states: bool = False + """ Turn on/off the optimizer states offloading """ + + double_buffer: bool = True + """ Turn on/off the double buffering """ + + symmetric_memory: bool = False + """ Turn on/off the symmetric memory """ + + debug_log: bool = False + """ Turn on/off the graph dumping """ + + offload_parameters: bool = False + """ Turn on/off the parameter offloading """ + + sync_before_reduce: bool = False + """ Turn on/off the sync before reduce """ + + sync_after_reduce: bool = False + """ Turn on/off the sync after reduce """ + + sync_before_allgather: bool = False + """ Turn on/off the sync before allgather """ + + sync_after_allgather: bool = False + """ Turn on/off the sync after allgather """ diff --git a/lib/python3.12/site-packages/deepspeed/compile/fx.py b/lib/python3.12/site-packages/deepspeed/compile/fx.py new file mode 100644 index 0000000000000000000000000000000000000000..3506aef7062b0710f95e4b52e5016ee6c6d28126 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/fx.py @@ -0,0 +1,139 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Callable, Any, List +from collections import defaultdict + +import torch +from torch.fx import Node, Graph + +from .util import get_last_uses + + +def get_output_node(graph: Graph): + for v in graph.nodes: + if v.target == "output": + return v + raise ValueError("No output node found") + + +def move_primals_to_head(graph: Graph): + + # Move primals to the head of the graph + primals = [n for n in graph.nodes if n.op == "placeholder"] + non_primals = [n for n in graph.nodes if n.op != "placeholder"] + all_nodes = primals + non_primals + + new_graph = Graph() + env = {} + for node in all_nodes: + new_node = new_graph.node_copy(node, lambda n: env[n.name]) + env[node.name] = new_node + new_graph.lint() + + return new_graph + + +def add_args_process(graph: Graph, + node: Node, + fn: Callable[..., Any], + extra_args: List[int] = [], + name=None, + meta={}) -> List[Node]: + # Apply fn to all args of node + new_nodes = [] + with graph.inserting_before(node): + target_args = [arg for arg in node.args if isinstance(arg, Node)] + + for arg in target_args: + new_node = graph.create_node('call_function', fn, (arg, ) + tuple(extra_args), name=name) + for k, v in meta.items(): + new_node.meta[k] = v + node.replace_input_with(arg, new_node) + new_nodes.append(new_node) + + return new_nodes + + +def add_postprocess(graph: Graph, + node: Node, + fn: Callable[..., Any], + extra_args: List[int] = [], + name=None, + meta={}) -> Node: + # https://github.com/pytorch/examples/blob/main/fx/wrap_output_dynamically.py + with graph.inserting_after(node): + args = (node, ) + for a in extra_args: # To add ds_id + args += (a, ) + + node_users = node.users.keys() + new_node = graph.create_node('call_function', fn, args, {}, name=name) + users = {} + for u in node_users: + if u != new_node: + users[u] = (node, new_node) + for u, (old_in, new_in) in users.items(): + u.replace_input_with(old_in, new_in) + + for k, v in meta.items(): + new_node.meta[k] = v + + return new_node + + +def _make_node_meta(node: Node, ds_id: int, comm: bool): + meta = {"param_name": node.name, "ds_id": ds_id, "comm": comm} + if "tensor_meta" in node.meta: + meta["tensor_meta"] = node.meta["tensor_meta"] + return meta + + +def add_free_activations(graph_id: int, graph: Graph, activation_node_names: List[str]): + node_to_last_use, _ = get_last_uses(graph) + activation_nodes_set = set([n for n in graph.nodes if n.op == "placeholder" and n.name in activation_node_names]) + + offload_id_to_node = {} + node_to_wait_reload = {} + for node in graph.nodes: + if node.target == torch.ops.dc.reload_tensor.default: + offload_act = node.args[0] + # node_to_offload_id[offload_act] = node.args[2] + offload_id_to_node[node.args[2]] = offload_act + elif node.target == torch.ops.dc.wait_reload.default: + offload_id = node.args[2] + node_to_wait_reload[offload_id_to_node[offload_id]] = node + + activation_nodes_set = set(node_to_wait_reload[n] if n in node_to_wait_reload else n for n in activation_nodes_set) + + last_user_to_uses = defaultdict(list) + for node, last_user in node_to_last_use.items(): + last_user_to_uses[last_user].append(node) + + def _should_free(node: Node) -> bool: + if not hasattr(node, "meta"): + return False + if not "tensor_meta" in node.meta: + return False + return True + + def free_tensors(tensors: List[torch.Tensor]): + for a in tensors: + if a.numel() > 10_000_000: + a.data = torch.empty([0], device=a.device, dtype=a.dtype) + + for last_user, used_nodes in last_user_to_uses.items(): + activation_args = [an for an in used_nodes if an in activation_nodes_set and _should_free(an)] + + if len(activation_args) == 0: + continue + + node_name = f"free_activations_{[n.name for n in used_nodes]}" + with graph.inserting_after(last_user): + args = (activation_args, ) + graph.create_node('call_function', torch.ops.dc.free_tensors.default, args, {}, name=node_name) + + # Python version for debugging + # graph.create_node('call_function', free_tensors, args, {}, name=node_name) diff --git a/lib/python3.12/site-packages/deepspeed/compile/graph_param.py b/lib/python3.12/site-packages/deepspeed/compile/graph_param.py new file mode 100644 index 0000000000000000000000000000000000000000..445af97374f099d650c5d9e8a501c52fade1c3c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/graph_param.py @@ -0,0 +1,84 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Tuple +from functools import reduce + +import torch +from torch.fx import Graph, Node + +from .fx import get_output_node +from .util import get_param_nodes + + +@dataclass +class DSGraphParam: + name: str + shape: torch.Size + dtype: torch.dtype + device: torch.device + node: Node + allgather_node: Node + release_node: Node + param: torch.Tensor + numel: int = field(init=False) + + def __post_init__(self): + self.numel = reduce(lambda x, y: x * y, self.shape) + + +class DSGraphParamManager: + + def __init__(self, fw_graph: Graph, sample_inputs: Any, index_to_ds_ids: List[Tuple[int, int, int]]): + self._fw_graph = fw_graph + self._bw_graph = None + self._params: Dict[str, DSGraphParam] = {} + self._param_name_to_grad: Dict[str, Node] = {} + self._ds_ids: Dict[str, int] = {} + + param_nodes = get_param_nodes(fw_graph, index_to_ds_ids) + self._param_names = [pn.name for pn in param_nodes] + self._param_indices = [i for i, _, _ in index_to_ds_ids] + + param_inputs = [sample_inputs[i] for i, _, _ in index_to_ds_ids] + ds_ids = [ds_id for _, ds_id, _ in index_to_ds_ids] + ds_shapes = [ds_shape for _, _, ds_shape in index_to_ds_ids] + + for pn, pi, ds_id, ds_shape in zip(param_nodes, param_inputs, ds_ids, ds_shapes): + self._params[pn.name] = DSGraphParam(name=pn.name, + shape=ds_shape, + dtype=pi.dtype, + device=pi.device, + node=pn, + allgather_node=None, + release_node=None, + param=pi) + self._ds_ids[pn.name] = ds_id + + def get_bwd_mapping(self, bw_graph: Graph): + self._bw_graph = bw_graph + + output_node = get_output_node(bw_graph) + param_nodes_bw = [n for n in self._bw_graph.nodes if n.name in self.param_names] + grad_outputs = [output_node.args[0][i] for i in self._param_indices] + param_name_to_grad = {param_name: grad for param_name, grad in zip(self.param_names, grad_outputs)} + return param_nodes_bw, param_name_to_grad + + @property + def param_names(self) -> List[str]: + return self._param_names + + @property + def params(self) -> Dict[str, DSGraphParam]: + return self._params + + @property + def ds_ids(self) -> Dict[str, int]: + return self._ds_ids + + def get_grad_name(self, param_name) -> str: + assert self._param_name_to_grad is not None, "Backward graph is not added yet" + return self._param_name_to_grad[param_name] diff --git a/lib/python3.12/site-packages/deepspeed/compile/inductor.py b/lib/python3.12/site-packages/deepspeed/compile/inductor.py new file mode 100644 index 0000000000000000000000000000000000000000..7ba1c17ee43be103dc94b72f72649ffaafdedd59 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/inductor.py @@ -0,0 +1,214 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +try: + import torch.utils._pytree as pytree + from torch._functorch.aot_autograd import create_aot_dispatcher_function + from torch._inductor.lowering import register_lowering, fallbacks, add_needs_realized_inputs + from torch._inductor.ir import TensorBox, FallbackKernel, Layout, IRNode + from torch._inductor.virtualized import V + from torch._inductor.scheduler import Scheduler + + original_create_aot_dispatcher_function = create_aot_dispatcher_function +except ImportError: + pass + +from .util import get_input_nodes +from .graph_param import DSGraphParamManager + + +def patch_compiler(original_compiler, dc_compiler, z3_partition: bool, graph_id, graph_param_manager, bwd: bool): + + def wrapped_compiler(gm, fake_inputs): + mod_graph = dc_compiler(gm, fake_inputs) + + # For symint case + if mod_graph is None: + return None + + if z3_partition: + # Inductor validates input size estimated by the first trace, where ds tensor is materialized. + # We need to patch the input tensors to avoid the validation error. + patched_inputs = [] + if bwd: + param_nodes_bw, _ = graph_param_manager[graph_id].get_bwd_mapping(gm.graph) + param_names = [n.name for n in param_nodes_bw] + else: + param_names = graph_param_manager[graph_id].param_names + input_nodes = get_input_nodes(gm.graph) + + for in_node, in_v in zip(input_nodes, fake_inputs): + ds_param = in_node.name in param_names + if ds_param: + from torch._subclasses.fake_tensor import is_fake + from torch._dynamo.utils import to_fake_tensor + assert is_fake(in_v), f"Input {in_v} should be fake tensor" + patched_inputs.append( + to_fake_tensor(torch.empty([0], dtype=in_v.dtype, device=in_v.device), in_v.fake_mode)) + else: + patched_inputs.append(in_v) + + patched_inputs = tuple(patched_inputs) + else: + patched_inputs = fake_inputs + + return original_compiler(gm, patched_inputs) + + return wrapped_compiler + + +def wrap_partition_fn(partition_fn, real_inputs, param_indices): + + def wrapped_partition_fn(*args, **kwargs): + + fw_module, bw_module = partition_fn(*args, **kwargs) + + # get parameter names + pm = DSGraphParamManager(fw_module.graph, real_inputs, param_indices) + + def fix_placeholder_meta(graph): + for n in graph.nodes: + if n.op == "placeholder" and n.name in pm.param_names: + n.meta["val"] = torch.empty([0], dtype=n.meta["val"].dtype, device=n.meta["val"].device) + + fix_placeholder_meta(fw_module.graph) + fix_placeholder_meta(bw_module.graph) + + return fw_module, bw_module + + return wrapped_partition_fn + + +def patch_create_aot_dispatcher_function(graph_id: int, z3_partition: bool, make_fw_graph, make_bw_graph, real_inputs, + param_indices, param_manager): + + from torch._dynamo.backends.common import AotAutograd + import functools + + def patch_aotautograd(): + # Unpatch if it was already patched + if hasattr(AotAutograd, "__original_init"): + AotAutograd.__init__ = AotAutograd.__original_init + + original_init = AotAutograd.__init__ + + @functools.wraps(original_init) + def patched_init(self, **kwargs): + kwargs["fw_compiler"] = patch_compiler(kwargs["fw_compiler"], + make_fw_graph, + z3_partition, + graph_id, + param_manager, + bwd=False) + kwargs["bw_compiler"] = patch_compiler(kwargs["bw_compiler"], + make_bw_graph, + z3_partition, + graph_id, + param_manager, + bwd=True) + kwargs["inference_compiler"] = kwargs["fw_compiler"] + + if z3_partition: + kwargs["partition_fn"] = wrap_partition_fn(kwargs["partition_fn"], real_inputs, param_indices) + + original_init(self, **kwargs) + + AotAutograd.__original_init = original_init + AotAutograd.__init__ = patched_init + + patch_aotautograd() + + +def register_custom_ops(): + + def fallback_handler_no_reuse(kernel, + never_reuse_input, + never_reuse_output, + force_free_input, + add_to_fallback_set=True): + if add_to_fallback_set: + fallbacks.add(kernel) + + def handler(*args, **kwargs): + + def wrap_tensors(x): + out = TensorBox.create(x) if isinstance(x, torch._inductor.ir.IRNode) else x + if out is not None and never_reuse_output: + V.graph.never_reuse_buffers.add(out.get_name()) + return out + + class CustomDCKernel(FallbackKernel): + + def __init__(self, op, *args, **kwargs): + super().__init__(op, *args, **kwargs) + + def add_to_never_reuse(x): + if isinstance(x, IRNode): + assert hasattr(x, "get_name"), f"x doesn't have get_name {x.__class__}" + V.graph.never_reuse_buffers.add(x.get_name()) + + if never_reuse_input: + pytree.tree_map(add_to_never_reuse, args) + + def get_var_name_for_arg(self, arg: str): + if arg.isidentifier(): + return arg + + import re + match = re.match(r"reinterpret_tensor\((\w+),", arg) + if match: + return match.group(1) + return None + + def codegen(self, wrapper): + if not force_free_input: + return super().codegen(wrapper) + + kernel = self.op_overload + self.codegen_comment(wrapper) + args = [*self.codegen_args(), *self.codegen_kwargs()] + + V.graph.wrapper_code.generate_fallback_kernel(self, args) + if isinstance(self.layout, Layout): + self.codegen_size_asserts(wrapper) + + var_name = self.get_var_name_for_arg(args[0]) + if var_name: + wrapper.writeline(f"{var_name} = None") + + self.codegen_unbacked_symbol_defs(wrapper) + + kernel_cls = CustomDCKernel if force_free_input else FallbackKernel + return pytree.tree_map(wrap_tensors, kernel_cls.create(kernel, *args, **kwargs)) + + return handler + + def register_fallback_no_reuse(op_overload, + never_reuse_input=False, + never_reuse_output=False, + force_free_input=False): + add_needs_realized_inputs(op_overload) + return register_lowering(op_overload, type_promotion_kind=None)(fallback_handler_no_reuse( + op_overload, + never_reuse_input=never_reuse_input, + never_reuse_output=never_reuse_output, + force_free_input=force_free_input)) + + # Inductor tries to reuse output buffer when possible. We need to disable this behavior for some custom ops. + # -> It seems that memory region is still reused in some cases. So we clone the inputs for some ops. + register_fallback_no_reuse(torch.ops.dc.allgather_param.default, never_reuse_input=False, never_reuse_output=True) + register_fallback_no_reuse(torch.ops.dc.wait_allgather.default, never_reuse_input=True, never_reuse_output=True) + register_fallback_no_reuse(torch.ops.dc.release_param.default, never_reuse_input=True, never_reuse_output=False) + register_fallback_no_reuse(torch.ops.dc.reduce_grad.default, + never_reuse_input=True, + never_reuse_output=True, + force_free_input=True) + register_fallback_no_reuse(torch.ops.dc.free_tensors.default, never_reuse_input=True, never_reuse_output=True) + + if not hasattr(Scheduler, "is_dc_patched") or not Scheduler.is_dc_patched: + Scheduler.is_dc_patched = True + Scheduler.dead_node_elimination = lambda _: None diff --git a/lib/python3.12/site-packages/deepspeed/compile/init_z1.py b/lib/python3.12/site-packages/deepspeed/compile/init_z1.py new file mode 100644 index 0000000000000000000000000000000000000000..2591e9db8e019d1bb6ea97f6dde60935ca158bcd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/init_z1.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import copy + +import torch + +from deepspeed.accelerator import get_accelerator +from .passes import zero1_compile, zero3_compile +from .backend import make_backend, launch_compile_passes, init_schedule +from .util import get_deepcompile_handle, add_pre_backward_hook, is_backend_inductor + +WARMUP = 5 + + +def init_z1(engine, backend, compile_config, compile_kwargs, schedule=None): + + optimizer = engine.optimizer + optimizer.contiguous_gradients = False # Avoid creating unnecessary buffer + for hook in optimizer._grad_acc_hooks: + hook.remove() + optimizer._grad_acc_hooks.clear() + + dc = get_deepcompile_handle() + dc.init(engine.data_parallel_group, + engine.zero_reduce_bucket_size(), compile_config.double_buffer, compile_config.symmetric_memory, + is_backend_inductor(backend), compile_config.sync_before_reduce, compile_config.sync_after_reduce, False, + False) + + grad_buffer = {} + + for i, group in enumerate(optimizer.bit16_groups): + + grad_buffer[i] = optimizer.get_flat_partition(optimizer.params_in_partition[i], + optimizer.first_offset[i], + optimizer.partition_size[i], + dtype=optimizer.gradient_accumulation_dtype, + device=get_accelerator().current_device_name(), + return_tensor_list=True) + grad_buffer[i] = [p.clone().detach() for p in grad_buffer[i]] # Maybe not necessary + + index_in_partition = 0 + first_in_partition = True + for p in group: + param_id = optimizer.get_param_id(p) + p.param_id = param_id + in_partition = optimizer.is_param_in_current_partition[param_id] + + if in_partition: + buf = grad_buffer[i][index_in_partition] + offset = optimizer.first_offset[i] if first_in_partition else 0 + # print(f"[r{dist.get_rank()}] Registering group {i} param {param_id} in_partition={in_partition} p={p.shape} buf={buf.shape} partition_offset={offset}") + dc.register_z1_param(p.param_id, p.shape, p, buf, int(offset)) + index_in_partition += 1 + first_in_partition = False + else: + # print(f"[r{dist.get_rank()}] Registering group {i} param {param_id} in_partition={in_partition} p={p.shape} buf=None") + dc.register_z1_param(p.param_id, p.shape, p, torch.empty([0], dtype=p.dtype, device=p.device), 0) + + def set_grad_buffer(): + optimizer.averaged_gradients = copy.copy(grad_buffer) + + add_pre_backward_hook(set_grad_buffer) + + if schedule is None: + schedule = [] + schedule.append((0, [zero1_compile.add_z1_reduce])) + else: + for opt in schedule: + # avoid typical misconfiguration + if zero3_compile.add_z3_gather_release in opt[1]: + raise ValueError("A pass for ZeRO3 is not specified though ZeRO1 is enabled") + + init_schedule(schedule) + + engine.launch_compile_passes = launch_compile_passes + return make_backend(backend, + compile_kwargs=compile_kwargs, + free_activation=False, + debug_log=compile_config.debug_log) diff --git a/lib/python3.12/site-packages/deepspeed/compile/init_z3.py b/lib/python3.12/site-packages/deepspeed/compile/init_z3.py new file mode 100644 index 0000000000000000000000000000000000000000..f05de840de03e597277d18513ef0d2b834e71559 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/init_z3.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.partition_parameters import InsertPostInitMethodToModuleSubClasses + +from .passes import zero3_compile, prefetch, selective_gather, offload_parameters +from .backend import make_backend, launch_compile_passes, init_schedule +from .patch_fake_tensor import patch_fake_tensor +from .util import get_deepcompile_handle, add_pre_backward_hook, is_backend_inductor + +WARMUP = 5 + + +def init_z3(engine, backend, compile_config, compile_kwargs, schedule=None): + + optimizer = engine.optimizer + if optimizer is not None and hasattr(optimizer, '_DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer'): + optimizer._DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer = None + get_accelerator().empty_cache() + + dc = get_deepcompile_handle() + dc.init(engine.data_parallel_group, + engine.zero_reduce_bucket_size(), compile_config.double_buffer, compile_config.symmetric_memory, + is_backend_inductor(backend), compile_config.sync_before_reduce, compile_config.sync_after_reduce, + compile_config.sync_before_allgather, compile_config.sync_after_allgather) + + # Unset hooks + for m in engine.module.modules(): + m._parameters = m._original_parameters + optimizer.parameter_offload._remove_module_hooks() + + for hook in optimizer._grad_acc_hooks: + hook.remove() + optimizer._grad_acc_hooks.clear() + + # Unpatch linear + if hasattr(InsertPostInitMethodToModuleSubClasses, "linear_bk"): + torch.nn.functional.linear = InsertPostInitMethodToModuleSubClasses.linear_bk + + if compile_config.symmetric_memory: + group_name = engine.data_parallel_group.group_name + dist.enable_symm_mem_for_group(group_name) + + for p in engine.module.parameters(): + grad_buffer = optimizer._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[p.ds_id] + + # Disable persistent param + p.ds_persist = False + dc.register_z3_param(p.ds_id, p.ds_shape, p.ds_tensor, grad_buffer, p.ds_persist) + + def set_grad_buffer(): + for i, sub_group in enumerate(optimizer.fp16_groups): + optimizer.averaged_gradients[i] = [ + optimizer._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[param.ds_id] + if param.requires_grad else torch.zeros_like(param.ds_tensor) for param in sub_group + ] + + add_pre_backward_hook(set_grad_buffer) + + if schedule is None: + schedule = [] + if (compile_config.offload_parameters): + schedule.append((0, [zero3_compile.add_z3_gather_release, offload_parameters.offload_parameter_fwd])) + else: + schedule.append((0, [zero3_compile.add_z3_gather_release])) + schedule.append( + (WARMUP, + [zero3_compile.add_z3_gather_release, prefetch.schedule_prefetch, selective_gather.selective_gather])) + + init_schedule(schedule) + + # offloading opt states need additional setup + from .passes.offload_adam_states import move_opt_states, move_opt_states_sync, init_offload_opt_states + for _, passes in schedule: + if move_opt_states in passes or move_opt_states_sync in passes: + init_offload_opt_states(optimizer, dc) + + engine.launch_compile_passes = launch_compile_passes + + patch_fake_tensor() + free_activation = compile_config.free_activation and not is_backend_inductor(backend) + + torch._inductor.config.size_asserts = False + + return make_backend(backend, + compile_kwargs=compile_kwargs, + free_activation=free_activation, + debug_log=compile_config.debug_log) diff --git a/lib/python3.12/site-packages/deepspeed/compile/list_schedule.py b/lib/python3.12/site-packages/deepspeed/compile/list_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..8df84498a5818d9ea7a377d2313e58dc6145dc4b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/list_schedule.py @@ -0,0 +1,431 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from collections import defaultdict +from typing import List, Dict +from copy import copy +from dataclasses import dataclass + +import torch +from torch.fx import Graph, Node +from torch.fx.node import map_arg + +try: + from torch.utils._pytree import tree_iter +except ImportError: + pass + +from .util import get_last_uses, is_release_node +from .fx import get_output_node + + +def make_graph_from_schedule(scheduled: List[Node]): + new_graph = Graph() + env = {} + for node in scheduled: + new_node = new_graph.node_copy(node, lambda n: env[n.name]) + env[node.name] = new_node + + return new_graph + + +def get_original_args_num(node: Node): + if node.name.startswith("allgather_ds_param") \ + or node.name.startswith("release_ds_param") \ + or node.name.startswith("wait_allgather_ds_param") \ + or node.name.startswith("reduce_ds_param"): + return 1 + + return len(node.args) + + +def flat_nodes_in_args(args: List[Node]): + return [a for a in tree_iter(args) if isinstance(a, Node)] + + +def filter_args(node: Node): + args = node.args[:get_original_args_num(node)] + return flat_nodes_in_args(args) + + +def init_schedule(graph: Graph): + mem_table = create_mem_table(graph) + remaining_users = defaultdict(set) + user_to_producer = {} + + scheduled = [] + unscheduled = [] + edges = defaultdict(list) + for node in graph.nodes: + filtered_args = filter_args(node) + # print(f"Node: {node} args: {node.args}") + if len(filtered_args) == 0: + scheduled.append(node) + + remaining_users[node] = set(node.users.keys()) + for user in node.users.keys(): + user_to_producer[user] = node + else: + unscheduled.append(node) + for a in filtered_args: + for elem_a in tree_iter(a): + if isinstance(elem_a, Node): + if node not in edges[elem_a]: + edges[elem_a].append(node) + + return scheduled, unscheduled, edges, mem_table, remaining_users, user_to_producer + + +def get_runnable_nodes(scheduled: List[Node], unscheduled: List[Node]): + scheduled = set(scheduled) + return [node for node in unscheduled if all(arg in scheduled for arg in filter_args(node))] + + +def choose_next_node(scheduled: List[Node], unscheduled: List[Node], mem_table: Dict[str, int]): + runnable_nodes = get_runnable_nodes(scheduled, unscheduled) + + # sort by memory usage + runnable_nodes = sorted(runnable_nodes, key=lambda n: mem_table[n.name]) + return runnable_nodes[0] + + +def create_mem_table(graph: Graph) -> Dict[str, int]: + mem_table = {} + for node in graph.nodes: + if node.name.startswith("allgather_ds_param"): + mem_table[node.name] = node.meta["tensor_size"] + elif node.name.startswith("release_ds_param") or node.name.startswith("reduce_ds_param"): + mem_table[node.name] = -node.meta["tensor_size"] + else: + mem_table[node.name] = 0 + + return mem_table + + +def list_schedule(graph: Graph) -> Graph: + + scheduled, unscheduled, mem_table = init_schedule(graph) + + while len(unscheduled) > 0: + next_node = choose_next_node(scheduled, unscheduled, mem_table) + scheduled.append(next_node) + unscheduled.remove(next_node) + + return make_graph_from_schedule(scheduled) + + +############################### + + +def get_new_runnable_nodes_with(scheduled: List[Node], edges: Dict[Node, List[Node]], new_scheduled: Node): + scheduled = set(scheduled) + new_runnables = [] + for node in edges[new_scheduled]: + if all(arg in scheduled for arg in filter_args(node) if arg != new_scheduled): + new_runnables.append(node) + + return new_runnables + + +def _do_schedule_without_allgather(scheduled: List[Node], unscheduled: List[Node], edges: Dict[Node, List[Node]], + non_ag_runnable: List[Node]): + + while len(non_ag_runnable) > 0: + next_node = non_ag_runnable.pop() + + new_runnables = get_new_runnable_nodes_with(scheduled, edges, next_node) + non_ag_runnable += [n for n in new_runnables if not n.name.startswith("allgather_ds_param")] + + scheduled.append(next_node) + unscheduled.remove(next_node) + + return scheduled, unscheduled + + +def schedule_without_allgather(scheduled: List[Node], unscheduled: List[Node], edges: Dict[Node, List[Node]]): + runnable = get_runnable_nodes(scheduled, unscheduled) + non_ag_runnable = [n for n in runnable if not n.name.startswith("allgather_ds_param")] + + tmp_scheduled = copy(scheduled) + tmp_unscheduled = copy(unscheduled) + + return _do_schedule_without_allgather(tmp_scheduled, tmp_unscheduled, edges, non_ag_runnable) + + +def try_schedule_with_new_allgather(scheduled: List[Node], unscheduled: List[Node], edges: Dict[Node, List[Node]], + new_scheduled: Node): + new_runnables = get_new_runnable_nodes_with(scheduled, edges, new_scheduled) + non_ag_runnable = [n for n in new_runnables if not n.name.startswith("allgather_ds_param")] + + tmp_scheduled = copy(scheduled) + tmp_unscheduled = copy(unscheduled) + + tmp_scheduled.append(new_scheduled) + tmp_unscheduled.remove(new_scheduled) + + return _do_schedule_without_allgather(tmp_scheduled, tmp_unscheduled, edges, non_ag_runnable) + + +def simple_prefetch(graph: Graph, available_mem: int, output_size: int, debug_log: bool) -> Graph: + + scheduled, unscheduled, edges, mem_table, remaining_users, user_to_producer = init_schedule(graph) + tmp_scheduled, tmp_unscheduled = schedule_without_allgather(scheduled, unscheduled, edges) + + while len(tmp_unscheduled) > 0: + + runnable = get_runnable_nodes(tmp_scheduled, tmp_unscheduled) + ag_with_unblock_time = [] + + for ag_node in runnable: + ag_scheduled, ag_unscheduled = try_schedule_with_new_allgather(tmp_scheduled, tmp_unscheduled, edges, + ag_node) + unblock_time = sum(n.meta["device_time"] for n in ag_scheduled[len(tmp_scheduled) + 1:]) + ag_with_unblock_time.append((ag_node, unblock_time, ag_scheduled, ag_unscheduled)) + + ag_with_unblock_time = sorted(ag_with_unblock_time, key=lambda x: x[1], reverse=True) + best_ag_node = ag_with_unblock_time[0][0] + best_ag_scheduled = ag_with_unblock_time[0][2] + + no_ag_runnables = tmp_scheduled[len(scheduled):] + after_ag_runnables = best_ag_scheduled[len(tmp_scheduled) + 1:] + + scheduled.append(best_ag_node) + unscheduled.remove(best_ag_node) + for n in no_ag_runnables: + scheduled.append(n) + unscheduled.remove(n) + + tmp_scheduled = copy(scheduled) + tmp_unscheduled = copy(unscheduled) + for n in after_ag_runnables: + tmp_scheduled.append(n) + tmp_unscheduled.remove(n) + + return make_graph_from_schedule(tmp_scheduled) + + +############################### + + +def init_schedule_with_placeholders(graph: Graph): + mem_table = create_mem_table(graph) + remaining_users = defaultdict(set) + user_to_producer = {} + + scheduled = [] + unscheduled = [] + edges = defaultdict(list) + for node in graph.nodes: + if node.op == 'placeholder': + scheduled.append(node) + + remaining_users[node] = set(node.users.keys()) + for user in node.users.keys(): + user_to_producer[user] = node + else: + unscheduled.append(node) + + return scheduled, unscheduled, edges, mem_table, remaining_users, user_to_producer + + +def get_node_requirements(target_node: Node, scheduled: List[Node]): + scheduled = set(scheduled) + visited = set() + ordered_nodes = [] + + def dfs(node: Node): + if node in scheduled: + return + if node in visited: + return + visited.add(node) + + args = [] + + def register_arg(n: Node): + args.append(n) + + map_arg(node.args, register_arg) + + for arg in args: + dfs(arg) + ordered_nodes.append(node) + + dfs(target_node) + + return ordered_nodes + + +@dataclass +class AllgatherTask: + node: Node + allgather_cost: float + free_cost: float + allgathered_mem: int + allgather_acc_mem: int + free_acc_mem: int + last_use: Node + n_scheduled_ags: int + schedule_until_ag: List[Node] + schedule_until_free: List[Node] + + +def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug_log: bool) -> Graph: + node_to_last_use, user_to_last_uses = get_last_uses(graph) + + # check tensor size + for node in graph.nodes: + if "tensor_size" not in node.meta: + # Our profiler may not visit all nodes because of the control flow. + node.meta["tensor_size"] = 0 + + scheduled, unscheduled, edges, mem_table, remaining_users, user_to_producer = init_schedule_with_placeholders( + graph) + + unscheduled_ags = [n for n in unscheduled if n.target == torch.ops.dc.allgather_param.default] + + release_nodes = defaultdict(list) + for n in unscheduled: + if is_release_node(n): + release_nodes[n.args[2]].append(n) + + ag_nodes_in_path = {} + for ag_node in unscheduled_ags: + last_use = node_to_last_use[ag_node] + required_nodes = get_node_requirements(last_use, scheduled) + ag_nodes_in_path[ag_node] = set(n for n in required_nodes if n.target == torch.ops.dc.allgather_param.default) + + reduce_nodes = [n for n in unscheduled if n.target == torch.ops.dc.reduce_grad.default] + ag_nodes_in_path_to_reduce_nodes = {} + for reduce_node in reduce_nodes: + ag_nodes_in_path_to_reduce_nodes[reduce_node] = set(n for n in get_node_requirements(reduce_node, scheduled) + if n.target == torch.ops.dc.allgather_param.default) + + output_nodes = [ + n for n in get_output_node(graph).args[0] + if isinstance(n, Node) and n.target != torch.ops.dc.reduce_grad.default + ] + ag_nodes_in_path_to_output_nodes = {} + for output_node in output_nodes: + ag_nodes_in_path_to_output_nodes[output_node] = set(n for n in get_node_requirements(output_node, scheduled) + if n.target == torch.ops.dc.allgather_param.default) + + while len(unscheduled_ags) > 0: + + ag_nodes_count = {ag_node: len(nodes) for ag_node, nodes in ag_nodes_in_path.items()} + count_list = sorted(set(ag_nodes_count.values())) + + runnable_ags = [] + for ag_count in count_list: + + target_unscheduled_ags = [ag for ag in unscheduled_ags if ag_nodes_count[ag] == ag_count] + + for node in target_unscheduled_ags: + ds_id = node.args[2] + + schedule_until_ag = get_node_requirements(node, scheduled) + if schedule_until_ag is None: + continue + + last_use = node_to_last_use[node] + + diff_required_nodes = get_node_requirements(last_use, scheduled + schedule_until_ag) + + allgather_cost = sum(n.meta["device_time"] for n in schedule_until_ag) + free_cost = sum(n.meta["device_time"] for n in diff_required_nodes) + allgathered_mem = node.meta["tensor_size"] + allgather_acc_mem = sum(n.meta["tensor_size"] for n in schedule_until_ag + if n.target == torch.ops.dc.allgather_param.default) + free_acc_mem = sum(n.meta["tensor_size"] for n in diff_required_nodes + if n.target == torch.ops.dc.allgather_param.default) + + schedule_until_free = schedule_until_ag + diff_required_nodes + for release_node in release_nodes[ds_id]: + if release_node not in schedule_until_free: + schedule_until_free.append(release_node) + + n_scheduled_ags = len( + [n for n in schedule_until_free if n.target == torch.ops.dc.allgather_param.default]) + + task = AllgatherTask(node, allgather_cost, free_cost, allgathered_mem, allgather_acc_mem, free_acc_mem, + last_use, n_scheduled_ags, schedule_until_ag, schedule_until_free) + + # print(f" ag_count {ag_count} allgather runnable {i}: {node} last_use: {node_to_last_use[node]} t: {t2-t1:.2f}") + runnable_ags.append(task) + + if len(runnable_ags) > 0: + break + + assert len(runnable_ags) > 0, "No runnable allgather nodes" + + # Criteria of the choice: + # We want to choose allgather that does not require additional allgather until releasing the param. + # When we can find such a node, free_acc_mem will be zero. In that case, we choose the one with the smallest cost until free to minimize the period of occupying memory for the gathered param. + # If there is no such node, we choose the one with the smallest free_cost to minimize the period of occupying memory for the gathered param. + ags_with_no_additional_ag = [ag for ag in runnable_ags if ag.free_acc_mem == 0] + if len(ags_with_no_additional_ag) > 0: + sorted_ags = sorted(runnable_ags, key=lambda x: x.free_cost) + next_ag = sorted_ags[0] + nodes_to_schedule = next_ag.schedule_until_free + else: + # sorted_ags = sorted(runnable_ags, key=lambda x: x.allgathered_mem) + sorted_ags = sorted(runnable_ags, key=lambda x: x.free_acc_mem) + next_ag = sorted_ags[0] + nodes_to_schedule = next_ag.schedule_until_ag + + # print(f" next_ag {next_ag}") + for n in nodes_to_schedule: + scheduled.append(n) + unscheduled.remove(n) + + unscheduled_ags.remove(next_ag.node) + + ag_nodes_in_path.pop(next_ag.node) + for ag_node, nodes in ag_nodes_in_path.items(): + if next_ag.node in nodes: + nodes.remove(next_ag.node) + + # Schedule reduce nodes when possible to free memory earlier + reduces_to_schedule = [] + for reduce_node in reduce_nodes: + if next_ag.node in ag_nodes_in_path_to_reduce_nodes[reduce_node]: + ag_nodes_in_path_to_reduce_nodes[reduce_node].remove(next_ag.node) + if len(ag_nodes_in_path_to_reduce_nodes[reduce_node]) == 0: + reduces_to_schedule.append(reduce_node) + + for n in reduces_to_schedule: + need_to_schedule = get_node_requirements(n, scheduled) + for nn in need_to_schedule: + scheduled.append(nn) + unscheduled.remove(nn) + + # Do the same for output nodes + outputs_to_schedule = [] + for output_node in output_nodes: + if next_ag.node in ag_nodes_in_path_to_output_nodes[output_node]: + ag_nodes_in_path_to_output_nodes[output_node].remove(next_ag.node) + if len(ag_nodes_in_path_to_output_nodes[output_node]) == 0: + outputs_to_schedule.append(output_node) + + for n in outputs_to_schedule: + need_to_schedule = get_node_requirements(n, scheduled) + for nn in need_to_schedule: + scheduled.append(nn) + unscheduled.remove(nn) + + # print(f"After ag scheduled: scheduled: {scheduled}") + + scheduled_set = set(scheduled) + for node in graph.nodes: + if node in scheduled_set: + continue + scheduled.append(node) + unscheduled.remove(node) + + assert len(unscheduled) == 0, f"There are unscheduled nodes: {unscheduled}" + + ret_graph = make_graph_from_schedule(scheduled) + ret_graph.lint() + return ret_graph diff --git a/lib/python3.12/site-packages/deepspeed/compile/partitioner.py b/lib/python3.12/site-packages/deepspeed/compile/partitioner.py new file mode 100644 index 0000000000000000000000000000000000000000..f60170d0b56b32b40ed58ba8065c4a85b04c25f5 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/partitioner.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# This file was copied from PyTorch and modified for DeepSpeed. + +from typing import Tuple, List +import operator + +import torch +from torch.fx import GraphModule, Graph, Node + +try: + from torch._functorch.partitioners import is_sym_node, _is_primal, _is_fwd_seed_offset, _extract_fwd_bwd_outputs, _extract_graph_with_inputs_outputs, _extract_fwd_bwd_modules, has_recomputable_ops, min_cut_rematerialization_partition, choose_saved_values_set +except ImportError: + pass + +from .util import get_no_copy_ops + +_recompute_ops = {torch.ops.aten.t.default} + + +def _find_recompute_nodes(graph: Graph, ds_param_node: Node) -> List[Node]: + """ + Given a graph and a node that represents a parameter that was allgathered, + find all nodes that use the parameter and require recomputation. + """ + no_copy_ops = get_no_copy_ops() + recompute_nodes = set() + for node in graph.nodes: + if node.target in no_copy_ops: + if ds_param_node in node.args: + recompute_nodes.add(node) + if any(a in recompute_nodes for a in node.args): + recompute_nodes.add(node) + + return recompute_nodes + + +def _get_values_from_ds_params(joint_graph, param_indices): + primal_inputs = list(filter(_is_primal, joint_graph.nodes)) + ds_param_inputs = [primal_inputs[arg_idx] for arg_idx, _, _ in param_indices] + + no_copy_ops = get_no_copy_ops() + + ds_param_inputs = set(ds_param_inputs) + ds_param_users = {} + + for node in joint_graph.nodes: + if node.target in no_copy_ops and any((a in ds_param_inputs or a in ds_param_users) for a in node.args): + for a in node.args: + if a in ds_param_inputs: + ds_param_users[node] = a + elif a in ds_param_users: + ds_param_users[node] = ds_param_users[a] + + return ds_param_users + + +def get_wrapped_choose_saved_values_set(param_indices: List[Tuple[int, int, torch.Size]]): + + def ds_choose_saved_values_set(joint_graph: torch.fx.Graph, node_info, memory_budget=1) -> List[Node]: + saved_values = choose_saved_values_set(joint_graph, node_info, memory_budget) + ds_param_users = _get_values_from_ds_params(joint_graph, param_indices) + + new_saved_values = [] + for v in saved_values: + if v in ds_param_users: + ds_val = ds_param_users[v] + if ds_val not in new_saved_values: + new_saved_values.append(ds_val) + else: + new_saved_values.append(v) + + return new_saved_values + + return ds_choose_saved_values_set + + +def get_wrapped_partitioner(param_indices: List[Tuple[int, int, torch.Size]]): + + def partition_recompute_ds_params(joint_module: GraphModule, _joint_inputs, *, + num_fwd_outputs) -> Tuple[GraphModule, GraphModule]: + """ + This is basically the same as the default_partition function, but + it doesn't save the gathered params and values computed from them. + """ + if has_recomputable_ops(joint_module): + return min_cut_rematerialization_partition(joint_module, _joint_inputs, num_fwd_outputs=num_fwd_outputs) + + primal_inputs = list(filter(_is_primal, joint_module.graph.nodes)) + fwd_seed_offset_inputs = list(filter(_is_fwd_seed_offset, joint_module.graph.nodes)) + inputs = primal_inputs + fwd_seed_offset_inputs + fwd_outputs, bwd_outputs = _extract_fwd_bwd_outputs(joint_module, num_fwd_outputs=num_fwd_outputs) + forward_only_graph = _extract_graph_with_inputs_outputs(joint_module.graph, inputs, fwd_outputs, "forward") + forward_node_names = {node.name for node in forward_only_graph.nodes if node.op != "output"} + saved_values = [] + saved_sym_nodes = [] + + fwd_inputs = list(filter(_is_primal, forward_only_graph.nodes)) + ds_param_inputs = [fwd_inputs[arg_idx] for arg_idx, _, _ in param_indices] + ds_param_input_names = {node.name for node in ds_param_inputs} + + ds_param_recompute_nodes = set() + + for node in joint_module.graph.nodes: + if node.name not in forward_node_names: + continue + + if is_sym_node(node): + # Symints must be kept separate from tensors so that PythonFunction only calls + # save_for_backward on tensors and stashes symints in autograd .ctx + saved_sym_nodes.append(node) + elif "tensor_meta" not in node.meta and node.op == "call_function": + # Since we can't save tuple of tensor values, we need to flatten out what we're saving + users = node.users + assert all(user.target == operator.getitem for user in users) + saved_values.extend(users) + else: + backward_usages = [n for n in node.users if n.name not in forward_node_names] + + if "tensor_meta" in node.meta and all(is_sym_node(n) for n in backward_usages): + # If we have a tensor in the forward, where only its sizes/strides are needed in the backward, + # and not the actual tensor data, + # then it will be a lot cheaper to save only the sizes/strides, and not the actual tensor. + # + # Note that saving the tensor could also cause compilation problems: + # If the user mutated an input in the forward and uses its sizes/strides in the backward, + # then we would be obligated to clone the input before saving it to appease autograd. + # (This is how we originally found this bug). + saved_sym_nodes.extend(backward_usages) + + if node.name in ds_param_input_names: + saved_values.append(node) + recompute_nodes = _find_recompute_nodes(joint_module.graph, node) + recompute_nodes = [n for n in recompute_nodes if n.name in forward_node_names] + for recompute_node in recompute_nodes: + ds_param_recompute_nodes.add(recompute_node) + + if len(recompute_nodes) > 0: + saved_values.append(node) + else: + if node not in ds_param_recompute_nodes: + saved_values.append(node) + saved_values = list(dict.fromkeys(saved_values).keys()) + saved_sym_nodes = list(dict.fromkeys(saved_sym_nodes).keys()) + + f_gm, b_gm = _extract_fwd_bwd_modules( + joint_module, + saved_values, + saved_sym_nodes=saved_sym_nodes, + num_fwd_outputs=num_fwd_outputs, + ) + + return f_gm, b_gm + + return partition_recompute_ds_params diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__init__.py b/lib/python3.12/site-packages/deepspeed/compile/passes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..620e99147647f1ab1057d16248d4b41e3cd860a5 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/__init__.py @@ -0,0 +1,48 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..profilers.graph_profile import MemoryProfilingInterpreter + +import deepspeed.comm as dist + + +def run_opt_passes(nz3, + graph_index, + graph_id, + gm, + create_inputs_fn, + opt_passes, + graph_order, + profiling_results, + param_manager, + bwd, + debug_log=False): + profile = profiling_results[graph_id] + rank = dist.get_rank() + + for i, opt_pass in enumerate(opt_passes): + + opt_pass_fn, mem_budget = opt_pass + + graph = opt_pass_fn(gm.graph, graph_id, graph_order, profiling_results, mem_budget, param_manager, bwd) + graph.lint() + gm.graph = graph + gm.recompile() + + if debug_log: + print(f"Prefetching enabled for {'bwd' if bwd else 'fwd'} graph_id={graph_id} {graph}") + + mem_prof = MemoryProfilingInterpreter(nz3, gm) + mem_prof.run(*create_inputs_fn()) + if debug_log and rank == 0: + mem_prof.dump(f"mem_prof_r{rank}_{'bwd' if bwd else 'fwd'}_{graph_index}_{graph_id}_pass_{i}.csv") + + mem = [(name, current_alloc, delta, peak) for name, current_alloc, delta, peak in mem_prof.mem_record] + if bwd: + profile.bwd_mem = mem + else: + profile.fwd_mem = mem + + return gm diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f7e9e31e333ce23bdfd6ad9c3540d80b55beef6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_activation.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_activation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7aab220c2abcfa150b0aad9a2d442ad24610af3c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_activation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_adam_states.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_adam_states.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00914d03aaed614723c944943554296fa509e630 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_adam_states.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_parameters.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_parameters.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c4f822c9728861178cdc366d467b69c95483ec3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/offload_parameters.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/prefetch.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/prefetch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..647178b1351b7c40a5a2628e061649e286b102ec Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/prefetch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/selective_gather.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/selective_gather.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e6ed5135817fe9564196586f3e1fddc5b0cd270 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/selective_gather.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/zero1_compile.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/zero1_compile.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30f4a5c4db048259867decda57ca57bcfa034993 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/zero1_compile.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/zero3_compile.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/zero3_compile.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8728c9f1bf87eee7c83b241838286220da197b07 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/passes/__pycache__/zero3_compile.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/offload_activation.py b/lib/python3.12/site-packages/deepspeed/compile/passes/offload_activation.py new file mode 100644 index 0000000000000000000000000000000000000000..496e7351f218eb103336646d1335c16d2f687d38 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/offload_activation.py @@ -0,0 +1,116 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import List, Dict, Set, Tuple +import random +from collections import defaultdict + +import torch +from torch.fx import Graph, Node + +from ..fx import get_output_node, move_primals_to_head +from ..graph_param import DSGraphParamManager + +value_to_id: Dict[int, Dict[str, int]] = defaultdict(dict) +used_ids: Set[int] = set() + + +def get_random_id() -> int: + + def _gen(): + # generate random int + return random.randint(10000, 2**31) + + global used_ids + v = _gen() + while v in used_ids: + v = _gen() + used_ids.add(v) + return v + + +def _should_offload(node: Node) -> bool: + if not hasattr(node, "meta"): + return False + if not "tensor_meta" in node.meta: + return False + + return True + + +def offload_activation_fwd(graph: Graph, graph_id: int, nodes_to_offload_with_names: List[Tuple[str, Node]], + graph_order: List[int], mem_budget: float, param_manager: DSGraphParamManager) -> Graph: + param_names = set(param_manager.param_names) + + import copy + cl_graph = copy.deepcopy(graph) + cl_graph.erase_node(get_output_node(cl_graph)) + + global value_to_id + for name, node in nodes_to_offload_with_names: + if node.name in param_names: + continue + + if not _should_offload(node): + continue + + val_id = get_random_id() + with graph.inserting_after(node): + offload_node = graph.create_node('call_function', + torch.ops.dc.offload_tensor.default, (node, graph_id, val_id), {}, + name=f"offload_{node.name}_{val_id}") + with graph.inserting_after(offload_node): + wait_node = graph.create_node('call_function', + torch.ops.dc.wait_offload.default, (offload_node, graph_id, val_id), {}, + name=f"wait_copy_{node.name}_{val_id}") + + output_node = get_output_node(graph) + output_node.replace_input_with(node, wait_node) + + value_to_id[graph_id][name] = val_id + + graph = move_primals_to_head(graph) + + graph.lint() + return graph + + +def reload_activation_bwd(graph: Graph, graph_id: int, graph_order: List[int], mem_budget: float, + param_manager: DSGraphParamManager) -> Graph: + + graph_value_to_id = value_to_id[graph_id] + name_to_node = {n.name: n for n in graph.nodes} + act_nodes = [name_to_node[n] for n in graph_value_to_id.keys()] + + node_to_first_user = {} + for act in act_nodes: + for node in graph.nodes: + if act in node.args: + node_to_first_user[act] = node + break + + for node in act_nodes: + val_id = graph_value_to_id[node.name] + + with graph.inserting_before(node_to_first_user[node]): + reload_node = graph.create_node('call_function', + torch.ops.dc.reload_tensor.default, (node, graph_id, val_id), {}, + name=f"reload_{node.name}_{val_id}") + with graph.inserting_after(reload_node): + wait_node = graph.create_node('call_function', + torch.ops.dc.wait_reload.default, (reload_node, graph_id, val_id), {}, + name=f"wait_copy_{node.name}_{val_id}") + + # replace all uses of node with wait_node + users = {} + for u in node.users.keys(): + if u != reload_node: + users[u] = (node, wait_node) + for u, (old_in, new_in) in users.items(): + u.replace_input_with(old_in, new_in) + + graph = move_primals_to_head(graph) + graph.lint() + return graph diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/offload_adam_states.py b/lib/python3.12/site-packages/deepspeed/compile/passes/offload_adam_states.py new file mode 100644 index 0000000000000000000000000000000000000000..458d07f39d161391b7bd4f47d595cb4298fb5e2c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/offload_adam_states.py @@ -0,0 +1,546 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import copy +from typing import List + +import torch +from torch.fx import Graph, GraphModule + +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.offload_states import _make_offload_state_key + +try: + from torch._subclasses.fake_tensor import unset_fake_temporarily +except ImportError: + # Unsupported torch version + pass + +from ..profilers import ProfilingResult +from ..graph_param import DSGraphParamManager +from ..fx import move_primals_to_head + +import deepspeed.comm as dist + +NAME = "offload_adam_states" + + +def print_r0(msg): + if dist.get_rank() == 0: + print(msg) + + +MARGIN = 0.2 + +copy_stream = None +offload_event = None +reload_event = None + +offload_key_events = {} +reload_key_events = {} + +max_memory = 0 + + +def lazy_init(): + global copy_stream + global offload_event + global reload_event + + if copy_stream is None: + + copy_stream = get_accelerator().Stream() + offload_event = get_accelerator().Event() + reload_event = get_accelerator().Event() + + +optimizer = None +device = None +nz3 = None + + +def move_key(state, key, key_event=None): + offload_buf_key = _make_offload_state_key(key) + if offload_buf_key not in state: + state[offload_buf_key] = get_accelerator().pin_memory(torch.empty_like(state[key], device="cpu")) + + if key not in state: + return + + with get_accelerator().stream(copy_stream): + state[offload_buf_key].copy_(state[key], non_blocking=True) + + if key_event is None: + offload_event.record(stream=copy_stream) + else: + key_event.record(stream=copy_stream) + + +def move_back_key(state, key, key_event=None): + + with get_accelerator().stream(copy_stream): + state[key] = torch.empty_like(state[_make_offload_state_key(key)], device=device) + state[key].copy_(state[_make_offload_state_key(key)], non_blocking=True) + + if key_event is None: + reload_event.record(stream=copy_stream) + else: + key_event.record(stream=copy_stream) + + +def move_hp_param(src_tensor, dest_buf, key_event=None): + with get_accelerator().stream(copy_stream): + dest_buf.copy_(src_tensor, non_blocking=True) + src_tensor.data = dest_buf + + if key_event is None: + reload_event.record(stream=copy_stream) + else: + key_event.record(stream=copy_stream) + + +def move_back_hp_param(src_tensor, dest_buf, key_event=None): + with get_accelerator().stream(copy_stream): + dest_buf.data = torch.empty_like(src_tensor, device=device) + dest_buf.copy_(src_tensor, non_blocking=True) + + if key_event is None: + reload_event.record(stream=copy_stream) + else: + key_event.record(stream=copy_stream) + + +def offload_adam_states_sync(): + + with unset_fake_temporarily(): + + if not hasattr(optimizer, "hp_params_pin_buffers"): + optimizer.hp_params_pin_buffers = [ + get_accelerator().pin_memory(torch.empty_like(t, device="cpu")) + for t in optimizer.fp32_partitioned_groups_flat + ] + + for i, (k, state) in enumerate(optimizer.state.items()): + if "exp_avg" in state: + move_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_key(state, "exp_avg_sq") + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + del state["exp_avg"] + if "exp_avg_sq" in state: + del state["exp_avg_sq"] + + for src_tensor, dest_buf in zip(optimizer.fp32_partitioned_groups_flat, optimizer.hp_params_pin_buffers): + move_hp_param(src_tensor, dest_buf) + + get_accelerator().synchronize() + + +def reload_adam_states_sync(): + + with unset_fake_temporarily(): + # print_r0("Reloading Adam states") + + for _, state in optimizer.state.items(): + if _make_offload_state_key("exp_avg") in state: + move_back_key(state, "exp_avg") + if _make_offload_state_key("exp_avg_sq") in state: + move_back_key(state, "exp_avg_sq") + + for src, dest in zip(optimizer.hp_params_pin_buffers, optimizer.fp32_partitioned_groups_flat): + move_back_hp_param(src, dest) + + get_accelerator().synchronize() + + +def sync_offload_states(event=None): + if nz3.is_profiling(): + offload_adam_states_sync() + else: + if event is None: + offload_event.wait(copy_stream) + else: + event.wait(copy_stream) + + +def sync_reload_states(event=None): + if nz3.is_profiling(): + reload_adam_states_sync() + else: + if event is None: + reload_event.wait(copy_stream) + else: + event.wait(copy_stream) + + +def make_offload_task(task): + + def run_offload_task(): + # if not nz3.is_profiling(): + # print_r0(f"run_offload_task {task[0]} {task[2]} {task[3]} {task[4]}") + + if offload_key_events.get(task[1]) is None: + offload_key_events[task[1]] = get_accelerator().Event() + + if task[2] == "hp_param": + move_hp_param(task[1][0], task[1][1], offload_key_events[task[1][0]]) + else: + assert task[1] in optimizer.state, f"State {task[1]} not found in optimizer" + state = optimizer.state[task[1]] + # if offload_key_events.get(task[1]) is None: + # offload_key_events[task[1]] = get_accelerator().Event() + move_key(state, task[2], offload_key_events[task[1]]) + + return run_offload_task + + +def make_offload_sync(task): + + def run_offload_sync(): + # if not nz3.is_profiling(): + event = offload_key_events[task[1]] + event.synchronize() + + if task[2] != "hp_param": + state = optimizer.state[task[1]] + key = task[2] + if key in state: + del state[key] + # print_r0(f"run_offload_sync {task[0]} {task[2]} alloc_mem={get_accelerator().memory_allocated()}") + + return run_offload_sync + + +def make_reload_task(task): + + def run_reload_task(): + if not nz3.is_profiling(): + if reload_key_events.get(task[1]) is None: + reload_key_events[task[1]] = get_accelerator().Event() + + if task[2] == "hp_param": + move_back_hp_param(task[1][1], task[1][0], reload_key_events[task[1]]) + else: + state = optimizer.state[task[1]] + # print_r0(f"run_reload_task {task[0]} {task[2]} {task[3]} {task[4]}") + move_back_key(state, task[2], reload_key_events[task[1]]) + + return run_reload_task + + +def update_max_memory(name): + + global max_memory + mem = get_accelerator().max_memory_allocated() + max_memory = max(max_memory, mem) + + +def empty_cache(): + get_accelerator().empty_cache() + + +offload_tasks = [] +offload_tasks_remaining = [] +offload_tasks_scheduled = [] +reload_task_remaining = [] +total_reload_mem = 0 + + +def offload_opt_states_inc(graph: Graph, graph_id: int, graph_order: List[int], profiling_results: ProfilingResult, + mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> Graph: + + to_remove = [] + for node in graph.nodes: + if node.op == 'call_function' and \ + node.target in [offload_adam_states_sync, sync_offload_states, reload_adam_states_sync, sync_reload_states, update_max_memory]: + to_remove.append(node) + + for node in to_remove: + graph.erase_node(node) + + accelerator = get_accelerator() + total_mem = accelerator.total_memory() * (1 - MARGIN) + print_r0(f"offload_opt_states_inc start graph {graph_id} bwd={bwd} max_memory={max_memory} total_mem={total_mem}") + + mem = profiling_results[graph_id].bwd_mem if bwd else profiling_results[graph_id].fwd_mem + mem_dict = {name: peak for name, alloc_mem, delta, peak in mem} + + current_peak_mem = 0 + peak_mem = {} + + ordered_node = reversed(graph.nodes) if bwd else graph.nodes + for node in ordered_node: + # print(f"Node: {node.name} mem: {mem_dict[node.name]}") + if mem_dict[node.name] > current_peak_mem: + current_peak_mem = mem_dict[node.name] + peak_mem[node.name] = current_peak_mem + + # fwd_max_mem = max(m[3] for m in prof.fwd_mem) + # bwd_max_mem = max(m[3] for m in prof.bwd_mem) if len(prof.bwd_mem) > 0 else 0 + # peak_mem = max(peak_mem, fwd_max_mem, bwd_max_mem) + + global offload_tasks_remaining, reload_tasks_remaining, offload_tasks_scheduled + + if not bwd: + is_first_graph = graph_id == graph_order[0][0] + # print_r0( + # f"offload_opt_states_inc start graph {graph_id} graph_order {graph_order} fwd is_first_graph {is_first_graph}" + # ) + + # At the beginning of the first graph, we schedule offload tasks to launch all offloading + if is_first_graph: + # print_r0( + # f"offload_opt_states_inc fwd before reload graph {graph_id} allocated_mem={get_accelerator().memory_allocated()}" + # ) + + with unset_fake_temporarily(): + offload_adam_states_sync() + reload_adam_states_sync() + sync_reload_states() + + reload_size = 0 + + for i, ((k, state), hp_param, hp_param_cpu) in enumerate( + zip(optimizer.state.items(), optimizer.fp32_partitioned_groups_flat, + optimizer.hp_params_pin_buffers)): + # print_r0( + # f"Checking key for offloading {i} {k.shape} has_key {_make_offload_state_key('exp_avg') in state}") + + if _make_offload_state_key("exp_avg") in state: + key = _make_offload_state_key("exp_avg") + size = state[key].numel() * state[key].element_size() + + # if total_mem < max_memory + reload_size + size: + offload_tasks.append( + (i, k, "exp_avg", state[key].numel() * state[key].element_size(), state[key].dtype)) + # print_r0( + # f"Offloading task {i} exp_avg reload_size={reload_size} size={size} estimated_mem={max_memory + reload_size + size}" + # ) + + if _make_offload_state_key("exp_avg_sq") in state: + key = _make_offload_state_key("exp_avg_sq") + size = state[key].numel() * state[key].element_size() + + # if total_mem < max_memory + reload_size + size: + offload_tasks.append( + (i, k, "exp_avg_sq", state[key].numel() * state[key].element_size(), state[key].dtype)) + # print_r0( + # f"Offloading task {i} exp_avg_sq reload_size={reload_size} size={size} estimated_mem={max_memory + reload_size + size}" + # ) + + hp_param_size = hp_param.numel() * hp_param.element_size() + # if total_mem < max_memory + reload_size + hp_param_size: + offload_tasks.append((i, (hp_param, hp_param_cpu), "hp_param", + hp_param.numel() * hp_param.element_size(), hp_param.dtype)) + # print_r0( + # f"Offloading task {i} hp_param reload_size={reload_size} size={hp_param_size} estimated_mem={max_memory + reload_size + hp_param_size}" + # ) + + # print_r0(f"offload_opt_states_inc fwd graph {graph_id} allocated_mem={get_accelerator().memory_allocated()}") + + for node in graph.nodes: + # print_r0(f"checking sync node insert node: {node.name}") + + if node.name not in peak_mem \ + or node.op == 'placeholder' \ + or "offload_opt_" in node.name: + continue + + to_offload = [] + optim_size = sum([task[3] for task in offload_tasks]) + + # print_r0( + # f" optim_size: {optim_size} total_mem: {total_mem} peak_mem: {peak_mem[node.name]} available: {total_mem - peak_mem[node.name] - optim_size} #tasks={len(offload_tasks)}" + # ) + while total_mem - peak_mem[node.name] - optim_size < 0: + if len(offload_tasks) == 0: + break + + task = offload_tasks.pop(0) + to_offload.append(task) + optim_size = sum([task[3] for task in offload_tasks]) + # print_r0( + # f" scheduled task {task[0]} {task[2]} {task[3]} optim_size: {optim_size} peak_mem: {peak_mem[node.name]} available: {total_mem - peak_mem[node.name] - optim_size} #tasks={len(offload_tasks)}" + # ) + + for task in to_offload: + with graph.inserting_before(node): + graph.create_node('call_function', + make_offload_sync(task), (), {}, + name=f"offload_opt_sync_{task[0]}_{task[2]}") + print_r0(f"Inserting fwd offload_opt_sync_{task[0]}_{task[2]}") + offload_tasks_scheduled.append(task) + + for node in graph.nodes: + # print(f"Node: {node.name} mem: {mem_dict[node.name]}") + if node.op != 'placeholder': + print_r0(f"Inserting all offload tasks before {node.name}") + for task in offload_tasks_scheduled: + name = f"offload_opt_{task[0]}_{task[2]}" + with graph.inserting_before(node): + offload_node = graph.create_node('call_function', make_offload_task(task), (), {}, name=name) + break + + # print_r0(f"offload_opt_states_inc finish graph {graph_id} fwd graph {graph}") + print_r0(f"offload_opt_states_inc finish graph {graph_id}") + else: + + graph_order_with_backward = [g[0] for g in graph_order if g[1]] + is_first_graph = graph_id == graph_order_with_backward[-1] + is_last_graph = graph_id == graph_order_with_backward[0] + + # print_r0( + # f"offload_opt_states_inc bwd graph {graph_id} graph_order_with_backward {graph_order_with_backward} is_first_graph {is_first_graph} is_last_graph {is_last_graph}" + # ) + + if is_first_graph: + inserted_sync = False + for node in graph.nodes: + if node.op != 'placeholder' and not inserted_sync: + # print(f"Inserting offload_sync before {node.name}") + with graph.inserting_before(node): + graph.create_node('call_function', empty_cache, (), {}, name="empty_cache") + + inserted_sync = True + reload_tasks_remaining = copy.copy(offload_tasks_scheduled) + + global total_reload_mem + for node in graph.nodes: + if node.name not in peak_mem \ + or node.op == 'placeholder' \ + or node.op == 'output' \ + or "offload_opt_sync_" in node.name: + continue + + if len(reload_tasks_remaining) > 0: + task = reload_tasks_remaining[0] + next_reload_mem = task[3] + + insert_pos = node + while total_mem > peak_mem[node.name] + total_reload_mem + next_reload_mem: + expected_mem = peak_mem[node.name] + total_reload_mem + print_r0( + f" Inserting reload_opt reload_opt_{task[0]}_{task[2]} after {insert_pos.name} next_inc={next_reload_mem} peak_mem[{node.name}]={peak_mem[node.name]} inc_total={total_reload_mem} expected_mem={expected_mem}" + ) + + with graph.inserting_after(insert_pos): + insert_pos = graph.create_node('call_function', + make_reload_task(task), (), {}, + name=f"reload_opt_{task[0]}_{task[2]}") + + total_reload_mem += next_reload_mem + reload_tasks_remaining.pop(0) + if len(reload_tasks_remaining) == 0: + break + + task = reload_tasks_remaining[0] + next_reload_mem = task[3] + + # prev_node = node + + if is_last_graph: + for node in graph.nodes: + # print(f"Node: {node.name} mem: {mem_dict[node.name]}") + if node.op == 'output': + for task in reload_tasks_remaining: + with graph.inserting_before(node): + graph.create_node('call_function', + make_reload_task(task), (), {}, + name=f"reload_opt_{task[0]}_{task[2]}") + + sync_fn = lambda: copy_stream.synchronize() + with graph.inserting_before(node): + graph.create_node('call_function', sync_fn, (), {}, name="sync_offload_copy_stream") + + print_r0( + f"offload_opt_states_inc graph {graph_id} graph_order {graph_order} bwd is_first_graph {is_first_graph} is_last_graph {is_last_graph}" + ) + + return graph + + +def add_record_max_mem_nodes(graph: Graph): + + nodes = list(graph.nodes) + for node in nodes: + if node.op == "output" or node.op == "placeholder": + continue + + with graph.inserting_after(node): + name = f"update_max_memory_{node.name}" + graph.create_node('call_function', update_max_memory, (name, ), {}, name=name) + + +def insert_offload_opt_states(graph: Graph, graph_id: int, graph_order: List[int], profiling_results: ProfilingResult, + mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> Graph: + + if bwd: + graph_order_with_backward = [g[0] for g in graph_order if g[1]] + is_last_graph = graph_id == graph_order_with_backward[0] + + inserted_reload = False + for node in graph.nodes: + # print(f"Node: {node.name} mem: {mem_dict[node.name]}") + if node.op == 'output' and not inserted_reload and is_last_graph: + # print(f"Inserting reload_opt before {node.name}") + with graph.inserting_before(node): + graph.create_node('call_function', reload_adam_states_sync, (), {}, name="reload_opt") + inserted_reload = True + + # add_record_max_mem_nodes(graph) + + else: + is_first_graph = graph_id == graph_order[0][0] + + graph = move_primals_to_head(graph) + + inserted_offload = False + for node in graph.nodes: + # print(f"Node: {node.name} mem: {mem_dict[node.name]}") + if node.op != 'placeholder' and not inserted_offload and is_first_graph: + print(f"Inserting offload_opt before {node.name}") + with graph.inserting_before(node): + graph.create_node('call_function', offload_adam_states_sync, (), {}, name="offload_opt") + inserted_offload = True + + add_record_max_mem_nodes(graph) + + return graph + + +def move_opt_states(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, create_inputs_fn, + mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> GraphModule: + gm.graph = offload_opt_states_inc(gm.graph, graph_id, graph_order, profiling_results, mem_budget, param_manager, + bwd) + return gm + + +def move_opt_states_sync(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, create_inputs_fn, + mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> GraphModule: + gm.graph = insert_offload_opt_states(gm.graph, graph_id, graph_order, profiling_results, mem_budget, param_manager, + bwd) + return gm + + +def offload_adam_states_for_init(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, + create_inputs_fn, mem_budget: float, param_manager: DSGraphParamManager, + bwd: bool) -> GraphModule: + if not bwd and graph_id == graph_order[0][0]: + with unset_fake_temporarily(): + offload_adam_states_sync() + # returns None, and profiling will be skipped + + +def init_offload_opt_states(adam_optimizer, _nz3): + lazy_init() + + global optimizer + optimizer = adam_optimizer + global device + device = torch.device(get_accelerator().current_device()) + global nz3 + nz3 = _nz3 diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/offload_parameters.py b/lib/python3.12/site-packages/deepspeed/compile/passes/offload_parameters.py new file mode 100644 index 0000000000000000000000000000000000000000..29468f4970d90a4e44b505acbaaa3fe5e2d613fc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/offload_parameters.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import List + +import torch +from torch.fx import Node, GraphModule +from deepspeed.compile.util import get_last_uses +from ..graph_param import DSGraphParamManager + + +def add_offload_parameter(graph_id: int, gm: GraphModule, node: Node, ds_id: int): + new_node = None + with gm.graph.inserting_after(node): + args = (node, ) + for a in [graph_id, ds_id]: # To add ds_id + args += (a, ) + new_node = gm.graph.create_node('call_function', + torch.ops.dc.offload_parameter.default, + args, {}, + name="offload_parameter") + + return new_node + + +def add_reload_parameter(graph_id: int, gm: GraphModule, node: Node, ds_id: int): + new_node = None + with gm.graph.inserting_after(node): + args = (node, ) + for a in [graph_id, ds_id]: # To add ds_id + args += (a, ) + new_node = gm.graph.create_node('call_function', + torch.ops.dc.reload_parameter.default, + args, {}, + name=f"reload_parameter") + return new_node + + +def get_ds_id(node: Node): + assert node.target == torch.ops.dc.allgather_param.default + return node.args[2] + + +def offload_parameter_fwd(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, create_inputs_fn, + mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> GraphModule: + node_to_last_use, user_to_last_uses = get_last_uses(gm.graph) + for node in gm.graph.nodes: + if (isinstance(node, Node) and node.target == torch.ops.dc.allgather_param.default): + add_reload_parameter(graph_id, gm, node.args[0], get_ds_id(node)) + add_offload_parameter(graph_id, gm, node_to_last_use[node], get_ds_id(node)) + gm.graph.lint() + return gm diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/prefetch.py b/lib/python3.12/site-packages/deepspeed/compile/passes/prefetch.py new file mode 100644 index 0000000000000000000000000000000000000000..ce0d721f8d5800442e70ab772c0f2eafc3f15479 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/prefetch.py @@ -0,0 +1,174 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import List + +import torch +from torch.fx import Graph, Node, GraphModule + +from deepspeed.accelerator import get_accelerator +import deepspeed.comm as dist + +from ..profilers.comm_profile import create_predictor +from ..graph_param import DSGraphParamManager + +NAME = "prefetch" + +FUSE_FACTOR = 0.8 +MARGIN = 0.1 +MAX_FUSE_SIZE = 1e9 +MAX_BUFFERED_SIZE = 4e9 + +run_prefetch_pass = False + + +def print_rank_0(message): + if dist.get_rank() == 0: + print(message) + + +def get_ds_id(node: Node): + assert node.target == torch.ops.dc.allgather_param.default + return node.args[2] + + +def schedule_prefetch(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, create_inputs_fn, + mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> GraphModule: + + max_mem = get_accelerator().total_memory() * (1 - MARGIN) + vals_to_bcast = torch.tensor([max_mem], device=torch.device(get_accelerator().current_device())) + dist.all_reduce(vals_to_bcast, dist.ReduceOp.MIN) + max_mem = vals_to_bcast[0].item() + + mem = profiling_results[graph_id].bwd_mem if bwd else profiling_results[graph_id].fwd_mem + op_time = profiling_results[graph_id].bwd_time if bwd else profiling_results[graph_id].fwd_time + tensor_sizes = profiling_results[graph_id].bwd_tensor_sizes if bwd else profiling_results[graph_id].fwd_tensor_sizes + + mem_dict = {name: (alloc_mem, peak) for name, alloc_mem, delta, peak in mem} + time_dict = {name: (device_time, wall_time) for name, device_time, wall_time in op_time} + tensor_size_dict = {name: size for name, size in tensor_sizes} + + graph = gm.graph + total_param_size = sum( + [tensor_size_dict[n.name] for n in graph.nodes if n.target == torch.ops.dc.allgather_param.default]) + + print_rank_0( + f"schedule_prefetch graph_id={graph_id} max_mem={max_mem} available_memory={get_accelerator().available_memory()} memory_allocated={get_accelerator().memory_allocated()} max_allocated={get_accelerator().max_memory_allocated()} total_param_size={total_param_size} margin={MARGIN}" + ) + + # Fill missing values + prev_mem = 0 + prev_peak = 0 + for node in graph.nodes: + if node.name in mem_dict: + prev_mem = mem_dict[node.name][0] + prev_peak = mem_dict[node.name][1] + else: + print_rank_0(f"node {node.name} not in mem_dict") + mem_dict[node.name] = (prev_mem, prev_peak) + + comm_predictor = create_predictor() + + order_rev = list(reversed(graph.nodes)) + new_order_rev = [] + prefetch_ags = [] + prefetch_ag_groups = [] + ag_tensor_size_sum = 0 + for i, node in enumerate(order_rev): + # print_rank_0( + # f"Checking node reverse order {node.name} {node.target} ag_tensor_size_sum={ag_tensor_size_sum} max_mem={max_mem}" + # ) + + if node.op != "placeholder": + assert i < len(order_rev) - 1 + assert node.name in mem_dict + next_node = order_rev[i + 1] + next_alloc_mem, next_peak = mem_dict[next_node.name] + + # Free up memory + while next_peak + ag_tensor_size_sum > max_mem or ag_tensor_size_sum > MAX_BUFFERED_SIZE: + if len(prefetch_ag_groups) > 0: + # launch prefetch + fused_ag_nodes = prefetch_ag_groups.pop(0) + total_ag_tensor_size = sum([tensor_size_dict[ag_node.name] for ag_node in fused_ag_nodes]) + ag_tensor_size_sum -= total_ag_tensor_size + new_order_rev.append(fused_ag_nodes) + assert len(fused_ag_nodes) > 0 + # print_rank_0( + # f"Free up memory fused_ag_nodes={fused_ag_nodes} next_alloc_mem={next_alloc_mem} total_ag_tensor_size={total_ag_tensor_size} ag_tensor_size_sum={ag_tensor_size_sum} max_mem={max_mem}" + # ) + elif len(prefetch_ags) > 0: + prefetch_ag_groups.append(prefetch_ags) + prefetch_ags = [] + # print_rank_0( + # f"Free up memory prefetch_ags={prefetch_ag_groups} next_alloc_mem={next_alloc_mem} ag_tensor_size_sum={ag_tensor_size_sum} max_mem={max_mem}" + # ) + else: + break + + if node.target == torch.ops.dc.allgather_param.default: + + current_ag_size = sum([tensor_size_dict[ag_node.name] for ag_node in prefetch_ags]) + pred_time_current = comm_predictor(current_ag_size) + pred_time_next = comm_predictor(tensor_size_dict[node.name]) + pred_time_fused = comm_predictor(current_ag_size + tensor_size_dict[node.name]) + + do_fuse = max(pred_time_current, pred_time_next) * 1.2 > pred_time_fused and ( + current_ag_size + tensor_size_dict[node.name]) < MAX_FUSE_SIZE + # print_rank_0( + # f"found allgather_param do_fuse={do_fuse} current_ag_size={current_ag_size} tensor_size_dict[node.name]={tensor_size_dict[node.name]} pred_time_current={pred_time_current} pred_time_next={pred_time_next} pred_time_fused={pred_time_fused}" + # ) + + if len(prefetch_ags) > 0 and not do_fuse: + # stop fusing here + prefetch_ag_groups.append(prefetch_ags) + prefetch_ags = [] + # print_rank_0( + # f"stop fusing prefetch_ags={prefetch_ag_groups} ag_tensor_size_sum={ag_tensor_size_sum}") + # else: + # print_rank_0( + # f"continue fusing ag_tensor_size_sum={ag_tensor_size_sum} ag_size={tensor_size_dict[node.name]} prefetch_ags={prefetch_ags} prefetch_ag_groups={prefetch_ag_groups}" + # ) + prefetch_ags.append(node) + ag_tensor_size_sum += tensor_size_dict[node.name] + + new_order_rev.append(node) + + if (node.op != "placeholder" + and node.target != torch.ops.dc.reload_parameter) and order_rev[i + 1].op == "placeholder": + for ag_group in prefetch_ag_groups: + assert len(ag_group) > 0 + new_order_rev.append(ag_group) + total_ag_tensor_size = sum([tensor_size_dict[ag_node.name] for ag_node in ag_group]) + ag_tensor_size_sum -= total_ag_tensor_size + if len(prefetch_ags) > 0: + new_order_rev.append(prefetch_ags) + ag_tensor_size_sum -= sum([tensor_size_dict[ag_node.name] for ag_node in prefetch_ags]) + assert ag_tensor_size_sum == 0 + + # print_rank_0( + # f"node={node} next_alloc_mem={next_alloc_mem} pending_ags={len(prefetch_ags)} ag_tensor_size_sum={ag_tensor_size_sum}" + # ) + + assert ag_tensor_size_sum >= 0 + + new_graph = Graph() + env = {} + for node in reversed(new_order_rev): + if isinstance(node, Node): + #print(f"reconstruct {node.name} {node.target}") + new_node = new_graph.node_copy(node, lambda n: env[n.name]) + env[node.name] = new_node + else: + param_nodes = [ag_node.args[0] for ag_node in node] + param_nodes_copy = [env[param_node.name] for param_node in param_nodes] + + ds_ids = [get_ds_id(ag_node) for ag_node in node] + new_graph.call_function(torch.ops.dc.prefetch_params_fused.default, + args=(graph_id, param_nodes_copy, ds_ids)) + new_graph.lint() + gm.graph = new_graph + + return gm diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/selective_gather.py b/lib/python3.12/site-packages/deepspeed/compile/passes/selective_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..83306872ce062febe68a399c53bc14b584086586 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/selective_gather.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from collections import defaultdict +from typing import List + +import torch +from torch.fx import GraphModule + +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator + +from ..util import get_deepcompile_handle +from ..graph_param import DSGraphParamManager + +NAME = "selective_gather" + +max_alloc_mem = 0 +last_optimize_step = 0 + + +def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, create_inputs_fn, + mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> GraphModule: + + if not bwd: + return gm + + last_backward_graph_id = None + for g_id, needs_bwd in graph_order: + if needs_bwd: + last_backward_graph_id = g_id + break + + # Run only on the last backward graph + if last_backward_graph_id is None or graph_id != last_backward_graph_id: + return gm + + peak_mem = 0 + for graph_id, prof in profiling_results.items(): + # Use peak memory + fwd_max_mem = max(m[3] for m in prof.fwd_mem) + bwd_max_mem = max(m[3] for m in prof.bwd_mem) if len(prof.bwd_mem) > 0 else 0 + peak_mem = max(peak_mem, fwd_max_mem, bwd_max_mem) + if dist.get_rank() == 0: + print( + f"selective_gather graph_id={graph_id} max_mem={peak_mem} fwd_max_mem={fwd_max_mem} bwd_max_mem={bwd_max_mem}" + ) + + persistent_ds_ids = set() + for graph_id, pm in param_manager.items(): + for name, ds_param in pm.params.items(): + if ds_param.param.ds_persist: + persistent_ds_ids.add(pm.ds_ids[name]) + + ds_id_to_size = {} + ds_id_to_time = defaultdict(float) + ds_id_to_prof_dtime = defaultdict(float) + ds_id_to_prof_wtime = defaultdict(float) + + for graph_id, pm in param_manager.items(): + params = pm.params + for param_name, param in params.items(): + ds_id = pm.ds_ids[param_name] + ds_id_to_size[ds_id] = param.numel * param.dtype.itemsize + + profile = profiling_results[graph_id] + for n in profile.fwd_graph.nodes: + if n.target == torch.ops.dc.allgather_param.default: + assert "tensor_size" in n.meta + ds_id_to_size[n.args[2]] = n.meta["tensor_size"] + assert "device_time" in n.meta + ds_id_to_time[n.args[2]] += n.meta["device_time"] + + ds_id_to_prof_dtime[n.args[2]] = n.meta["device_time"] + ds_id_to_prof_wtime[n.args[2]] = n.meta["wall_time"] + + if profile.bwd_graph is not None: + for n in profile.bwd_graph.nodes: + if n.target == torch.ops.dc.allgather_param.default: + assert "tensor_size" in n.meta + ds_id_to_size[n.args[2]] = n.meta["tensor_size"] + assert "device_time" in n.meta + ds_id_to_time[n.args[2]] += n.meta["device_time"] + + ds_ids = [ds_id for ds_id in ds_id_to_size if ds_id not in persistent_ds_ids] + ds_ids.sort(key=lambda ds_id: ds_id_to_time[ds_id] / ds_id_to_size[ds_id], reverse=True) + + # print(f"ds_id_to_size={ds_id_to_size}") + # print(f"ds_id_to_time={ds_id_to_time}") + + # if dist.get_rank() == 0: + # for ds_id in ds_ids: + # dtime_in_sec = ds_id_to_prof_dtime[ds_id] + # wtime_in_sec = ds_id_to_prof_wtime[ds_id] + # size_in_mb = ds_id_to_size[ds_id] / 1024 / 1024 + # print( + # f"ds_id={ds_id} time_per_size={ds_id_to_time[ds_id] / ds_id_to_size[ds_id]:.5f} dtime={dtime_in_sec:.3f} wtime={wtime_in_sec:.3f} size={size_in_mb:.2f}MB bw={size_in_mb/dtime_in_sec:.2f}MB/s" + # ) + + sorted_ds_ids = {ds_id: ds_id_to_size[ds_id] for ds_id in ds_ids} + + accelerator = get_accelerator() + total_mem = accelerator.total_memory() + vals_to_bcast = torch.tensor([total_mem], device=torch.device(get_accelerator().current_device())) + dist.all_reduce(vals_to_bcast, dist.ReduceOp.MIN) + total_mem = vals_to_bcast[0].item() + + MEM_MARGIN = 0.1 + available_mem = total_mem * (1 - MEM_MARGIN) - peak_mem + + if dist.get_rank() == 0: + print( + f"selective_gather max_mem={peak_mem} total_mem={total_mem} MEM_MARGIN={MEM_MARGIN} available_mem={available_mem}" + ) + + ds_id_to_param = {} + for g_id, g_pm in param_manager.items(): + for name, ds_param in g_pm.params.items(): + ds_id_to_param[g_pm.ds_ids[name]] = ds_param.param + + persistent_mem = 0 + nz3 = get_deepcompile_handle() + for ds_id, size in sorted_ds_ids.items(): + if persistent_mem + size > available_mem: + break + persistent_mem += size + + param_obj = ds_id_to_param[ds_id] + + nz3.set_persistent(ds_id) + if dist.get_rank() == 0: + print(f"Set persistent: {ds_id} size: {size} persistent_mem: {persistent_mem} shape: {param_obj.ds_shape}") + + return gm + + +# def make_selective_gather(z3_optimizer, nz3): + +# def selective_gather_wrapper(graph: Graph, graph_id: int, graph_order: List[int], profiling_results, +# mem_budget: float, param_manager, bwd: bool) -> Graph: +# return selective_gather(graph, graph_id, graph_order, profiling_results, mem_budget, param_manager, bwd, +# z3_optimizer, nz3) + +# return selective_gather_wrapper diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/zero1_compile.py b/lib/python3.12/site-packages/deepspeed/compile/passes/zero1_compile.py new file mode 100644 index 0000000000000000000000000000000000000000..fb331cc6bca32b2ed725c09ff86aec1f45650980 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/zero1_compile.py @@ -0,0 +1,55 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import List + +import torch +from torch.fx import GraphModule + +from ..util import get_deepcompile_handle +from ..fx import add_postprocess, move_primals_to_head, _make_node_meta + +NAME = "zero1_compile" + + +def add_z1_reduce_fw(gm: GraphModule, graph_id: int, profiling_results, param_manager) -> GraphModule: + + dc = get_deepcompile_handle() + param_indices = profiling_results[graph_id].param_indices + dc.register_graph_z1(graph_id, [v[1] for v in param_indices]) # Need this before profiling + + return gm + + +def add_z1_reduce_bw(gm: GraphModule, graph_id: int, param_manager) -> GraphModule: + + graph = gm.graph + pm = param_manager[graph_id] + _, param_name_to_grad = pm.get_bwd_mapping(graph) + + for param_name in pm.param_names: + + grad_node = param_name_to_grad[param_name] + + assert param_name in pm.ds_ids, f"param_name={param_name} not in ds_ids" + ds_id = pm.ds_ids[param_name] + + new_node = add_postprocess(graph, + grad_node, + torch.ops.dc.reduce_grad.default, + extra_args=[graph_id, ds_id], + name=f"reduce_param_{param_name}", + meta=_make_node_meta(grad_node, param_name, True)) + new_node.meta["val"] = None + + gm.graph = move_primals_to_head(graph) + return gm + + +def add_z1_reduce(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, create_inputs_fn, + mem_budget: float, param_manager, bwd: bool) -> GraphModule: + if bwd: + return add_z1_reduce_bw(gm, graph_id, param_manager) + return add_z1_reduce_fw(gm, graph_id, profiling_results, param_manager) diff --git a/lib/python3.12/site-packages/deepspeed/compile/passes/zero3_compile.py b/lib/python3.12/site-packages/deepspeed/compile/passes/zero3_compile.py new file mode 100644 index 0000000000000000000000000000000000000000..1fe420081bd0d10f1780d1aa8349534e7530dfa7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/passes/zero3_compile.py @@ -0,0 +1,191 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import gc +from typing import List, Dict + +import torch +from torch.fx import Graph, Node, GraphModule + +from ..util import get_input_nodes, get_param_nodes, get_index_by_graph_id, get_deepcompile_handle, get_real_uses +from ..fx import add_postprocess, _make_node_meta, get_output_node, move_primals_to_head +from ..profilers.graph_profile import ProfilingInterpreter +from ..list_schedule import fast_free_schedule + +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator + +NAME = "zero3_compile" + + +def add_allgather(graph_id: int, graph: Graph, node: Node, ds_id: int): + new_ag_node = add_postprocess(graph, + node, + torch.ops.dc.allgather_param.default, + extra_args=[graph_id, ds_id], + name=f"allgather_ds_param_{node.target}_{ds_id}", + meta=_make_node_meta(node, ds_id, True)) + new_ag_node.meta["val"] = node.meta["val"] + + # Set the previous node back to output + # We don't want to change the output node to allgather + output_node = get_output_node(graph) + output_node.replace_input_with(new_ag_node, node) + + # Add wait as well + new_wait_node = add_postprocess(graph, + new_ag_node, + torch.ops.dc.wait_allgather.default, + extra_args=[graph_id, ds_id], + name=f"wait_allgather_ds_param__{node.target}_{ds_id}", + meta=_make_node_meta(node, ds_id, False)) + new_wait_node.meta["val"] = node.meta["val"] + + return new_ag_node + + +def add_release(graph_id: int, graph: Graph, node: Node, release_node: Node, ds_id: int, n_users: int): + new_node = add_postprocess(graph, + node, + torch.ops.dc.release_param.default, + extra_args=[graph_id, ds_id, n_users], + name=f"release_ds_param_{release_node.target}_{node.name}_{ds_id}", + meta=_make_node_meta(node, ds_id, False)) + new_node.meta["val"] = None + + +def add_reduce(graph_id: int, graph: Graph, grad_node: Node, param_name: str, ds_id: int): + new_node = add_postprocess(graph, + grad_node, + torch.ops.dc.reduce_grad.default, + extra_args=[graph_id, ds_id], + name=f"reduce_ds_param_{param_name}", + meta=_make_node_meta(grad_node, ds_id, True)) + new_node.meta["val"] = None + + +def add_gather_and_release(graph_id: int, graph: Graph, param_manager, param_nodes: List[Node]) -> Graph: + + node_to_uses = get_real_uses(graph) + for pn in param_nodes: + add_allgather(graph_id, graph, pn, param_manager.ds_ids[pn.name]) + ds_id = param_manager.ds_ids[pn.name] + users = node_to_uses[pn] + for user in users: + add_release(graph_id, graph, user, pn, ds_id, len(users)) + + return move_primals_to_head(graph) + + +def add_gather_and_reduce(graph_id: int, graph: Graph, param_manager, param_nodes_bw: List[Node], + param_name_to_grad: Dict[str, Node]) -> Graph: + + add_gather_and_release(graph_id, graph, param_manager, param_nodes_bw) + + for param_name in param_manager.param_names: + add_reduce(graph_id, graph, param_name_to_grad[param_name], param_name, param_manager.ds_ids[param_name]) + + return move_primals_to_head(graph) + + +def add_z3_gather_release_fw(gm: GraphModule, + graph_id: int, + graph_order: List[int], + profiling_results, + create_inputs_fn, + param_manager, + debug_log=False) -> GraphModule: + + nz3 = get_deepcompile_handle() + + real_inputs = create_inputs_fn() + param_indices = profiling_results[graph_id].param_indices + + gm.graph = add_gather_and_release(graph_id, gm.graph, param_manager[graph_id], + get_param_nodes(gm.graph, param_indices)) + + nz3.register_graph_z3(graph_id, [v[1] for v in param_indices]) # Need this before profiling + + profiler = ProfilingInterpreter(gm, debug_log=debug_log) + profiler.run(*real_inputs) + del profiler + gc.collect() + get_accelerator().empty_cache() + + rank = dist.get_rank() + graph_index = get_index_by_graph_id(graph_order, graph_id) + if rank == 0 and debug_log: + print(f"Fwd before scheduling graph {graph_index} graph_id={graph_id} {gm.graph}") + + for n in gm.graph.nodes: + is_ds_param = n.name in param_manager[graph_id].ds_ids + if "val" in n.meta and is_ds_param: + # Used for Inductor's validation + n.meta["val"] = torch.empty([0], dtype=n.meta['val'].dtype, device=n.meta['val'].device) + + gm.graph = fast_free_schedule( + gm.graph, + get_accelerator().available_memory(), + 0, # unused + debug_log=debug_log) + + if rank == 0 and debug_log: + print(f"Fwd after scheduling graph {graph_index} graph_id={graph_id} {gm.graph}") + + return gm + + +def add_z3_gather_release_bw(gm: GraphModule, + graph_id: int, + graph_order: List[int], + profiling_results, + create_inputs_fn, + param_manager, + debug_log=False) -> GraphModule: + + param_nodes_bw, param_name_to_grad = param_manager[graph_id].get_bwd_mapping(gm.graph) + gm.graph = add_gather_and_reduce(graph_id, gm.graph, param_manager[graph_id], param_nodes_bw, param_name_to_grad) + + input_nodes = get_input_nodes(gm.graph) + real_inputs = create_inputs_fn() + assert len(input_nodes) == len(real_inputs), f"Expected {len(real_inputs)} inputs, got {len(input_nodes)}" + + real_outputs = ProfilingInterpreter(gm, debug_log=debug_log).run(*real_inputs) + + del real_outputs + gc.collect() + get_accelerator().empty_cache() + + rank = dist.get_rank() + graph_index = get_index_by_graph_id(graph_order, graph_id) + if rank == 0 and debug_log: + print(f"Bwd before scheduling graph {graph_index} graph_id={graph_id} {gm.graph}") + + gm.graph = fast_free_schedule( + gm.graph, + get_accelerator().available_memory(), + 0, # unused + debug_log=debug_log) + + return gm + + +def add_z3_gather_release(gm: GraphModule, graph_id: int, graph_order: List[int], profiling_results, create_inputs_fn, + mem_budget: float, param_manager, bwd: bool) -> GraphModule: + if bwd: + return add_z3_gather_release_bw(gm, + graph_id, + graph_order, + profiling_results, + create_inputs_fn, + param_manager, + debug_log=False) + return add_z3_gather_release_fw(gm, + graph_id, + graph_order, + profiling_results, + create_inputs_fn, + param_manager, + debug_log=False) diff --git a/lib/python3.12/site-packages/deepspeed/compile/patch_compiled_func.py b/lib/python3.12/site-packages/deepspeed/compile/patch_compiled_func.py new file mode 100644 index 0000000000000000000000000000000000000000..c77d529a64ac42cfe30b54a7a5422b23b27d05c3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/patch_compiled_func.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.utils.torch import required_torch_version + +backward_inputs = [] + +enabled_patched_func = False +original_grad_fn = None +base_meta = type(torch.autograd.Function) + +if required_torch_version(min_version=2.7): + + class FunctionMeta(base_meta): + + def __new__(cls, name, bases, dct): + if name == "CompiledFunction": + original_backward_impl = dct.get("_backward_impl") + + def wrapped_backward_impl(ctx, all_args): + assert original_backward_impl is not None + + if enabled_patched_func: + backward_inputs.append(all_args) + wrapped_backward_impl.owner_class.compiled_bw = None + + return original_backward_impl(ctx, all_args) + + wrapped_backward_impl.owner_class = None + dct["_backward_impl"] = staticmethod(wrapped_backward_impl) + new_class = super().__new__(cls, name, bases, dct) + wrapped_backward_impl.owner_class = new_class + + return new_class + + return super().__new__(cls, name, bases, dct) + +elif required_torch_version(min_version=2.6): + + class FunctionMeta(base_meta): + + def __new__(cls, name, bases, dct): + if name == "CompiledFunction": + original_backward_prologue = dct.get("_backward_prologue") + + def wrapped_backward_prologue(ctx, *grad_outputs): + assert original_backward_prologue is not None + + all_args = original_backward_prologue(ctx, *grad_outputs) + if enabled_patched_func: + backward_inputs.append(all_args) + wrapped_backward_prologue.owner_class.compiled_bw = None + + return all_args + + wrapped_backward_prologue.owner_class = None + dct["_backward_prologue"] = staticmethod(wrapped_backward_prologue) + new_class = super().__new__(cls, name, bases, dct) + wrapped_backward_prologue.owner_class = new_class + + return new_class + + return super().__new__(cls, name, bases, dct) + + +def patch_compiled_func(): + + global enabled_patched_func + enabled_patched_func = True + + class PatchedFunction(torch.autograd.Function, metaclass=FunctionMeta): + pass + + global original_grad_fn + original_grad_fn = torch.autograd.Function + torch.autograd.Function = PatchedFunction + + return backward_inputs + + +def unpatch_compiled_func(): + global enabled_patched_func + enabled_patched_func = False + + global original_grad_fn + torch.autograd.Function = original_grad_fn + + +def get_backward_inputs(): + return backward_inputs diff --git a/lib/python3.12/site-packages/deepspeed/compile/patch_fake_tensor.py b/lib/python3.12/site-packages/deepspeed/compile/patch_fake_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..1924e886192163c97232bf1c3f450e8fe4e0b28e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/patch_fake_tensor.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +try: + from torch._subclasses import FakeTensorMode + from torch._subclasses.fake_tensor import unset_fake_temporarily + from torch._dynamo.variables.builder import wrap_to_fake_tensor_and_record +except ImportError: + # Unsupported torch version + pass + + +def wrap_if_ds_param(t): + if hasattr(t, 'ds_id'): + data = torch.rand(t.ds_shape, + dtype=t.dtype, + layout=t.layout, + device=t.device, + pin_memory=t.is_pinned(), + requires_grad=t.requires_grad) + if isinstance(t, torch.nn.Parameter): + t = torch.nn.Parameter(data, requires_grad=t.requires_grad) + else: + t = data + return t + + +def patch_fake_tensor(): + # dynamo tracer uses wrap_to_fake_tensor_and_record + # Wrapping FakeTensorMode.from_tensor is not sufficient as dynamo generates SymbolicContext before calling from_tensor + original_wrap_to_fake_tensor_and_record = wrap_to_fake_tensor_and_record + + def wrap_to_fake_tensor_and_record_wrapper(t, *args, **kwargs): + dummy_tensor = wrap_if_ds_param(t) + ret = original_wrap_to_fake_tensor_and_record(dummy_tensor, *args, **kwargs) + if tracing_context := torch._guards.TracingContext.try_get(): + tracing_context.tensor_to_context[t] = tracing_context.tensor_to_context.pop(dummy_tensor) + return ret + + torch._dynamo.variables.builder.wrap_to_fake_tensor_and_record = wrap_to_fake_tensor_and_record_wrapper + + # aot_module_simplified uses fake_mode.from_tensor to process inputs + original_from_tensor = FakeTensorMode.from_tensor + + def from_tensor_wrapper(self, t, *args, **kwargs): + with unset_fake_temporarily(): + return original_from_tensor(self, wrap_if_ds_param(t), *args, **kwargs) + + FakeTensorMode.from_tensor = from_tensor_wrapper diff --git a/lib/python3.12/site-packages/deepspeed/compile/profilers/__init__.py b/lib/python3.12/site-packages/deepspeed/compile/profilers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7adb54f11872d8e7cd7df6d72b336fb863d342d5 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/profilers/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import List, Tuple +from dataclasses import dataclass, field + +from torch.fx import Graph + + +@dataclass +class ProfilingResult: + fwd_graph: Graph = None + bwd_graph: Graph = None + needs_backward: bool = False + fwd_mem: List[Tuple[str, int, int, int]] = field(default_factory=list) # name, current_alloc, delta, peak + bwd_mem: List[Tuple[str, int, int, int]] = field(default_factory=list) + fwd_time: List[Tuple[str, int, int]] = field(default_factory=list) # name, device_time, wall_time + bwd_time: List[Tuple[str, int, int]] = field(default_factory=list) + fwd_tensor_sizes: List[Tuple[str, int]] = field(default_factory=list) # name, size + bwd_tensor_sizes: List[Tuple[str, int]] = field(default_factory=list) + param_indices: List[Tuple[int, int, Tuple[int, ...]]] = field(default_factory=list) # index, ds_id, ds_shape diff --git a/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c77c64a73dd6d43288fc39ebb8e8df5f030319f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/comm_profile.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/comm_profile.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de094c35f5e3100b0436503bc027c09005dc7f43 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/comm_profile.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/graph_profile.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/graph_profile.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cad03998a1417d85136f2388635fc6c92792053b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compile/profilers/__pycache__/graph_profile.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compile/profilers/comm_profile.py b/lib/python3.12/site-packages/deepspeed/compile/profilers/comm_profile.py new file mode 100644 index 0000000000000000000000000000000000000000..18bd517c1e8f8c17fc469c4ac9f3b54db65d62d1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/profilers/comm_profile.py @@ -0,0 +1,171 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch + +try: + from torch._subclasses.fake_tensor import unset_fake_temporarily +except ImportError: + # Unsupported torch version + pass + +import deepspeed +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator + + +def sync_all(): + get_accelerator().synchronize() + dist.barrier() + + +def get_bw(comm_op, size, duration): + n = dist.get_world_size() + tput = 0 + busbw = 0 + + if duration == 0: + raise ValueError("Error. Duration is 0.") + + if comm_op == "all_to_all": + tput = (size / duration) + busbw = (size / duration) * ((n - 1) / n) + elif comm_op == "all_gather": + size *= n + tput = (size / duration) + busbw = (size / duration) * ((n - 1) / n) + elif comm_op == "all_reduce": + tput = (size * 2 / duration) + busbw = (size / duration) * (2 * (n - 1) / n) + elif comm_op == "pt2pt" or comm_op == "broadcast": + tput = (size / duration) + busbw = tput + else: + raise ValueError("wrong comm_op specified") + + return tput, busbw + + +# Run all_gather and print metrics +def timed_all_gather(device, input, output, start_event, end_event, warmup, trials, async_op): + sync_all() + # Warmups, establish connections, etc. + for i in range(warmup): + dist.all_gather_into_tensor(output, input, async_op=async_op) + sync_all() + + # time the actual comm op trials times and average it + start_event.record() + for i in range(trials): + dist.all_gather_into_tensor(output, input, async_op=async_op) + end_event.record() + sync_all() + duration = start_event.elapsed_time(end_event) / 1000 + + # maintain and clean performance data + avg_duration = duration / trials + size = input.element_size() * input.nelement() * dist.get_world_size() + # tput, busbw = get_bw('all_gather', size, avg_duration) + + avg_duration_ten = torch.tensor([avg_duration], device=device) + if dist.get_world_size() > 1: + dist.all_reduce(avg_duration_ten, dist.ReduceOp.AVG) + + return size, avg_duration_ten.item() + + +def run_all_gather(device, dtype, maxsize, warmup=5, trials=10, async_op=False): + + # Prepare benchmark header + global_rank = dist.get_rank() + world_size = dist.get_world_size() + + start_event = get_accelerator().Event(enable_timing=True) + end_event = get_accelerator().Event(enable_timing=True) + + # Create list of message sizes + M_LIST = [] + for x in (2**p for p in range(1, maxsize)): + m = x // world_size + if m > 0: + M_LIST.append(m) + + results = [(0, 0)] + sync_all() + # loop over various tensor sizes + for M in M_LIST: + global_rank = dist.get_rank() + try: + mat = torch.ones(M, dtype=dtype, device=device) + sync_all() + input = ((mat.mul_(float(global_rank))).view(-1)) + # Delete original mat to avoid OOM + del mat + get_accelerator().empty_cache() + output = torch.zeros(input.nelement() * world_size, dtype=dtype, device=device) + except RuntimeError as e: + if 'out of memory' in str(e): + if dist.get_rank() == 0: + print('WARNING: Ran out of GPU memory. Exiting comm op.') + sync_all() + break + else: + raise e + sync_all() + results.append(timed_all_gather(device, input, output, start_event, end_event, warmup, trials, async_op)) + + return results + + +profile_results = None + + +def create_predictor(): + global profile_results + if profile_results is None: + with unset_fake_temporarily(): + device = get_accelerator().current_device() + profile_results = run_all_gather(device, torch.bfloat16, 31) + if dist.get_rank() == 0: + for size, avg_duration in profile_results: + print(f"size: {size}, avg_duration: {avg_duration}") + + # Extract size and avg_duration from results + sizes = [result[0] for result in profile_results] + durations = [result[1] for result in profile_results] + + try: + from scipy.interpolate import interp1d + except ImportError: + raise RuntimeError("Please install scipy to use communication profiler in DeepCompile") + + predictor = interp1d(sizes, durations, kind='linear', fill_value="extrapolate") + + def f(size): + if size == 0: + return 0 + return predictor(size) + + # Create an interpolation function + return f + + +if __name__ == "__main__": + local_rank = int(os.environ['LOCAL_RANK']) + get_accelerator().set_device(local_rank) + print(f"local_rank={local_rank}") + + deepspeed.init_distributed(dist_backend='nccl') + + # Create predictor function + predictor = create_predictor() + + # Predict time for a specific data size + example_size = 1e9 + predicted_time = predictor(example_size) + print(f"Predicted time for size {example_size}: {predicted_time:.6f} seconds") + + dist.destroy_process_group() diff --git a/lib/python3.12/site-packages/deepspeed/compile/profilers/graph_profile.py b/lib/python3.12/site-packages/deepspeed/compile/profilers/graph_profile.py new file mode 100644 index 0000000000000000000000000000000000000000..6cb3d83e485a0ba570f78312fba1f8e776b54215 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/profilers/graph_profile.py @@ -0,0 +1,295 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import time +from typing import Any, Tuple, Dict +import statistics + +import torch +from torch.fx import GraphModule, Interpreter +from torch.fx.node import map_aggregate + +try: + from torch.utils._pytree import tree_all, tree_leaves + from torch._subclasses.fake_tensor import unset_fake_temporarily, is_fake +except ImportError: + # Unsupported torch version + pass + +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from ..util import is_comm_op, is_release_node, get_deepcompile_handle + + +def _all_real_if_tensor(args): + return tree_all(lambda x: not torch.is_tensor(x) or not is_fake(x), args) + + +def _to(v, device): + if torch.is_tensor(v): + with unset_fake_temporarily(): + return v.to(device) + return v + + +def _args_to_key(v): + + def _tensor_to_key(v) -> str: + if torch.is_tensor(v): + if v.numel() == 1: + return f"{v.dtype}{v.device}{v.item()}" + else: + return f"{v.dtype}{v.device}{v.shape}" + return str(v) + + return map_aggregate(v, _tensor_to_key) + + +def _node_size(out): + return sum([v.element_size() * v.numel() for v in tree_leaves(out) if torch.is_tensor(v)]) + + +def _get_mem_usage_out_of_torch(): + + adjust = 0 + try: + import pynvml + pynvml.nvmlInit() + + current_dev_id = get_accelerator().current_device() + handle = pynvml.nvmlDeviceGetHandleByIndex(current_dev_id) + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + + torch_alloc = get_accelerator().memory_allocated() + adjust = info.used - torch_alloc + except: + # pynvml not available + pass + + return adjust + + +# https://pytorch.org/tutorials/intermediate/fx_profiling_tutorial.html +class ProfilingInterpreter(Interpreter): + + def __init__(self, gm: GraphModule, iteration: int = 10, warmup: int = 5, debug_log=False): + super().__init__(gm) + + self.nz3 = get_deepcompile_handle() + + assert iteration > 0 + assert warmup >= 0 + self.iteration = iteration + self.warmup = warmup + self.device = torch.device(get_accelerator().current_device()) + self.cache: Dict[Tuple, Any] = {} + self.distributed = dist.is_initialized() + self.allgather_mem: Dict[int, int] = {} + self.debug_log = debug_log + self.mem_usage_out_of_torch = 0 + + def run(self, *args) -> Any: + """Run the graph with profiling enabled. + + args: inputs to the graph. Tensors in the inpusts must be real tensors, not fake tensors. args can contain ds parameters. + returns: The output of the graph. Tensor in the output is real tensors. + """ + try: + assert _all_real_if_tensor(args), "Inputs must be real tensors" + self.nz3.enable_profiling(True) + + with unset_fake_temporarily(): + with get_accelerator().random().fork_rng(devices=[self.device]): + self.mem_usage_out_of_torch = _get_mem_usage_out_of_torch() + return_val = super().run(*args) + except Exception as e: + msg = e.msg if "msg" in dir(e) else str(e) + print(f"Profiling error {msg}") + finally: + self.nz3.clear_all_gathered_params() + self.nz3.enable_profiling(False) + return return_val + + def run_node(self, n: torch.fx.Node) -> Any: + + if n.op in {"placeholder", "output"}: + n.meta["device_time"] = 0.0 + n.meta["wall_time"] = 0.0 + n.meta["alloc_mem"] = 0 + n.meta["max_memory"] = 0 + n.meta["tensor_size"] = _node_size(n) + return super().run_node(n) + + args, kwargs = self.fetch_args_kwargs_from_env(n) + assert isinstance(args, tuple) + assert isinstance(kwargs, dict) + + def rebuild_param_if_necessary(v): + if hasattr(v, "ds_id"): + v.all_gather(param_list=[v]) + return v + + args = map_aggregate(args, lambda x: rebuild_param_if_necessary(x)) + + args = map_aggregate(args, lambda x: _to(x, self.device)) + kwargs = map_aggregate(kwargs, lambda x: _to(x, self.device)) + + cache_key = (n.target, _args_to_key(args), _args_to_key(kwargs)) + cache_hit = cache_key in self.cache + + cache_hit_flag = torch.tensor([0 if cache_hit else 1], device=self.device, dtype=torch.int) + if self.distributed: + dist.all_reduce(cache_hit_flag, dist.ReduceOp.SUM) + cache_hit = cache_hit_flag.item() == 0 + + if cache_hit: + device_time, wall_time, alloc_mem, max_mem, tensor_size = self.cache[cache_key] + n.meta["device_time"] = device_time + n.meta["wall_time"] = wall_time + n.meta["alloc_mem"] = alloc_mem + n.meta["max_mem"] = max_mem + n.meta["tensor_size"] = tensor_size + + is_release_op = is_release_node(n) + run_only_once = cache_hit or is_release_op + iteration = 1 if run_only_once else self.iteration + accelerator = get_accelerator() + start_events = [accelerator.Event(enable_timing=True) for _ in range(iteration)] + end_events = [accelerator.Event(enable_timing=True) for _ in range(iteration)] + + get_accelerator().reset_peak_memory_stats() + alloc_mem_start = get_accelerator().memory_allocated() + max_mem_start = get_accelerator().max_memory_allocated() + + if not run_only_once: + for i in range(self.warmup): + out = getattr(self, n.op)(n.target, args, kwargs) + + if is_comm_op(n): + assert self.distributed, f"Distributed environment is not initialized but comm operator {n.name} {n.target} is used." + dist.barrier() + + start = time.time() + for i in range(iteration): + start_events[i].record() + out = getattr(self, n.op)(n.target, args, kwargs) + end_events[i].record() + accelerator.synchronize() + walltime_sum = time.time() - start + + if is_comm_op(n): + dist.barrier() + + alloc_mem = get_accelerator().memory_allocated() - alloc_mem_start + self.mem_usage_out_of_torch + max_memory = get_accelerator().max_memory_allocated() - max_mem_start + self.mem_usage_out_of_torch + tensor_size = _node_size(out) + + def partition_param_if_necessary(v): + if hasattr(v, "ds_id") and not v.ds_persist: + v.partition(param_list=[v], has_been_updated=False) + return v + + args = map_aggregate(args, lambda x: partition_param_if_necessary(x)) + + if not cache_hit: + device_time = statistics.mean([s.elapsed_time(e) for s, e in zip(start_events, end_events)]) + wall_time = walltime_sum / iteration * 1000 + + with unset_fake_temporarily(): + vals_to_bcast = torch.tensor([device_time, wall_time, alloc_mem, max_memory, tensor_size], + device=self.device) + if self.distributed: + dist.all_reduce(vals_to_bcast, dist.ReduceOp.AVG) + n.meta["device_time"] = vals_to_bcast[0].item() + n.meta["wall_time"] = vals_to_bcast[1].item() + n.meta["alloc_mem"] = int(vals_to_bcast[2].item()) + n.meta["max_mem"] = int(vals_to_bcast[3].item()) + n.meta["tensor_size"] = int(vals_to_bcast[4].item()) + self.cache[cache_key] = (n.meta["device_time"], n.meta["wall_time"], n.meta["alloc_mem"], + n.meta["max_mem"], n.meta["tensor_size"]) + + if is_release_op: + n.meta["alloc_mem"] = -self.allgather_mem.get(args[2], 0) + + if dist.get_rank() == 0 and self.debug_log: + print( + f"{n.target} {n.meta['device_time']:.2f}ms {n.meta['wall_time']:.2f}ms alloc_mem={n.meta['alloc_mem'] / 1024 / 1024:.2f}MB max_mem={n.meta['max_mem'] / 1024 / 1024:.2f}MB tensor_size={n.meta['tensor_size']}" + ) + + if n.target == torch.ops.dc.allgather_param.default: + out = args[0] + assert hasattr(out, "ds_id") + if not out.ds_persist: + self.nz3.invalidate_gathered_param(args[2]) + self.allgather_mem[out.ds_id] = n.meta["alloc_mem"] + + return out + + +class MemoryProfilingInterpreter(Interpreter): + + def __init__(self, gm: GraphModule, debug_log=False): + super().__init__(gm) + self.nz3 = get_deepcompile_handle() + self.device = torch.device(get_accelerator().current_device()) + self.mem_record = [] + self.last_alloc = get_accelerator().memory_allocated() + + self.node_counter = 0 + self.node_num = len(gm.graph.nodes) + self.debug_log = debug_log + + def run(self, *args) -> Any: + try: + assert _all_real_if_tensor(args), "Inputs must be real tensors" + self.nz3.enable_profiling(True) + self.mem_usage_out_of_torch = _get_mem_usage_out_of_torch() + + with unset_fake_temporarily(): + with get_accelerator().random().fork_rng(devices=[self.device]): + return_val = super().run(*args) + except Exception as e: + print(f"MemoryProfiling error {e}") + finally: + self.nz3.enable_profiling(False) + + return return_val + + def run_node(self, n: torch.fx.Node) -> Any: + get_accelerator().reset_peak_memory_stats() + + if n.op in {"placeholder", "output"}: + ret = super().run_node(n) + else: + args, kwargs = self.fetch_args_kwargs_from_env(n) + args = map_aggregate(args, lambda x: _to(x, self.device)) + kwargs = map_aggregate(kwargs, lambda x: _to(x, self.device)) + ret = getattr(self, n.op)(n.target, args, kwargs) + + del args, kwargs + + current_alloc = get_accelerator().memory_allocated() + self.mem_usage_out_of_torch + max_alloc = get_accelerator().max_memory_allocated() + self.mem_usage_out_of_torch + vals_to_bcast = torch.tensor([current_alloc, max_alloc], device=self.device) + dist.all_reduce(vals_to_bcast, dist.ReduceOp.MAX) + current_alloc = vals_to_bcast[0].item() + max_alloc = vals_to_bcast[1].item() + + self.mem_record.append((n.name, current_alloc, current_alloc - self.last_alloc, max_alloc)) + + self.node_counter += 1 + if self.debug_log and dist.get_rank() == 0: + print( + f"Mem prof Node {self.node_counter}/{self.node_num} {n.name} memory {current_alloc / 1024 / 1024:.2f}MB delta {(current_alloc - self.last_alloc) / 1024 / 1024:.2f}MB" + ) + + self.last_alloc = current_alloc + + return ret + + def dump(self, path): + import pandas as pd + df = pd.DataFrame(self.mem_record, columns=["node", "memory", "delta", "max_mem"]) + df.to_csv(path, index=False) diff --git a/lib/python3.12/site-packages/deepspeed/compile/util.py b/lib/python3.12/site-packages/deepspeed/compile/util.py new file mode 100644 index 0000000000000000000000000000000000000000..fdeb5f4347a9d1ccba2bfc1643c5d68029825b75 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/compile/util.py @@ -0,0 +1,429 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import functools +import operator +from typing import List, Tuple, Dict +from collections import defaultdict + +import torch +from torch.fx import Node, Graph +from torch.fx.node import map_aggregate, Argument, map_arg + +try: + from torch._subclasses.fake_tensor import unset_fake_temporarily +except ImportError: + # Unsupported torch version + pass + +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.utils.torch import required_torch_version +from deepspeed.ops.op_builder.dc import DeepCompileBuilder + + +def is_deepcompile_supported() -> bool: + return required_torch_version(min_version=2.6, max_version=2.7) and get_accelerator().device_name() == "cuda" + + +dc_handle = None + +if is_deepcompile_supported(): + sym_size_ops = { + operator.ge, + operator.le, + operator.eq, + operator.ne, + operator.gt, + operator.lt, + torch.ops.aten.sym_size.int, + operator.getitem, + } + + +def get_deepcompile_handle(): + global dc_handle + if dc_handle is None: + dc_handle = DeepCompileBuilder().load() + return dc_handle + + +def is_backend_inductor(backend): + return backend == "inductor" + + +backward_started = False +pre_backward_hooks = [] + + +def add_pre_backward_hook(hook): + pre_backward_hooks.append(hook) + + +def deepcompile_backward_prologue(is_gradient_accumulation_boundary): + + for hook in pre_backward_hooks: + hook() + + dc = get_deepcompile_handle() + dc.start_backward(is_gradient_accumulation_boundary) + + +def log_rank0(msg: str, enable: bool = False): + if dist.get_rank() == 0 and enable: + print(msg) + + +def get_no_copy_ops(): + # Need to compile custom ops + get_deepcompile_handle() + return { + torch.ops.aten.t.default, torch.ops.aten.view.default, torch.ops.aten.detach.default, + torch.ops.aten.permute.default, torch.ops.dc.wait_allgather.default + } + + +def get_input_nodes(graph: Graph) -> List[Node]: + return [n for n in graph.nodes if n.op == "placeholder"] + + +def get_param_nodes(graph: Graph, index_to_ds_ids: List[Tuple[int, int]]) -> List[Node]: + all_input_nodes = get_input_nodes(graph) + return [all_input_nodes[i] for i, _, _ in index_to_ds_ids] + + +def is_comm_op(node: Node) -> bool: + return "comm" in node.meta and node.meta["comm"] + + +def exclude_from_act_offload(node: Node) -> bool: + return node.target in sym_size_ops + + +def dtype_to_elem_size(dtype: torch.dtype) -> int: + if dtype == torch.float32: + elem_size = 4 + elif dtype == torch.float64: + elem_size = 8 + elif dtype == torch.float16: + elem_size = 2 + else: + raise ValueError(f"Unsupported dtype: {dtype}") + return elem_size + + +def tensor_meta_size(tensor_meta) -> int: + numel = 1 if len(tensor_meta.shape) == 0 else functools.reduce(operator.mul, tensor_meta.shape) + + dtype = tensor_meta.dtype + if dtype == torch.float32: + elem_size = 4 + elif dtype == torch.float64 or dtype == torch.int64: + elem_size = 8 + elif dtype == torch.float16 or dtype == torch.bfloat16: + elem_size = 2 + elif dtype == torch.bool: + elem_size = 1 + else: + raise ValueError(f"Unsupported dtype: {dtype}") + + return numel * elem_size + + +class NodeValueOffloadHelper: + + def __init__(self, device): + self.device = device + self.env_values: Dict[str, Argument] = {} + self.original_device: Dict[torch.Tensor, torch.device] = {} + + def _to_cpu(self, v): + if torch.is_tensor(v): + with unset_fake_temporarily(): + device = v.device + offloaded = v.to('cpu').detach() + self.original_device[offloaded] = device + return offloaded + return v + + def _from_cpu(self, v): + if torch.is_tensor(v) and v in self.original_device: + return v.to(self.original_device[v]) + return v + + def save(self, name: str, v: Argument, offload) -> None: + self.env_values[name] = map_aggregate(v, lambda x: self._to_cpu(x) if offload else x) + + def load(self, name: str) -> Argument: + return map_aggregate(self.env_values[name], lambda x: self._from_cpu(x)) + + def get_offloaded_value(self, name: str) -> Argument: + return self.env_values[name] + + def has_value(self, name: str) -> bool: + return name in self.env_values + + def clear(self) -> None: + self.env_values.clear() + self.original_device.clear() + + +def materialize_fake(v, device=None): + from torch._subclasses.fake_tensor import is_fake + + def convert(t): + if is_fake(t): + with unset_fake_temporarily(): + if t.is_floating_point(): + return torch.randn(t.shape, + dtype=t.dtype, + device=t.device if device is None else device, + layout=t.layout, + requires_grad=t.requires_grad, + pin_memory=t.is_pinned()) + else: + return torch.zeros(t.shape, + dtype=t.dtype, + device=t.device if device is None else device, + requires_grad=t.requires_grad) + + return t + + return map_aggregate(v, lambda x: convert(x)) + + +def get_last_uses(graph: Graph): + position = {node: i for i, node in enumerate(graph.nodes)} + + node_to_last_use: Dict[Node, Node] = {} + user_to_last_uses: Dict[Node, List[Node]] = {} + no_copy_ops = get_no_copy_ops() + + def register_last_uses(n: Node, user: Node): + update = False + known_last_use = None + + if user.target in no_copy_ops and n in node_to_last_use: + last_user = node_to_last_use[user] + last_use_position = position[last_user] + + known_last_use = node_to_last_use[n] + known_last_use_position = position[known_last_use] + update = last_use_position > known_last_use_position + + if n not in node_to_last_use or update: + if user.target in no_copy_ops: + user = node_to_last_use[user] + + node_to_last_use[n] = user + user_to_last_uses.setdefault(user, []).append(n) + + if known_last_use: + user_to_last_uses[known_last_use].remove(n) + + for node in reversed(graph.nodes): + map_arg(node.args, lambda n: register_last_uses(n, node)) + map_arg(node.kwargs, lambda n: register_last_uses(n, node)) + + return node_to_last_use, user_to_last_uses + + +def get_real_uses(graph: Graph): + node_to_uses: Dict[Node, List[Node]] = defaultdict(list) + no_copy_ops = get_no_copy_ops() + + def register_last_uses(n: Node, user: Node): + if user.target == "output": + return + + if user.target in no_copy_ops: + users = node_to_uses[user] + node_to_uses[n].extend(users) + else: + node_to_uses[n].append(user) + + for node in reversed(graph.nodes): + map_arg(node.args, lambda n: register_last_uses(n, node)) + map_arg(node.kwargs, lambda n: register_last_uses(n, node)) + + return node_to_uses + + +def count_inflight_values(graph: Graph, file_path: str): + position = {node: i for i, node in enumerate(graph.nodes)} + + node_to_last_use, user_to_last_uses = get_last_uses(graph) + + max_inflight_size = 0 + inflight_values = set() + + # Output csv. + csv_filename = file_path + csv_data = [] + header = [ + 'Node', 'tensor_size', 'inflight_size', 'inflight_size_in_output', 'args', 'users', 'node_to_last_use', + 'lifetime', 'user_to_last_uses', 'inflight_values' + ] + csv_data.append(header) + + from .fx import get_output_node + output_node = get_output_node(graph) + values_in_output = set([n for n in output_node.args[0] if isinstance(n, Node)]) + + for node in graph.nodes: + inflight_values.add(node) + if node in user_to_last_uses: + for to_delete in user_to_last_uses[node]: + inflight_values.remove(to_delete) + + assert "tensor_size" in node.meta, f"Node {node} does not have tensor_size" + inflight_size = sum(n.meta["tensor_size"] for n in inflight_values) + inflight_size_in_output = sum(n.meta["tensor_size"] for n in inflight_values if n in values_in_output) + + lifetime = position[node_to_last_use[node]] - position[node] if node in node_to_last_use else 0 + + row = [ + node.name, node.meta["tensor_size"], inflight_size, inflight_size_in_output, + [a.name for a in node.args if isinstance(a, Node)], + list(node.users.keys()), node_to_last_use[node] if node in node_to_last_use else 'NA', lifetime, + user_to_last_uses[node] if node in user_to_last_uses else 'NA', + list(inflight_values) + ] + csv_data.append(row) + + # print( + # f"Node: {node.name} users: {list(node.users.keys())} node_to_last_use: {node_to_last_use[node] if node in node_to_last_use else 'NA'} user_to_last_uses: {user_to_last_uses[node] if node in user_to_last_uses else 'NA'} inflight_values: {inflight_values} inflight_size: {inflight_size}" + # ) + max_inflight_size = max(max_inflight_size, inflight_size) + + import csv + with open(csv_filename, mode='w', newline='') as file: + writer = csv.writer(file) + writer.writerows(csv_data) + + print(f"Max inflight size: {max_inflight_size}") + print(f"Data successfully written to {csv_filename}") + + +def get_activation_node_names(graph: Graph, param_nodes_bw: List[Node], fwd_output_names: List[str]): + + input_nodes = get_input_nodes(graph) + param_node_names = set([n.name for n in param_nodes_bw]) + + activation_node_names = [] + for in_node in input_nodes: + if in_node.name in fwd_output_names: + if in_node.name not in param_node_names: + activation_node_names.append(in_node.name) + + return activation_node_names + + +class TensorOffloadHelper(): + + def __init__(self): + self.devices = {} + self.base_tensors = {} + self.views = {} + self.arg_list = [] + self.offloaded = {} + self.non_tensor = {} + + def offload(self, argument): + + def is_base_tensor(tensor): + return torch.is_tensor(a) and not a._is_view() and not hasattr(tensor, "ds_id") + + base_tensor_ids = set() + for a in argument: + if is_base_tensor(a): + base_tensor_ids.add(id(a)) + + for a in argument: + a_id = id(a) + + if is_base_tensor(a): + # Base tensor + self.devices[a_id] = a.device + self.base_tensors[a_id] = a + # elif torch.is_tensor(a) and not hasattr(a, "ds_id") and id(a._base) in base_tensor_ids: + # # View + # self.views[a_id] = { + # "base_id": id(a._base), + # "size": a.size(), + # "stride": a.stride(), + # "offset": a.storage_offset(), + # } + else: + # other types or ds tensor + self.non_tensor[a_id] = a + + self.arg_list.append(a_id) + + for a in argument: + if is_base_tensor(a): + a.data = a.data.to("cpu") + + def reload(self, in_place): + + loaded_base_tensors = {} + for a_id in self.arg_list: + if a_id in self.base_tensors: + device = self.devices[a_id] + + if in_place: + self.base_tensors[a_id].data = self.base_tensors[a_id].to(device) + loaded_base_tensors[a_id] = self.base_tensors[a_id] + else: + loaded_base_tensors[a_id] = self.base_tensors[a_id].to(device) + + results = [] + for a_id in self.arg_list: + if a_id in self.base_tensors: + results.append(loaded_base_tensors[a_id]) + + # elif a_id in self.views: + # view_info = self.views[a_id] + # # print(f"load_args loading view {a_id} base_id={view_info['base_id']} size={view_info['size']} stride={view_info['stride']} offset={view_info['offset']}") + # base_tensor = loaded_base_tensors[view_info["base_id"]] + # view_tensor = base_tensor.as_strided( + # view_info["size"], view_info["stride"], view_info["offset"] + # ) + # results.append(view_tensor) + + elif a_id in self.non_tensor: + results.append(self.non_tensor[a_id]) + + return results + + +def add_mem_profile_nodes(graph: Graph, prefix: str): + + def show_memory(label: str): + if dist.get_rank() == 0: + print( + f"{prefix} {label} alloc_mem={get_accelerator().memory_allocated()} max_mem={get_accelerator().max_memory_allocated()}" + ) + + nodes = list(graph.nodes) + for node in nodes: + if node.op == "output": + continue + + with graph.inserting_after(node): + msg = f"Mem {node.name}" + name = f"show_memory_{node.name}" + graph.create_node('call_function', show_memory, (msg, ), {}, name=name) + + +def is_release_node(n: Node) -> bool: + return n.target == torch.ops.dc.release_param.default + + +def get_index_by_graph_id(graph_order, target_graph_id): + for index, (graph_id, _) in enumerate(graph_order): + if graph_id == target_graph_id: + return index + return -1 diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/__init__.py b/lib/python3.12/site-packages/deepspeed/elasticity/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..66bda96fa6bafa90c8a68f4033b665220880bcf6 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/elasticity/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .elasticity import compute_elastic_config, elasticity_enabled, ensure_immutable_elastic_config +from .utils import is_torch_elastic_compatible +from .constants import ENABLED, ENABLED_DEFAULT, ELASTICITY +if is_torch_elastic_compatible(): + from .elastic_agent import DSElasticAgent diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86e8b5f2e9b5bbea66f9a8f0d20d2c6b8139f8d1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..056561f6c8ae09bb15d9267bb46e016ea186354c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8727add057bb664e42f6f24e2b454b2c0503a21 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/elastic_agent.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/elastic_agent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eda9c8d8a37271e4ee5f163405ddac1976d1c055 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/elastic_agent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/elasticity.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/elasticity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70be7b75a5287cff3f5e99af83ef759262f3d6f7 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/elasticity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0568b05323940071a96132d5ba3260202005531d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/elasticity/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/config.py b/lib/python3.12/site-packages/deepspeed/elasticity/config.py new file mode 100644 index 0000000000000000000000000000000000000000..7c6bd42cdfd93e7979e942d3b353e55744b356d4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/elasticity/config.py @@ -0,0 +1,110 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import json +from .constants import * + + +class ElasticityError(Exception): + """ + Base exception for all elasticity related errors + """ + + +class ElasticityConfigError(ElasticityError): + """ + Elasticity configuration error + """ + + +class ElasticityIncompatibleWorldSize(ElasticityError): + """ + Attempting to run a world size that is incompatible with a given elastic config + """ + + +class ElasticityConfig: + """ + Elastic config object, constructed from a param dictionary that only contains elastic + config parameters, example below: + + If elasticity is enabled, user must specify (at least) max_train_batch_size + and micro_batch_sizes. + + { + "enabled": true, + "max_train_batch_size": 2000, + "micro_batch_sizes": [2,4,6], + "min_gpus": 1, + "max_gpus" : 10000 + "min_time": 20 + "ignore_non_elastic_batch_info": false + "version": 0.1 + } + """ + + def __init__(self, param_dict): + self.enabled = param_dict.get(ENABLED, ENABLED_DEFAULT) + if self.enabled: + if MAX_ACCEPTABLE_BATCH_SIZE in param_dict: + self.max_acceptable_batch_size = param_dict[MAX_ACCEPTABLE_BATCH_SIZE] + else: + raise ElasticityConfigError(f"Elasticity config missing {MAX_ACCEPTABLE_BATCH_SIZE}") + if MICRO_BATCHES in param_dict: + self.micro_batches = param_dict[MICRO_BATCHES] + else: + raise ElasticityConfigError(f"Elasticity config missing {MICRO_BATCHES}") + else: + self.max_acceptable_batch_size = param_dict.get(MAX_ACCEPTABLE_BATCH_SIZE, + MAX_ACCEPTABLE_BATCH_SIZE_DEFAULT) + self.micro_batches = param_dict.get(MICRO_BATCHES, MICRO_BATCHES_DEFAULT) + + if not isinstance(self.micro_batches, list): + raise ElasticityConfigError( + f"Elasticity expected value of {MICRO_BATCHES} to be a " + f"list of micro batches, instead is: {type(self.micro_batches)}, containing: {self.micro_batches}") + + if not all(map(lambda m: isinstance(m, int), self.micro_batches)): + raise ElasticityConfigError(f"Elasticity expected {MICRO_BATCHES} to only contain a list of integers, " + f"instead contains: f{self.micro_batches}") + + if not all(map(lambda m: m > 0, self.micro_batches)): + raise ElasticityConfigError(f"Elasticity expected {MICRO_BATCHES} to only contain positive integers, " + f"instead contains: f{self.micro_batches}") + + self.min_gpus = param_dict.get(MIN_GPUS, MIN_GPUS_DEFAULT) + self.max_gpus = param_dict.get(MAX_GPUS, MAX_GPUS_DEFAULT) + + if self.min_gpus < 1 or self.max_gpus < 1: + raise ElasticityConfigError("Elasticity min/max gpus must be > 0, " + f"given min_gpus: {self.min_gpus}, max_gpus: {self.max_gpus}") + if self.max_gpus < self.min_gpus: + raise ElasticityConfigError("Elasticity min_gpus cannot be greater than max_gpus, " + f"given min_gpus: {self.min_gpus}, max_gpus: {self.max_gpus}") + + self.model_parallel_size = param_dict.get(MODEL_PARALLEL_SIZE, MODEL_PARALLEL_SIZE_DEFAULT) + if self.model_parallel_size < 1: + raise ElasticityConfigError("Model-Parallel size cannot be less than 1, " + f"given model-parallel size: {self.model_parallel_size}") + + self.num_gpus_per_node = param_dict.get(NUM_GPUS_PER_NODE, NUM_GPUS_PER_NODE_DEFAULT) + if self.num_gpus_per_node < 1: + raise ElasticityConfigError("Number of GPUs per node cannot be less than 1, " + f"given number of GPUs per node: {self.num_gpus_per_node}") + + self.min_time = param_dict.get(MIN_TIME, MIN_TIME_DEFAULT) + if self.min_time < 0: + raise ElasticityConfigError(f"Elasticity min time needs to be >= 0: given {self.min_time}") + + self.version = param_dict.get(VERSION, VERSION_DEFAULT) + self.prefer_larger_batch_size = param_dict.get(PREFER_LARGER_BATCH, PREFER_LARGER_BATCH_DEFAULT) + self.ignore_non_elastic_batch_info = param_dict.get(IGNORE_NON_ELASTIC_BATCH_INFO, + IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT) + + def repr(self): + return self.__dict__ + + def __repr__(self): + return json.dumps(self.__dict__, sort_keys=True, indent=4) diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/constants.py b/lib/python3.12/site-packages/deepspeed/elasticity/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..b3134c54b4d6a127a04b9f633a51226b6713be20 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/elasticity/constants.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######################################### +# Elasticity +######################################### +''' Elasticity Utility in DeepSpeed can be used to create highly elastic jobs compatible +with a large number of GPUs. For elastic jobs, DeepSpeed will provide a batch size that +can support a large number of GPUs based on the user specified parameters +''' +FORMAT = ''' +Elasticity should be enabled as: +"elasticity": { + "enabled": true, + "max_train_batch_size": 2000, + "micro_batch_sizes": [2,4,6], + "min_gpus": 1, + "max_gpus" : 10000, + "min_time": 20, + "prefer_larger_batch": true, + "ignore_non_elastic_batch_info": false, + "version": 0.1 +} +''' + +ELASTICITY = 'elasticity' + +# Current elasticity version +LATEST_ELASTICITY_VERSION = 0.2 + +ENABLED = 'enabled' +ENABLED_DEFAULT = False + +# Max acceptable train_batch_size +MAX_ACCEPTABLE_BATCH_SIZE = 'max_train_batch_size' +MAX_ACCEPTABLE_BATCH_SIZE_DEFAULT = 2000 + +# Acceptable micro batch sizes, same as train_micro_batch_size_per_gpu +MICRO_BATCHES = 'micro_batch_sizes' +MICRO_BATCHES_DEFAULT = [2, 4, 6] + +# Min/max of GPUs to search over +MIN_GPUS = 'min_gpus' +MIN_GPUS_DEFAULT = 1 +MAX_GPUS = 'max_gpus' +MAX_GPUS_DEFAULT = 10000 + +NUM_GPUS_PER_NODE = 'num_gpus_per_node' +NUM_GPUS_PER_NODE_DEFAULT = 1 + +MODEL_PARALLEL_SIZE = "model_parallel_size" +MODEL_PARALLEL_SIZE_DEFAULT = 1 + +# Minimum running time (minutes) before the scheduler will scale us, 0 implies it's unknown +MIN_TIME = "min_time" +MIN_TIME_DEFAULT = 0 + +# When finding a suitable batch size, attempt to find one that is closest +# to the max train batch size given. +PREFER_LARGER_BATCH = 'prefer_larger_batch' +PREFER_LARGER_BATCH_DEFAULT = True + +# In order to reduce confusion, if elastic mode is enabled we +# require (via assert) that no batch info is set outside of the +# elastic config. You can turn off this assert via this config +# but keep in mind that all batch info defined outside the +# elastic mode *will be ignored*. +IGNORE_NON_ELASTIC_BATCH_INFO = 'ignore_non_elastic_batch_info' +IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT = False + +# Version of elastic logic to use +VERSION = "version" +VERSION_DEFAULT = LATEST_ELASTICITY_VERSION + +# Minimum deepspeed version to use elasticity +MINIMUM_DEEPSPEED_VERSION = "0.3.8" + +# Environment variable storing elastic config from resource scheduler +DEEPSPEED_ELASTICITY_CONFIG = "DEEPSPEED_ELASTICITY_CONFIG" diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/elastic_agent.py b/lib/python3.12/site-packages/deepspeed/elasticity/elastic_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..8fd4293d312c08e2d24a56e6ea0a678488febdc4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/elasticity/elastic_agent.py @@ -0,0 +1,190 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent +from typing import Any, Dict, Optional, Tuple +from datetime import datetime +from torch.distributed.elastic.utils.distributed import get_free_port +from torch.distributed.elastic.metrics import put_metric +from torch.distributed.elastic.agent.server.api import ( + RunResult, + WorkerGroup, + WorkerSpec, + WorkerState, +) +from torch.distributed import Store +import time +import os +from torch.distributed.elastic.multiprocessing import start_processes +from torch.distributed.elastic.utils import macros +import shutil +import copy +from contextlib import closing +import subprocess + +from torch.distributed.elastic.utils.logging import get_logger + +log = get_logger(__name__) + + +class DSElasticAgent(LocalElasticAgent): + + def __init__( + self, + spec: WorkerSpec, + env: Dict, + start_method="spawn", + exit_barrier_timeout: float = 300, + log_dir: Optional[str] = None, + ): + super().__init__(spec, start_method, exit_barrier_timeout, log_dir) + self.ds_env = env + + @staticmethod + def _set_master_addr_port(store: Store, + master_addr: Optional[str], + master_port: Optional[int], + local_addr: Optional[str] = None): + if master_port is None: + sock = get_free_port() + with closing(sock): + master_port = sock.getsockname()[1] + + if master_addr is None: + # master_addr = _get_fq_hostname() + import shlex + safe_cmd = shlex.split("hostname -I") + result = subprocess.check_output(safe_cmd) + master_addr = result.decode('utf-8').split()[0] + + store.set("MASTER_ADDR", master_addr.encode(encoding="UTF-8")) + store.set("MASTER_PORT", str(master_port).encode(encoding="UTF-8")) + + def _start_workers(self, worker_group: WorkerGroup) -> Dict[int, Any]: + spec = worker_group.spec + store = worker_group.store + assert store is not None + master_addr, master_port = super()._get_master_addr_port(store) + restart_count = spec.max_restarts - self._remaining_restarts + + use_agent_store = spec.rdzv_handler.get_backend() == "static" + + args: Dict[int, Tuple] = {} + envs: Dict[int, Dict[str, str]] = {} + for worker in worker_group.workers: + local_rank = worker.local_rank + + worker_env_ds = copy.deepcopy(self.ds_env) + worker_env_elastic = { + "LOCAL_RANK": str(local_rank), + "RANK": str(worker.global_rank), + "GROUP_RANK": str(worker_group.group_rank), + "ROLE_RANK": str(worker.role_rank), + "ROLE_NAME": spec.role, + "LOCAL_WORLD_SIZE": str(spec.local_world_size), + "WORLD_SIZE": str(worker.world_size), + "GROUP_WORLD_SIZE": str(worker_group.group_world_size), + "ROLE_WORLD_SIZE": str(worker.role_world_size), + "MASTER_ADDR": master_addr, + "MASTER_PORT": str(master_port), + "TORCHELASTIC_RESTART_COUNT": str(restart_count), + "TORCHELASTIC_MAX_RESTARTS": str(spec.max_restarts), + "TORCHELASTIC_RUN_ID": spec.rdzv_handler.get_run_id(), + "TORCHELASTIC_USE_AGENT_STORE": str(use_agent_store), + "NCCL_ASYNC_ERROR_HANDLING": os.getenv("NCCL_ASYNC_ERROR_HANDLING", str(1)), + } + worker_env_ds.update(worker_env_elastic) + if "OMP_NUM_THREADS" in os.environ: + worker_env_ds["OMP_NUM_THREADS"] = os.environ["OMP_NUM_THREADS"] + + envs[local_rank] = worker_env_ds + worker_args = list(spec.args) + worker_args = macros.substitute(worker_args, str(local_rank)) + args[local_rank] = tuple(worker_args) + + # scaling events do not count towards restarts (gets same attempt #) + # remove existing log dir if this restart is due to a scaling event + attempt_log_dir = os.path.join(self._log_dir, f"attempt_{restart_count}") + shutil.rmtree(attempt_log_dir, ignore_errors=True) + os.makedirs(attempt_log_dir) + + assert spec.entrypoint is not None + self._pcontext = start_processes( + name=spec.role, + entrypoint=spec.entrypoint, + args=args, + envs=envs, + log_dir=attempt_log_dir, + start_method=self._start_method, + redirects=spec.redirects, + tee=spec.tee, + ) + + return self._pcontext.pids() + + def _invoke_run(self, role: str = "default") -> RunResult: + # NOTE: currently only works for a single role + + spec = self._worker_group.spec + role = spec.role + + log.info(f"[{role}] starting workers for entrypoint: {spec.get_entrypoint_name()}") + + self._initialize_workers(self._worker_group) + monitor_interval = spec.monitor_interval + rdzv_handler = spec.rdzv_handler + + participants = rdzv_handler._state_holder.state.participants + + while True: + assert self._worker_group.state != WorkerState.INIT + time.sleep(monitor_interval) + run_result = self._monitor_workers(self._worker_group) + state = run_result.state + self._worker_group.state = state + + expire_time = datetime.utcnow() - (rdzv_handler._settings.keep_alive_interval * + rdzv_handler._settings.keep_alive_max_attempt) + _dead_nodes = [ + node for node, last_heartbeat in rdzv_handler._state_holder.state.last_heartbeats.items() + if last_heartbeat < expire_time + ] + + put_metric(f"workers.{role}.remaining_restarts", self._remaining_restarts) + put_metric(f"workers.{role}.{state.name.lower()}", 1) + + if state == WorkerState.SUCCEEDED: + log.info(f"[{role}] worker group successfully finished." + f" Waiting {self._exit_barrier_timeout} seconds for other agents to finish.") + self._exit_barrier() + return run_result + elif state in {WorkerState.UNHEALTHY, WorkerState.FAILED} or len(participants) > len( + rdzv_handler._state_holder.state.participants): + if self._remaining_restarts > 0: + log.info(f"[{role}] Worker group {state.name}. " + f"{self._remaining_restarts}/{spec.max_restarts} attempts left;" + f" will restart worker group") + self._remaining_restarts -= 1 + # rdzv_handler._state_holder.state.restart = False + self._restart_workers(self._worker_group) + participants = rdzv_handler._state_holder.state.participants + + else: + self._stop_workers(self._worker_group) + self._worker_group.state = WorkerState.FAILED + self._exit_barrier() + return run_result + elif state == WorkerState.HEALTHY: + # membership changes do not count as retries + num_nodes_waiting = rdzv_handler.num_nodes_waiting() + group_rank = self._worker_group.group_rank + if num_nodes_waiting > 0: + log.info(f"[{role}] Detected {num_nodes_waiting} " + f"new nodes from group_rank={group_rank}; " + f"will restart worker group") + self._restart_workers(self._worker_group) + participants = rdzv_handler._state_holder.state.participants + else: + raise Exception(f"[{role}] Worker group in {state.name} state") diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/elasticity.py b/lib/python3.12/site-packages/deepspeed/elasticity/elasticity.py new file mode 100644 index 0000000000000000000000000000000000000000..f26f13c98dfaf965597d699418d98c80f35f67b4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/elasticity/elasticity.py @@ -0,0 +1,378 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import json +import numpy as np +import math +from packaging import version as pkg_version + +from .config import ElasticityConfig, ElasticityConfigError, ElasticityError, \ + ElasticityIncompatibleWorldSize +from .constants import ELASTICITY, ENABLED, ENABLED_DEFAULT, LATEST_ELASTICITY_VERSION, \ + MINIMUM_DEEPSPEED_VERSION, DEEPSPEED_ELASTICITY_CONFIG +from ..git_version_info import version as __version__ +from ..utils import logger + +# Thirty eight smallest highly composite numbers. The list should +# be enough to support up to 720K batch size. +HCN_LIST = [ + 1, 2, 4, 6, 12, 24, 36, 48, 60, 120, 180, 240, 360, 720, 840, 1260, 1680, 2520, 5040, 7560, 10080, 15120, 20160, + 25200, 27720, 45360, 50400, 55440, 83160, 110880, 166320, 221760, 277200, 332640, 498960, 554400, 665280, 720720 +] + + +def get_candidate_batch_sizes(base_list, max_acceptable_batch_size): + candidate_batch_size = [] + for base in base_list: + if base >= max_acceptable_batch_size: + candidate_batch_size.append(base) + else: + value = max_acceptable_batch_size // base + index = np.argmax(np.asarray(HCN_LIST) > value) + candidate_batch_size.append(HCN_LIST[index - 1] * base) + candidate_batch_size = list(set(candidate_batch_size)) + logger.info(f"Candidate batch size: {candidate_batch_size}") + return candidate_batch_size + + +def get_valid_gpus(batch_size, micro_batches, min_valid_gpus, max_valid_gpus): + valid_gpus = [] + for micro_batch in micro_batches: + if batch_size % micro_batch == 0: + + max_gpus = batch_size // micro_batch + if min_valid_gpus <= max_gpus <= max_valid_gpus: + valid_gpus.append(max_gpus) + + # find all factors less than max_gpus / 2 + for i in range(1, max_gpus // 2 + 1): + if i > max_valid_gpus: + break + if i < min_valid_gpus: + continue + if max_gpus % i == 0: + valid_gpus.append(i) + valid_gpus = set(valid_gpus) + valid_gpus = sorted(list(valid_gpus)) + return valid_gpus + + +def get_best_candidates(candidate_batch_sizes, micro_batches, min_gpus, max_gpus, prefer_larger): + + max_valid_gpus = 0 + valid_gpus = None + final_batch_size = int(min(micro_batches)) + + for batch_size in candidate_batch_sizes: + + current_valid_gpus = get_valid_gpus(batch_size, micro_batches, min_gpus, max_gpus) + + if (len(current_valid_gpus) > max_valid_gpus or (len(current_valid_gpus) == max_valid_gpus and + ((prefer_larger and batch_size > final_batch_size) or + (not prefer_larger and batch_size < final_batch_size)))): + max_valid_gpus = len(current_valid_gpus) + valid_gpus = current_valid_gpus + final_batch_size = batch_size + + return final_batch_size, valid_gpus + + +def _get_compatible_gpus_v01(micro_batches, + max_acceptable_batch_size, + min_gpus=None, + max_gpus=None, + prefer_larger=True): + '''We use two heuristics to compute the batch size + 1. We use the Lowest Common Multiple of the micro-batches + as the base batch size and scale it by a HCN such that the result is + the largest batch size less than the max_acceptable batch size + 2. We use each of the micro batches as a base and scale it + by a HCN such that the result is the largest batch size less than the + max_acceptable batch size. + + We then use brute force to count the number of compatible GPU count for + each of the aforementioned cases, and return the batch size with the most number of + compatible GPU counts in the min-max GPU range if provided, other wise + we return the batch size with the most number of total compatible GPU counts. + + Returns: + final_batch_size + valid_gpus + ''' + min_gpus = min_gpus or 1 + max_gpus = max_gpus or max_acceptable_batch_size // min(micro_batches) + + if not all(mb <= max_acceptable_batch_size for mb in micro_batches): + raise ValueError(f"All micro batches must be less than \ + or equal to max_acceptable_batch_size: {max_acceptable_batch_size}") + + lcm = np.lcm.reduce(micro_batches) + + base_list = [] + base_list.extend(micro_batches) + base_list.append(lcm) + + candidate_batch_sizes = get_candidate_batch_sizes(base_list, max_acceptable_batch_size) + + final_batch_size, valid_gpus = get_best_candidates(candidate_batch_sizes, micro_batches, min_gpus, max_gpus, + prefer_larger) + + return final_batch_size, valid_gpus + + +def _get_compatible_gpus_v02(micro_batches, + max_acceptable_batch_size, + current_num_gpus, + min_gpus=None, + max_gpus=None, + prefer_larger=True, + num_gpus_per_node=1, + model_parallel_size=1): + ''' + Returns: + final_batch_size + valid_gpus + micro-batch size + ''' + if num_gpus_per_node % model_parallel_size != 0: + raise ElasticityError( + f"In Elasticity v0.2, number of GPUs per node:" \ + f"{num_gpus_per_node} should be divisible by " \ + f"model parallel size {model_parallel_size}") + + def get_microbatch(final_batch_size): + candidate_microbatch = None + + for micro_batch in micro_batches: + if final_batch_size // current_num_gpus % micro_batch == 0: + if candidate_microbatch is None: + candidate_microbatch = micro_batch + if prefer_larger and candidate_microbatch < micro_batch: + candidate_microbatch = micro_batch + return candidate_microbatch + + dp_size_per_node = num_gpus_per_node // model_parallel_size + + final_batch_size, valid_world_size = _get_compatible_gpus_v01( + micro_batches, + int(max_acceptable_batch_size / dp_size_per_node), + int(min_gpus / num_gpus_per_node), + int(max_gpus / num_gpus_per_node), # Passing number of max nodes as Elasticity v2 works at node level + prefer_larger=prefer_larger) + + final_batch_size = int(final_batch_size) * dp_size_per_node + valid_dp_world_size = [i * dp_size_per_node for i in valid_world_size] + if current_num_gpus // model_parallel_size in valid_dp_world_size: + candidate_microbatch = get_microbatch(final_batch_size) + return final_batch_size, valid_dp_world_size, candidate_microbatch + + current_dp_size = (current_num_gpus / num_gpus_per_node) * dp_size_per_node + candidate_batch_sizes = [] + for micro_batch in micro_batches: + min_batch_size = micro_batch * current_dp_size + + factor = math.floor(max_acceptable_batch_size / float(min_batch_size)) + candidate_batch_sizes.append(factor * min_batch_size) + + used_microbatch = None + if prefer_larger: + candidate_batch_size = max(candidate_batch_sizes) + else: + candidate_batch_size = min(candidate_batch_sizes) + + candidate_microbatch = get_microbatch(candidate_batch_size) + + return candidate_batch_size, [int(current_dp_size)], candidate_microbatch + + +def _compatible_ds_version_check(target_deepspeed_version: str): + min_version = pkg_version.parse(MINIMUM_DEEPSPEED_VERSION) + target_version = pkg_version.parse(target_deepspeed_version) + + err_str = f"Target deepspeed version of {target_deepspeed_version} is not compatible " \ + f"with minimum version {MINIMUM_DEEPSPEED_VERSION} supporting elasticity." + if target_version < min_version: + raise ElasticityError(err_str) + return True + + +def elasticity_enabled(ds_config: dict): + if ELASTICITY not in ds_config: + return False + return ds_config[ELASTICITY].get(ENABLED, ENABLED_DEFAULT) + + +def ensure_immutable_elastic_config(runtime_elastic_config_dict: dict): + """ + Ensure the resource scheduler saw the same elastic config we are using at runtime + """ + if DEEPSPEED_ELASTICITY_CONFIG in os.environ: + scheduler_elastic_config_dict = json.loads(os.environ[DEEPSPEED_ELASTICITY_CONFIG]) + scheduler_elastic_config = ElasticityConfig(scheduler_elastic_config_dict) + runtime_elastic_config = ElasticityConfig(runtime_elastic_config_dict) + err_str = "Elastic config '{}={}' seen by resource scheduler does not match config passed to runtime {}={}" + if runtime_elastic_config.max_acceptable_batch_size != scheduler_elastic_config.max_acceptable_batch_size: + raise ElasticityConfigError( + err_str.format('max_acceptable_batch_size', scheduler_elastic_config.max_acceptable_batch_size, + 'max_acceptable_batch_size', runtime_elastic_config.max_acceptable_batch_size)) + if runtime_elastic_config.micro_batches != scheduler_elastic_config.micro_batches: + raise ElasticityConfigError( + err_str.format('micro_batches', scheduler_elastic_config.micro_batches, 'micro_batches', + runtime_elastic_config.micro_batches)) + if runtime_elastic_config.version != scheduler_elastic_config.version: + raise ElasticityConfigError( + err_str.format('version', scheduler_elastic_config.version, 'version', runtime_elastic_config.version)) + else: + logger.warning("Unable to find DEEPSPEED_ELASTICITY_CONFIG environment variable, cannot " \ + "guarantee resource scheduler will scale this job using compatible GPU counts.") + + +def compute_elastic_config(ds_config: dict, target_deepspeed_version: str, world_size=0, return_microbatch=False): + """Core deepspeed elasticity API. Given an elastic config (similar to the example below) + DeepSpeed will compute a total train batch size corresponding valid GPU count list that + provides a high level of elasticity. Elasticity in this case means we are safe to scale + the training job up/down across the GPU count list *without* any negative impacts on + training convergence. This is achievable primarily due to DeepSpeed's gradient accumulation + feature which allows us to decompose a global training batch size into: + micro-batch-size * gradient-accumulation-steps * world-size. + + "elasticity": { + "enabled": true, + "max_train_batch_size": 2000, + "micro_batch_sizes": [2,4,6], + "min_gpus": 1, + "max_gpus" : 10000 + "min_time": 20 + "version": 0.1 + } + + Intended to be called both by scheduling infrastructure and deepspeed runtime. + For the same `ds_config` we should return deterministic results. + + Args: + ds_config (dict): DeepSpeed config dictionary/json + target_deepspeed_version (str): When called from scheduling + infrastructure we want to ensure that the target deepspeed version is + compatible with the elasticity version used in the backend. + world_size (int, optional): Intended/current DP world size, will do some sanity + checks to ensure world size is actually valid with the config. + return_microbatch (bool, optional): whether to return micro batch size or not. + + Raises: + ElasticityConfigError: Missing required elasticity config or elasticity disabled + ElasticityError: If target deepspeed version is not compatible with current version + + Returns: + final_batch_size (int): total batch size used for training + valid_gpus (list(int)): list of valid GPU counts with this config + micro_batch_size (int, optional): if world_size is provided will return + specific micro batch size + """ + if not isinstance(ds_config, dict): + raise ValueError("Expected ds_config to be a dictionary but received " \ + f"a {type(ds_config)}, containing: {ds_config}") + + if ELASTICITY not in ds_config: + raise ElasticityConfigError(f"'{ELASTICITY}' is missing from config json," \ + " please add it if running an elastic training job.") + + elastic_config_dict = ds_config[ELASTICITY] + if not elastic_config_dict.get(ENABLED, ENABLED_DEFAULT): + raise ElasticityConfigError("Elasticity is disabled, please enable it " \ + "('enabled':true) if running an elastic training job.") + + elastic_config = ElasticityConfig(elastic_config_dict) + model_parallel_size = elastic_config.model_parallel_size + num_gpus_per_node = elastic_config.num_gpus_per_node + + if model_parallel_size > 1 and float(elastic_config.version) != 0.2: + raise ElasticityConfigError(f"Elasticity V{elastic_config.version} " \ + f"does not support model-parallel training. Given model-parallel size: " \ + f"{model_parallel_size}") + + if float(elastic_config.version) > LATEST_ELASTICITY_VERSION: + raise ElasticityConfigError("Attempting to run elasticity version " \ + f"{elastic_config.version} but runtime only supports up " \ + f"to {LATEST_ELASTICITY_VERSION}") + + # Ensure target deepspeed version works with intended elasticity version + if not _compatible_ds_version_check(target_deepspeed_version): + raise ElasticityError("Unable to run elasticity on target deepspeed version of" \ + f" {target_deepspeed_version}, currently {__version__}") + + if float(elastic_config.version) == 0.1: + final_batch_size, valid_gpus = _get_compatible_gpus_v01( + micro_batches=elastic_config.micro_batches, + max_acceptable_batch_size=elastic_config.max_acceptable_batch_size, + min_gpus=elastic_config.min_gpus, + max_gpus=elastic_config.max_gpus, + prefer_larger=elastic_config.prefer_larger_batch_size) + # ensure batch size is int dtype + final_batch_size = int(final_batch_size) + elif float(elastic_config.version) == 0.2: + if world_size != 0: + current_num_gpus = world_size + else: + if "WORLD_SIZE" in os.environ and \ + os.getenv('WORLD_SIZE').isnumeric(): + current_num_gpus = int(os.getenv('WORLD_SIZE')) + else: + WORLD_SIZE = os.getenv('WORLD_SIZE') + raise ElasticityConfigError( + 'Elasticity V 0.2 needs WORLD_SIZE '\ + 'to compute valid batch size. '\ + 'Either give it as argument to function compute_elastic_config '\ + 'or set it as an environment variable. '\ + f'Value of WORLD_SIZE as environment variable is {WORLD_SIZE}') + + final_batch_size, valid_gpus, candidate_microbatch_size = _get_compatible_gpus_v02( + micro_batches=elastic_config.micro_batches, + max_acceptable_batch_size=elastic_config.max_acceptable_batch_size, + current_num_gpus=current_num_gpus, + min_gpus=elastic_config.min_gpus, + max_gpus=elastic_config.max_gpus, + prefer_larger=elastic_config.prefer_larger_batch_size, + num_gpus_per_node=num_gpus_per_node, + model_parallel_size=model_parallel_size) + # ensure batch size is int dtype + final_batch_size = int(final_batch_size) + else: + raise NotImplementedError(f"Unable to find elastic logic for version: {elastic_config.version}") + + logger.info(f"Valid World Size (GPUs / Model Parallel Size): {valid_gpus}") + + if world_size > 0: + if world_size not in valid_gpus: + raise ElasticityIncompatibleWorldSize(f"World size ({world_size}) is not valid " \ + f"with the current list of valid GPU counts: {valid_gpus}") + + # Pick largest valid micro batch size + micro_batch_size = None + for mbsz in sorted(list(set(elastic_config.micro_batches)), reverse=True): + if final_batch_size // world_size % mbsz == 0: + micro_batch_size = mbsz + break + assert micro_batch_size is not None, "Unable to find divisible micro batch size" \ + f" world_size={world_size}, final_batch_size={final_batch_size}, and " \ + f" micro_batches={elastic_config.micro_batches}." + return final_batch_size, valid_gpus, micro_batch_size + + if return_microbatch: + # Pick a valid micro batch size + if float(elastic_config.version) == 0.2: + return final_batch_size, valid_gpus, candidate_microbatch_size + else: + micro_batch_size = None + for mbsz in sorted(list(set(elastic_config.micro_batches)), reverse=True): + if final_batch_size // world_size % mbsz == 0: + micro_batch_size = mbsz + break + assert micro_batch_size is not None, "Unable to find divisible micro batch size" \ + f" world_size={world_size}, final_batch_size={final_batch_size}, and " \ + f" micro_batches={elastic_config.micro_batches}." + return final_batch_size, valid_gpus, micro_batch_size + + return final_batch_size, valid_gpus diff --git a/lib/python3.12/site-packages/deepspeed/elasticity/utils.py b/lib/python3.12/site-packages/deepspeed/elasticity/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..78ae0352cf6eaeab95eed20e9c8259a843f25b93 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/elasticity/utils.py @@ -0,0 +1,14 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.utils.torch import required_torch_version + + +def is_torch_elastic_compatible(): + ''' + Helper to lookup torch version. Elastic training is + introduced in 1.11.x + ''' + return required_torch_version(min_version=1.11) diff --git a/lib/python3.12/site-packages/deepspeed/launcher/__init__.py b/lib/python3.12/site-packages/deepspeed/launcher/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/launcher/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2645cd0c83b07bbc118af9d5211e72a1f7d9aedf Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0119dea2ab746bab1c7a23404266f3bf41712e53 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/launch.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/launch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ddab91e6a42a8cb296ed77429c779a4e2a370ba Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/launch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/launcher_helper.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/launcher_helper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb18e2350741e7f990776034550cedbb53701d89 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/launcher_helper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/multinode_runner.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/multinode_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82e6b4c645d1f75eed7296bbee27f76c1d97697c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/multinode_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/runner.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0488ae1d2afbaf70bf3b4b460f6eb77d6f5706b4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/launcher/__pycache__/runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/launcher/constants.py b/lib/python3.12/site-packages/deepspeed/launcher/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..366ae0b236f81a94b1a4ed2cecc01a1f3e1aa69e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/launcher/constants.py @@ -0,0 +1,16 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +PDSH_LAUNCHER = 'pdsh' +PDSH_MAX_FAN_OUT = 1024 + +OPENMPI_LAUNCHER = 'openmpi' +MPICH_LAUNCHER = 'mpich' +IMPI_LAUNCHER = 'impi' +SLURM_LAUNCHER = 'slurm' +MVAPICH_LAUNCHER = 'mvapich' +MVAPICH_TMP_HOSTFILE = '/tmp/deepspeed_mvapich_hostfile' + +ELASTIC_TRAINING_ID_DEFAULT = "123456789" diff --git a/lib/python3.12/site-packages/deepspeed/launcher/launch.py b/lib/python3.12/site-packages/deepspeed/launcher/launch.py new file mode 100644 index 0000000000000000000000000000000000000000..079182a3b45b86455c376511e239924ecff931dd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/launcher/launch.py @@ -0,0 +1,359 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +DeepSpeed launcher, this is similar to torch's distributed.launch but supports +additional features such as arbitrary gpu exclusion. + +deepspeed.launcher.launch is intended to be run on a single worker node and +will spawn several worker sub-processes depending on how many devices/ranks +are on the worker. +""" + +import sys +import subprocess +import os +import json +import base64 +import time +import signal +import psutil +from collections import defaultdict +from typing import Dict +from argparse import ArgumentParser, REMAINDER +from deepspeed.accelerator import get_accelerator +from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT +from ..nebula.constants import DLTS_POD_ENV_PATH +from ..utils import logger, get_numactl_cmd +from ..elasticity import is_torch_elastic_compatible +from .constants import ELASTIC_TRAINING_ID_DEFAULT + +PID_FILE_BASEPATH = "/tmp" + + +def parse_args(): + parser = ArgumentParser(description="DeepSpeed distributed training launch" + " utility that creates multiple distributed" + " processes on a single node") + + # Optional arguments for the launch helper + parser.add_argument("--node_rank", + type=int, + default=0, + help="The rank of the node for multi-node distributed " + "training") + parser.add_argument("--master_addr", + default="127.0.0.1", + type=str, + help="Master node (rank 0)'s address, should be either" + " the IP address or the hostname of node 0, for" + " single node multi-proc training, the" + " --master_addr can simply be 127.0.0.1") + parser.add_argument("--master_port", + default=TORCH_DISTRIBUTED_DEFAULT_PORT, + type=int, + help="Master node (rank 0)'s free port that needs to " + "be used for communication during distributed " + "training") + parser.add_argument("--world_info", default="None", type=str, help="world info base64 encoded dictionary") + + parser.add_argument("--module", + action="store_true", + help="Change each process to interpret the launch " + "script as a Python module, executing with the same " + "behavior as 'python -m'.") + + parser.add_argument("--no_python", + action="store_true", + help="Skip prepending the training script with " + "'python' - just execute it directly.") + + parser.add_argument("--enable_elastic_training", action="store_true", help="Enable elastic training support.") + + parser.add_argument("--min_elastic_nodes", type=int, default=-1, help="Min number of nodes in elastic training.") + + parser.add_argument("--max_elastic_nodes", type=int, default=-1, help="Max number of nodes in elastic training.") + + parser.add_argument("--no_local_rank", + action="store_true", + help="Do not pass local_rank as an argument when calling " + "the user's training script.") + + parser.add_argument("--save_pid", + type=int, + default=0, + help="main launching process pid, for internal pid tracking") + + parser.add_argument("--enable_each_rank_log", + default="None", + type=str, + help="redirect the stdout and stderr from each rank into different log files") + + parser.add_argument("--bind_cores_to_rank", + action="store_true", + help="Bind each rank to different cores of the host. " + "This improves host efficiency especially for CPU backend") + + parser.add_argument("--bind_core_list", + type=str, + default=None, + help="List of cores to bind to with comma separated list of " + "numbers and range. i.e. 1,3-5,7 => [1,3,4,5,7]. When not " + "specified, all cores on system would be used rank binding") + + # positional + parser.add_argument("training_script", + type=str, + help="The full path to the single GPU training " + "program/script to be launched in parallel, " + "followed by all the arguments for the " + "training script") + + # rest from the training program + parser.add_argument('training_script_args', nargs=REMAINDER) + return parser.parse_args() + + +# Adapted from https://psutil.readthedocs.io/en/latest/#kill-process-tree +def terminate_process_tree(pid): + process = psutil.Process(pid) + children = process.children(recursive=True) + children.append(process) + for child in children: + try: + child.terminate() + except psutil.NoSuchProcess: + pass + gone, alive = psutil.wait_procs(children, timeout=30) + for p in alive: + p.kill() + + +def main(): + args = parse_args() + current_env = os.environ.copy() + + for k in current_env.keys(): + if "NCCL" in k: + logger.info(f"{args.node_rank} {k}={current_env[k]}") + + if args.world_info == "None": + raise ValueError("world_info can not be None") + world_info = base64.urlsafe_b64decode(args.world_info) + world_info = json.loads(world_info) + + logger.info(f"WORLD INFO DICT: {world_info}") + node_list = list(world_info.keys()) + args.nnodes = len(node_list) + local_node = node_list[args.node_rank] + local_accelerator_ids = world_info[local_node] + num_local_procs = len(local_accelerator_ids) + logger.info(f"nnodes={args.nnodes}, num_local_procs={num_local_procs}, node_rank={args.node_rank}") + + global_rank_mapping = defaultdict(list) + curr_global_rank = 0 + dist_world_size = 0 + for node_id in node_list: + gids = world_info[node_id] + dist_world_size += len(gids) + for gid in gids: + global_rank_mapping[node_id].append(curr_global_rank) + curr_global_rank += 1 + logger.info(f"global_rank_mapping={global_rank_mapping}") + logger.info(f"dist_world_size={dist_world_size}") + + get_accelerator().set_visible_devices_envs(current_env, local_accelerator_ids) + for env in get_accelerator().visible_devices_envs(): + logger.info(f"Setting {env}={current_env[env]}") + + # set PyTorch distributed related environmental variables + current_env["MASTER_ADDR"] = args.master_addr + current_env["MASTER_PORT"] = str(args.master_port) + current_env["WORLD_SIZE"] = str(dist_world_size) + current_env["CROSS_RANK"] = str(args.node_rank) + current_env["CROSS_SIZE"] = str(args.nnodes) + current_env["LOCAL_SIZE"] = str(num_local_procs) + + if args.save_pid: + print(f"launcher pid: {os.getpid()}") + + pid_file = None + if args.save_pid: + launcher_pid = os.getpid() + pid_file = os.path.join(PID_FILE_BASEPATH, f"{args.save_pid}.deepspeed") + assert not os.path.isfile(pid_file), "pid file exists but shouldn't" + with open(pid_file, 'w') as fd: + fd.write(f"{launcher_pid}") + + if not is_torch_elastic_compatible(): + if args.enable_elastic_training: + logger.info(f"Disabling elastic training support as \ + PyTorch version should be greater than 1.11.x") + args.enable_elastic_training = False + + if os.path.exists(DLTS_POD_ENV_PATH): + with open(DLTS_POD_ENV_PATH) as file: + lines = file.readlines() + lines = [line.rstrip() for line in lines] + for line in lines: + if line.startswith('export FC_TASKROLE_NAME') or line.startswith('export FC_TASK_INDEX'): + key_val = line.split()[1] + key, val = key_val.split('=') + current_env[key] = val + + processes = [] + cmd = [] + + if not args.enable_elastic_training: + if args.enable_each_rank_log != "None": + # prepare the log path and the file name prefix + if os.path.isfile(args.enable_each_rank_log): + raise ValueError(f"{args.enable_each_rank_log} should not be a file, it should be a directory.") + if not os.path.exists(args.enable_each_rank_log): + try: + os.makedirs(args.enable_each_rank_log) + except Exception as e: + print(e) + raise ValueError(f"unable to create directory {args.enable_each_rank_log} for each rank log.") + log_name_prefix = time.strftime("%Y%m%d%H%M%S", time.localtime()) + + for local_proc in range(0, num_local_procs): + # each process's rank + dist_rank = global_rank_mapping[local_node][local_proc] + local_rank = dist_rank % num_local_procs + current_env["RANK"] = str(dist_rank) + current_env["LOCAL_RANK"] = str(local_rank) + + # spawn the processes + cmd = [] + if args.bind_cores_to_rank: + cores_per_rank, numactl_cmd = get_numactl_cmd(args.bind_core_list, num_local_procs, local_rank) + current_env["OMP_NUM_THREADS"] = f"{cores_per_rank}" + cmd = cmd + numactl_cmd + if not args.no_python: + cmd.append(sys.executable) + cmd.append("-u") + if args.module: + cmd.append("-m") + else: + if args.module: + raise ValueError("Don't use both the '--no_python' flag" + " and the '--module' flag at the same time.") + cmd.append(args.training_script) + # A user may not want to pass local_rank as a keyword arg so we make this optional. + if not args.no_local_rank: + cmd.append(f"--local_rank={local_rank}") + cmd += args.training_script_args + + if args.enable_each_rank_log != "None": + log_file = os.path.join(args.enable_each_rank_log, f"{log_name_prefix}_rank{dist_rank}.log") + log_fd = open(log_file, 'w') + process = subprocess.Popen(cmd, env=current_env, stdout=log_fd, stderr=log_fd) + else: + process = subprocess.Popen(cmd, env=current_env) + # logs the command from processes + logger.info(f"process {process.pid} spawned with command: {cmd}") + processes.append(process) + else: + from ..elasticity import DSElasticAgent + from torch.distributed.elastic.rendezvous import RendezvousParameters + from torch.distributed.elastic.agent.server.api import WorkerSpec + import torch.distributed.elastic.rendezvous.registry as rdzv_registry + from torch.distributed.elastic.multiprocessing import Std + + if args.min_elastic_nodes == -1: + args.min_elastic_nodes = 1 + if args.max_elastic_nodes == -1: + args.max_elastic_nodes = args.nnodes + assert args.max_elastic_nodes > 0 and args.min_elastic_nodes > 0, "Max and Min nodes should be positive" + + current_env["NCCL_ASYNC_ERROR_HANDLING"] = str(1) + + # Get config and arguments + cmd = [] + if not args.no_python: + cmd = [sys.executable, "-u"] + if args.module: + cmd.append("-m") + else: + if args.module: + raise ValueError("Don't use both the '--no_python' flag" + " and the '--module' flag at the same time.") + cmd.append(args.training_script) + cmd += args.training_script_args + cmd_args = cmd[1:] + + rdzv_configs: Dict[str, str] = {'timeout': 100} + run_id = os.environ.get("ELASTIC_RUN_ID", ELASTIC_TRAINING_ID_DEFAULT) + + # Creating config for rendezvous class + rdzv_parameters = RendezvousParameters(backend='c10d', + endpoint=args.master_addr + ":" + str(args.master_port), + run_id=run_id, + min_nodes=args.min_elastic_nodes, + max_nodes=args.max_elastic_nodes, + **rdzv_configs) + + spec = WorkerSpec( + role='trainer', + local_world_size=num_local_procs, + entrypoint=cmd[0], + args=cmd[1:], + rdzv_handler=rdzv_registry.get_rendezvous_handler(rdzv_parameters), + max_restarts=100, + monitor_interval=5, + redirects=Std.from_str("0"), + tee=Std.from_str("0"), + master_addr=None, + master_port=None, + ) + agent = DSElasticAgent(spec, current_env) + agent.run() + + sig_names = {2: "SIGINT", 15: "SIGTERM"} + last_return_code = None + + def sigkill_handler(signum, frame): + for process in processes: + logger.info(f"Killing subprocess {process.pid}") + try: + terminate_process_tree(process.pid) + except Exception: + pass + if last_return_code is not None: + logger.error(f"{cmd} exits with return code = {last_return_code}") + sys.exit(last_return_code) + if signum in sig_names: + logger.info(f"Main process received {sig_names[signum]}, exiting") + if args.save_pid: + if os.path.isfile(pid_file): + os.remove(pid_file) + sys.exit(1) + + # pass SIGINT/SIGTERM to children if the parent is being terminated + signal.signal(signal.SIGINT, sigkill_handler) + signal.signal(signal.SIGTERM, sigkill_handler) + + alive_processes = set(processes) + while len(alive_processes): + finished_processes = [] + for process in alive_processes: + if process.poll() is None: + # the process is still running + continue + else: + if process.returncode != 0: + last_return_code = process.returncode # for sigkill_handler + sigkill_handler(signal.SIGTERM, None) # not coming back + else: + # exited cleanly + logger.info(f"Process {process.pid} exits successfully.") + finished_processes.append(process) + alive_processes = set(alive_processes) - set(finished_processes) + + time.sleep(1) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/deepspeed/launcher/launcher_helper.py b/lib/python3.12/site-packages/deepspeed/launcher/launcher_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..05ce14bcc52ed9dd2c6182b30be541a1873569dc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/launcher/launcher_helper.py @@ -0,0 +1,108 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import sys +import argparse +import subprocess +from deepspeed.utils import logger +from deepspeed.launcher.constants import MPICH_LAUNCHER + + +def parse_args(args=None): + parser = argparse.ArgumentParser(description="DeepSpeed launcher helper to map environment variables for" + "multi-node/multi-gpu training jobs.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + parser.add_argument("--launcher", + default=MPICH_LAUNCHER, + type=str, + help="(optional) choose launcher backend for multi-node " + "training. Options currently include MPICH.") + + parser.add_argument("--module", + action="store_true", + help="Change each process to interpret the launch " + "script as a Python module, executing with the same " + "behavior as 'python -m'.") + + parser.add_argument("--no_python", + action="store_true", + help="Skip prepending the training script with " + "'python' - just execute it directly.") + + parser.add_argument("user_script", type=str, help="User script to launch, followed by any required " + "arguments.") + + parser.add_argument('user_args', nargs=argparse.REMAINDER) + + parser.add_argument("--bind_cores_to_rank", + action="store_true", + help="Bind each rank to different cores of the host") + + parser.add_argument("--bind_core_list", + type=str, + default=None, + help="List of cores to bind to with comma separated list of " + "numbers and range. i.e. 1,3-5,7 => [1,3,4,5,7]. When not " + "specified, all cores on system would be used rank binding") + + return parser.parse_args(args=args) + + +def env_mapping(env, rank_name_list=None, local_rank_name_list=None): + rank = None + for rank_name in rank_name_list: + if rank_name in env: + if rank == None: + rank = env.get(rank_name) + elif rank != env.get(rank_name): + raise EnvironmentError(f"rank number doesn't match!") + if rank == None: + raise EnvironmentError(f"rank number is not in current env!") + env['RANK'] = rank + + local_rank = None + for local_rank_name in local_rank_name_list: + if local_rank_name in env: + if local_rank == None: + local_rank = env.get(local_rank_name) + elif local_rank != env.get(local_rank_name): + raise EnvironmentError(f"local_rank number doesn't match!") + if local_rank == None: + raise EnvironmentError(f"rank number is not in current env!") + env['LOCAL_RANK'] = local_rank + + return env + + +def main(args=None): + args = parse_args(args) + + env = os.environ.copy() + + args.launcher = args.launcher.lower() + if args.launcher == MPICH_LAUNCHER: + rank_name_list = ["PMIX_RANK"] + ["PMI_RANK"] + local_rank_name_list = ["PALS_LOCAL_RANKID"] + ["MPI_LOCALRANKID"] + env = env_mapping(env, rank_name_list=rank_name_list, local_rank_name_list=local_rank_name_list) + else: + raise NotImplementedError(f"Unknown launcher {args.launcher}") + + python_exec = [] + if not args.no_python: + python_exec += [sys.executable, "-u"] + if args.module: + python_exec.append("-m") + cmd = python_exec + [args.user_script] + args.user_args + + logger.info(f"launcher_helper cmd = {' '.join(cmd)}") + + result = subprocess.Popen(cmd, env=env, close_fds=False) + result.wait() + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/deepspeed/launcher/multinode_runner.py b/lib/python3.12/site-packages/deepspeed/launcher/multinode_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..218fb02be766b140f39076aff3f70e1f74dcd0b1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/launcher/multinode_runner.py @@ -0,0 +1,493 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import sys +import shutil +import subprocess +import warnings +import re +from shlex import split +from abc import ABC, abstractmethod +from deepspeed.accelerator import get_accelerator +from ..utils import logger, get_numactl_cmd +from .constants import PDSH_MAX_FAN_OUT, MVAPICH_TMP_HOSTFILE + + +class MultiNodeRunner(ABC): + + def __init__(self, args, world_info_base64): + self.args = args + self.validate_args() + self.user_arguments = self.parse_user_args() + self.user_script = args.user_script + self.world_info_base64 = world_info_base64 + self.exports = {} + + @abstractmethod + def backend_exists(self): + """Return whether the corresponding backend exists""" + + @abstractmethod + def get_cmd(self, environment, active_resources): + """Return the command to execute on node""" + + def add_export(self, key, var): + var = var.strip() + if re.search(r'[^\w@%+=:,./-]', var): + var = f"\"{var}\"" + self.exports[key.strip()] = var + + def parse_user_args(self): + return self.args.user_args + + @property + def name(self): + """Return the name of the backend""" + return self.__class__.__name__ + + def validate_args(self): + """Validate self.args""" + + +class PDSHRunner(MultiNodeRunner): + + def __init__(self, args, world_info_base64): + super().__init__(args, world_info_base64) + + def backend_exists(self): + return shutil.which('pdsh') + + def parse_user_args(self): + processed_args = [] + for arg in self.args.user_args: + # With pdsh, if we are passing a string as an argument, it will get + # split on whitespace. To avoid this and support strings that + # contain '"', we do this extra processing step: + if " " in arg: + arg = '"{}"'.format(arg.replace('"', '\\"')) + processed_args.append(arg) + return processed_args + + @property + def name(self): + return "pdsh" + + def get_cmd(self, environment, active_resources): + environment['PDSH_RCMD_TYPE'] = 'ssh' + if self.args.ssh_port is not None: # only specify ssh port if it is specified + environment["PDSH_SSH_ARGS_APPEND"] = f"{environment.get('PDSH_SSH_ARGS_APPEND', '')} \ + -p {self.args.ssh_port}" + + active_workers = ",".join(active_resources.keys()) + logger.info("Running on the following workers: %s" % active_workers) + + # PDSH flags for max node fan out and specific hosts to launch on + # See https://linux.die.net/man/1/pdsh for flag details + pdsh_cmd_args = ['pdsh', '-S', '-f', str(PDSH_MAX_FAN_OUT), '-w', active_workers] + split( + self.args.launcher_args) + + exports = "" + for key, val in self.exports.items(): + exports += "export {}={}; ".format(key, val) + + # https://linux.die.net/man/1/pdsh + # %n will be replaced by pdsh command + deepspeed_launch = [ + exports, f"cd {os.path.abspath('.')};", sys.executable, "-u", "-m", "deepspeed.launcher.launch", + f'--world_info={self.world_info_base64}', "--node_rank=%n", f"--master_addr={self.args.master_addr}", + f"--master_port={self.args.master_port}" + ] + if self.args.no_python: + deepspeed_launch.append("--no_python") + if self.args.module: + deepspeed_launch.append("--module") + if self.args.no_local_rank: + deepspeed_launch.append("--no_local_rank") + if self.args.save_pid: + deepspeed_launch += ["--save_pid", f"{os.getpid()}"] + if self.args.enable_each_rank_log: + deepspeed_launch.append(f"--enable_each_rank_log={self.args.enable_each_rank_log}") + if self.args.elastic_training: + deepspeed_launch.append("--enable_elastic_training") + deepspeed_launch.append(f"--max_elastic_nodes={self.args.max_elastic_nodes}") + deepspeed_launch.append(f"--min_elastic_nodes={self.args.min_elastic_nodes}") + + cmd_to_search = [i + "\\" for i in deepspeed_launch[2:6]] + + kill_command = pdsh_cmd_args + ["pkill -f ", " ".join(cmd_to_search)[:-2]] + return pdsh_cmd_args + deepspeed_launch + [self.user_script] + self.user_arguments, kill_command, environment + + +class OpenMPIRunner(MultiNodeRunner): + + def __init__(self, args, world_info_base64, resource_pool): + super().__init__(args, world_info_base64) + self.resource_pool = resource_pool + self.add_export('UCX_TLS', 'tcp') + + def backend_exists(self): + #TODO: if IB is available we should suggestion mvapich + return shutil.which('ompi_info') + + @property + def name(self): + return "openmpi" + + def validate_args(self): + super().validate_args() + + # Validate and set MPI environment variables + self._setup_mpi_environment() + + #TODO: Allow for include/exclude at node-level but not gpu-level + if self.args.include != "" or self.args.exclude != "": + raise ValueError(f"{self.name} backend does not support worker include/exclusion") + if self.args.num_nodes != -1 or self.args.num_gpus != -1: + raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") + + def _setup_mpi_environment(self): + """Sets up MPI-related environment variables or raises an error if they're missing.""" + + required_vars = ['OMPI_COMM_WORLD_LOCAL_RANK', 'OMPI_COMM_WORLD_RANK', 'OMPI_COMM_WORLD_SIZE'] + + # Check if all these are present + if not all(var in os.environ for var in required_vars): + raise EnvironmentError("MPI environment variables are not set. " + "Ensure you are running the script with an MPI-compatible launcher.") + + # Now safe to read all + os.environ['LOCAL_RANK'] = os.environ['OMPI_COMM_WORLD_LOCAL_RANK'] + os.environ['RANK'] = os.environ['OMPI_COMM_WORLD_RANK'] + os.environ['WORLD_SIZE'] = os.environ['OMPI_COMM_WORLD_SIZE'] + + def get_cmd(self, environment, active_resources): + total_process_count = sum(self.resource_pool.values()) + + launcher_args = split(self.args.launcher_args) + + # If btl_tcp_if_include option is provided through launcher_args, we use it. Otherwise, we add + # `--mca btl_tcp_if_include eth0` option as a default value for compatibility. + btl_tcp_opt = ['--mca', 'btl_tcp_if_include', 'eth0'] + if len(launcher_args) >= 2: + for i in range(len(launcher_args) - 1): + if launcher_args[i] in ['-mca', '--mca'] and launcher_args[i + 1] == 'btl_tcp_if_include': + btl_tcp_opt = [] + break + + mpirun_cmd = [ + 'mpirun', + '-n', + f'{total_process_count}', + '-hostfile', + f'{self.args.hostfile}', + '--mca', + 'btl', + '^openib', + ] + btl_tcp_opt + launcher_args + + export_cmd = [] + for k, v in self.exports.items(): + export_cmd += ['-x', "{}={}".format(k, v)] + + python_exec = [] + if not self.args.no_python: + python_exec = [sys.executable, "-u"] + if self.args.module: + python_exec.append("-m") + + return mpirun_cmd + export_cmd + python_exec + [self.user_script] + self.user_arguments + + +class MPICHRunner(MultiNodeRunner): + + def __init__(self, args, world_info_base64, resource_pool): + super().__init__(args, world_info_base64) + self.resource_pool = resource_pool + + def backend_exists(self): + #TODO: if IB is available we should suggestion mpich + return shutil.which('mpirun') #mpich_info + + @property + def name(self): + return "mpich" + + def validate_args(self): + super().validate_args() + #TODO: Allow for include/exclude at node-level but not gpu-level + if self.args.include != "" or self.args.exclude != "": + raise ValueError(f"{self.name} backend does not support worker include/exclusion") + + if self.args.num_nodes != -1 or self.args.num_gpus != -1: + raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") + + def get_cmd(self, environment, active_resources): + devices_per_node = self.resource_pool.values() + total_process_count = sum(devices_per_node) + process_per_node = list(devices_per_node)[0] + if not all([n == process_per_node for n in devices_per_node]): + raise ValueError("MPICH requires same number of devices per node") + + mpirun_cmd = [ + 'mpirun', + '-n', + f'{total_process_count}', + '-ppn', + f'{process_per_node}', + ] + split(self.args.launcher_args) + export_cmd = [] + + for k, v in self.exports.items(): + export_cmd += ['-genv', "{}={}".format(k, v)] + + export_cmd += ['-genv', 'MASTER_ADDR', str(self.args.master_addr)] + export_cmd += ['-genv', 'MASTER_PORT', str(self.args.master_port)] + export_cmd += ['-genv', 'WORLD_SIZE', str(total_process_count)] + export_cmd += ['-genv', 'LOCAL_SIZE', str(process_per_node)] + + export_cmd += ['-hosts'] + hosts = "" + for i, host in enumerate(self.resource_pool.keys()): + if i == 0: + hosts = f"{host}" + else: + hosts += f",{host}" + export_cmd += [hosts] + + helper_args = ["--launcher"] + [self.args.launcher] + python_exec = [] + if not self.args.no_python: + python_exec += [sys.executable, "-u"] + if self.args.module: + python_exec.append("-m") + helper_args.append("--module") + else: + helper_args.append("--no_python") + + helper_cmd = str(os.path.dirname(os.path.realpath(__file__))) + '/launcher_helper.py' + helper_cmd = [helper_cmd] + helper_args + [self.user_script] + self.user_arguments + + return mpirun_cmd + export_cmd + python_exec + helper_cmd + + +class IMPIRunner(MultiNodeRunner): + + def __init__(self, args, world_info_base64, resource_pool): + super().__init__(args, world_info_base64) + self.resource_pool = resource_pool + + def backend_exists(self): + #TODO: if IB is available we should suggestion mpich + return shutil.which('mpirun') #mpich_info + + @property + def name(self): + return "impi" + + def validate_args(self): + super().validate_args() + #TODO: Allow for include/exclude at node-level but not gpu-level + if self.args.include != "" or self.args.exclude != "": + raise ValueError(f"{self.name} backend does not support worker include/exclusion") + + if self.args.num_nodes != -1 or self.args.num_gpus != -1: + raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") + + def get_cmd(self, environment, active_resources): + devices_per_node = self.resource_pool.values() + total_process_count = sum(devices_per_node) + process_per_node = list(devices_per_node)[0] + if not all([n == process_per_node for n in devices_per_node]): + raise ValueError("Intel MPI requires same number of devices per node") + + mpirun_cmd = [ + 'mpirun', + '-ppn', + f'{process_per_node}', + ] + split(self.args.launcher_args) + export_cmd = [] + + for k, v in self.exports.items(): + export_cmd += ['-genv', f'{k}', f'{v}'] + + if self.args.bind_cores_to_rank: + cores_per_rank, _ = get_numactl_cmd(self.args.bind_core_list, process_per_node, 0) + export_cmd += ['-genv', 'OMP_NUM_THREADS', str(cores_per_rank)] + + export_cmd += ['-genv', 'MASTER_ADDR', str(self.args.master_addr)] + export_cmd += ['-genv', 'MASTER_PORT', str(self.args.master_port)] + export_cmd += ['-genv', 'WORLD_SIZE', str(total_process_count)] + export_cmd += ['-genv', 'LOCAL_SIZE', str(process_per_node)] + + # turn off IMPI core binding, use deepspeed's own core binding + export_cmd += ['-genv', 'I_MPI_PIN', '0'] + + export_cmd += ['-hosts'] + hosts = "" + for i, host in enumerate(self.resource_pool.keys()): + if i == 0: + hosts = f"{host}" + else: + hosts += f",{host}" + export_cmd += [hosts] + + per_host_cmd = [] + + for i in range(total_process_count): + local_rank = i % process_per_node + python_exec = [] + if self.args.bind_cores_to_rank: + _, numactl_cmd = get_numactl_cmd(self.args.bind_core_list, process_per_node, local_rank) + python_exec += numactl_cmd + + if not self.args.no_python: + python_exec += [sys.executable, "-u"] + if self.args.module: + python_exec.append("-m") + env_mapping = ['-env', 'RANK', str(i)] + env_mapping += ['-env', 'LOCAL_RANK', str(local_rank)] + if i == 0: + per_host_cmd = ['-n', '1'] + env_mapping + python_exec + [self.user_script] + self.user_arguments + else: + per_host_cmd = per_host_cmd + [':', '-n', '1'] + env_mapping + python_exec + [self.user_script + ] + self.user_arguments + print(mpirun_cmd + export_cmd + per_host_cmd) + return mpirun_cmd + export_cmd + per_host_cmd + + +class SlurmRunner(MultiNodeRunner): + + def __init__(self, args, world_info_base64, resource_pool): + super().__init__(args, world_info_base64) + self.resource_pool = resource_pool + + def backend_exists(self): + return shutil.which('sinfo') + + @property + def name(self): + return 'slurm' + + def get_cmd(self, environment, active_resources): + assert not getattr(self.args, 'detect_nvlink_pairs', + False), "slurm backend does not support remapping visible devices" + total_process_count = sum(self.resource_pool.values()) + srun_cmd = [ + 'srun', + '-n', + f'{total_process_count}', + ] + split(self.args.launcher_args) + + if getattr(self.args, 'slurm_comment', ''): + srun_cmd += ['--comment', self.args.slurm_comment] + + if self.args.include != "": + srun_cmd.append('--include') + srun_cmd.append(f'{self.args.include}') + if self.args.exclude != "": + srun_cmd.append('--exclude') + srun_cmd.append(f'{self.args.exclude}') + if self.args.num_nodes > 0: + srun_cmd.append('--nodes') + srun_cmd.append(f'{self.args.num_nodes}') + if self.args.num_gpus > 0: + srun_cmd.append('--gpus') + srun_cmd.append(f'{self.args.num_gpus}') + + exports = '--export=ALL' + for key, val in self.exports.items(): + exports += f",{key}={val}" + + python_exec = [sys.executable, "-u"] + command = srun_cmd + [exports] + python_exec + [self.user_script] + self.user_arguments + return command + + +class MVAPICHRunner(MultiNodeRunner): + + def __init__(self, args, world_info_base64, resource_pool): + super().__init__(args, world_info_base64) + self.resource_pool = resource_pool + + # Disable the CMA kernel module, not available on Ubuntu systems + self.add_export('MV2_SMP_USE_CMA', '0') + + # If we fail this will output more verbose logging + self.add_export('MV2_DEBUG_SHOW_BACKTRACE', '1') + + # Enabled cuda-aware communication + if get_accelerator().device_name() == 'cuda': + self.add_export('MV2_USE_CUDA', '1') + + # Support deep learning frameworks: http://hidl.cse.ohio-state.edu/userguide/horovod/ + self.add_export('MV2_SUPPORT_DL', '1') + + # Support MPI_THREAD_MULTIPLE + self.add_export('MV2_ENABLE_AFFINITY', '0') + + # Performance tuning flags for allgather + self.add_export('MV2_INTER_ALLGATHER_TUNING', '5') + self.add_export('MV2_CUDA_USE_NAIVE', '0') + + def backend_exists(self): + #TODO: if IB is available we should suggestion mvapich + mpiname_exists = shutil.which('mpiname') + exists = False + if not mpiname_exists: + warnings.warn("mpiname does not exist, mvapich is not installed properly") + else: + results = subprocess.check_output(['mpiname']) + mpiname_results = results.decode('utf-8').strip() + if "MVAPICH2-GDR" in mpiname_results: + exists = True + else: + warnings.warn(f"Expected MVAPICH2-GDR as return for mpiname but received {mpiname_results}") + return exists + + @property + def name(self): + return "mvapich" + + def validate_args(self): + super().validate_args() + #TODO: Allow for include/exclude at node-level but not gpu-level + if self.args.include != "" or self.args.exclude != "": + raise ValueError(f"{self.name} backend does not support worker include/exclusion") + if self.args.num_nodes != -1 or self.args.num_gpus != -1: + raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") + + def get_cmd(self, environment, active_resources): + devices_per_node = self.resource_pool.values() + total_process_count = sum(devices_per_node) + process_per_node = list(devices_per_node)[0] + if not all([n == process_per_node for n in devices_per_node]): + raise ValueError("mvapich requires same number of devices per node") + + with open(MVAPICH_TMP_HOSTFILE, 'w') as fd: + for host in self.resource_pool.keys(): + fd.write(f'{host}\n') + + mpirun_cmd = [ + 'mpirun', + '-np', + f'{total_process_count}', + '-ppn', + f'{process_per_node}', + '--hostfile', + f'{MVAPICH_TMP_HOSTFILE}', + ] + split(self.args.launcher_args) + + export_cmd = [] + for k, v in self.exports.items(): + export_cmd += ['-env', "{}={}".format(k, v)] + + python_exec = [] + if not self.args.no_python: + python_exec = [sys.executable, "-u"] + if self.args.module: + python_exec.append("-m") + + return mpirun_cmd + export_cmd + python_exec + [self.user_script] + self.user_arguments diff --git a/lib/python3.12/site-packages/deepspeed/launcher/runner.py b/lib/python3.12/site-packages/deepspeed/launcher/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..2f7daed4ef437d3f07e6a4629f79e3f297fb344e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/launcher/runner.py @@ -0,0 +1,631 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +DeepSpeed runner is the main front-end to launching multi-worker +training jobs with DeepSpeed. By default this uses pdsh to parallel +ssh into multiple worker nodes and launch all the necessary processes +per rank for training. +""" + +import os +import re +import sys +import json +import base64 +import argparse +import subprocess +import collections +from copy import deepcopy +import signal +import time +from typing import Tuple, List, Dict +from collections import defaultdict +import shlex + +from .multinode_runner import PDSHRunner, OpenMPIRunner, MVAPICHRunner, SlurmRunner, MPICHRunner, IMPIRunner +from .constants import PDSH_LAUNCHER, OPENMPI_LAUNCHER, MVAPICH_LAUNCHER, SLURM_LAUNCHER, MPICH_LAUNCHER, IMPI_LAUNCHER +from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT +from ..nebula.constants import NEBULA_EXPORT_ENVS +from ..utils import logger + +from ..autotuning import Autotuner +from deepspeed.accelerator import get_accelerator + +DLTS_HOSTFILE = "/job/hostfile" +EXPORT_ENVS = ['MLFLOW', 'PYTHON', 'MV2', 'UCX'] +EXPORT_ENVS += NEBULA_EXPORT_ENVS +DEEPSPEED_ENVIRONMENT_NAME = os.getenv("DS_ENV_FILE", ".deepspeed_env") +DEEPSPEED_ENVIRONMENT_PATHS = [os.path.expanduser("~"), '.'] +PDSH_MAX_FAN_OUT = 1024 + +# On AISC compute, each node sets environment variables independently, want to prevent +# exporting rank-0 env variables in case of heterogeneous compute. +EXCLUDE_ENVS = {'AISC_JOB_NAME': ['NCCL_IB_HCA', 'UCX_NET_DEVICES']} + + +def parse_args(args=None): + parser = argparse.ArgumentParser(description="DeepSpeed runner to help launch distributed " + "multi-node/multi-gpu training jobs.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + parser.add_argument("-H", + "--hostfile", + type=str, + default=DLTS_HOSTFILE, + help="Hostfile path (in MPI style) that defines the " + "resource pool available to the job (e.g., " + "worker-0 slots=4)") + + parser.add_argument("-i", + "--include", + type=str, + default="", + help='''Specify hardware resources to use during execution. + String format is + NODE_SPEC[@NODE_SPEC ...], + where + NODE_SPEC=NAME[:SLOT[,SLOT ...]]. + If :SLOT is omitted, include all slots on that host. + Example: -i "worker-0@worker-1:0,2" will use all slots + on worker-0 and slots [0, 2] on worker-1. + ''') + + parser.add_argument("-e", + "--exclude", + type=str, + default="", + help='''Specify hardware resources to NOT use during execution. + Mutually exclusive with --include. Resource formatting + is the same as --include. + Example: -e "worker-1:0" will use all available + resources except slot 0 on worker-1. + ''') + + parser.add_argument("--num_nodes", + type=int, + default=-1, + help="Total number of worker nodes to run on, this will use " + "the top N hosts from the given hostfile.") + + parser.add_argument("--min_elastic_nodes", + type=int, + default=-1, + help="Minimum number of nodes to run elastic training on. " + "Default is 1 when elastic training is enabled") + + parser.add_argument("--max_elastic_nodes", + type=int, + default=-1, + help="Maximum number of nodes to run elastic training on. " + "Default is num_nodes when elastic training is enabled") + + parser.add_argument("--num_gpus", + "--num_accelerators", + type=int, + default=-1, + help="Max number of GPUs to use on each node, will use " + "[0:N) GPU ids on each node.") + + parser.add_argument("--master_port", + default=TORCH_DISTRIBUTED_DEFAULT_PORT, + type=int, + help="(optional) Port used by PyTorch distributed for " + "communication during training.") + + parser.add_argument("--master_addr", + default="", + type=str, + help="(optional) IP address of node 0, will be " + "inferred via 'hostname -I' if not specified.") + + parser.add_argument("--node_rank", + default=-1, + type=int, + help="ID of each node in the range [0:N). " + "Only required when --no_ssh is set.") + + parser.add_argument("--launcher", + default=PDSH_LAUNCHER, + type=str, + help="(optional) choose launcher backend for multi-node " + "training. Options currently include PDSH, OpenMPI, MVAPICH, SLURM, MPICH, IMPI.") + + parser.add_argument("--launcher_args", + default="", + type=str, + help="(optional) pass launcher specific arguments as a " + "single quoted argument.") + + parser.add_argument("--module", + action="store_true", + help="Change each process to interpret the launch " + "script as a Python module, executing with the same " + "behavior as 'python -m'.") + + parser.add_argument("--no_python", + action="store_true", + help="Skip prepending the training script with " + "'python' - just execute it directly.") + + parser.add_argument("--no_local_rank", + action="store_true", + help="Do not pass local_rank as an argument when calling " + "the user's training script.") + + parser.add_argument("--no_ssh", + action="store_true", + help="Launch training independently on each node without ssh setup.") + + parser.add_argument("--no_ssh_check", + action="store_true", + help="Do not perform ssh check in multi-node launcher model") + + parser.add_argument("--force_multi", + action="store_true", + help="Force multi-node launcher mode, helps in cases where user " + "wants to launch on single remote node.") + + parser.add_argument("--save_pid", + action="store_true", + help="Save file containing launcher process id (pid) at /tmp/.ds, " + "where is the pid of the first process that invoked `deepspeed`. " + "Useful when launching deepspeed processes programmatically.") + + parser.add_argument("--enable_each_rank_log", + default="None", + type=str, + help="redirect the stdout and stderr from each rank into different log files") + + parser.add_argument("--autotuning", + default="", + choices=["tune", "run"], + type=str, + help="Run DeepSpeed autotuner to discover optimal configuration parameters " + "before running job.") + + parser.add_argument("--elastic_training", + action="store_true", + help="Enable elastic training support in DeepSpeed.") + + parser.add_argument("user_script", type=str, help="User script to launch, followed by any required " + "arguments.") + + parser.add_argument('user_args', nargs=argparse.REMAINDER) + + parser.add_argument("--bind_cores_to_rank", + action="store_true", + help="Bind each rank to different cores of the host") + + parser.add_argument("--bind_core_list", + type=str, + default=None, + help="List of cores to bind to with comma separated list of " + "numbers and range. i.e. 1,3-5,7 => [1,3,4,5,7]. When not " + "specified, all cores on system would be used rank binding") + + parser.add_argument("--ssh_port", type=int, default=None, help="SSH port to use for remote connections") + + return parser.parse_args(args=args) + + +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: + hostfile_text = fd.readlines() + + return _parse_hostfile(hostfile_text) + + +def _parse_hostfile(hostfile_lines): + # Regex matches one or more non-whitespace characters (\S+) at the start of + # the line, followed by one or more whitespace characters (\s+), followed + # by the string "slots=", followed by one or more digits (\d+). + pattern = r'^(\S+)\s+slots=(\d+)' + + resource_pool = collections.OrderedDict() + + for line in hostfile_lines: + line = line.strip() + match = re.search(pattern, line) + if line.startswith("#") or line == "": + # hostfile comment or empty line, ignore + continue + elif match: + host = match.group(1) + num_slots = int(match.group(2)) + if host in resource_pool: + logger.error(f"Bad hostfile text: {hostfile_lines}") + raise ValueError(f"Hostfile contains multiple entries for {host}, unable to proceed with launching") + resource_pool[host] = num_slots + else: + logger.error(f"Bad hostfile text: {hostfile_lines}") + raise ValueError(f"Hostfile contains a bad entry: {line}, unable to proceed with launching") + + if len(resource_pool) == 0: + logger.error(f"Bad hostfile text: {hostfile_lines}") + raise ValueError("Hostfile is empty or not formatted correctly, unable to proceed with launching.") + + return resource_pool + + +def _stable_remove_duplicates(data): + # Create a new list in the same order as original but with duplicates + # removed, should never be more than ~16 elements so simple is best + new_list = [] + for x in data: + if x not in new_list: + new_list.append(x) + return new_list + + +def parse_node_config(node_config: str) -> Tuple[str, List[int]]: + SLOT_LIST_START = ':' + SLOT_SEP = ',' + + if SLOT_LIST_START not in node_config: + return node_config, [] + + hostname, slots = node_config.split(SLOT_LIST_START) + slots = [int(x) for x in slots.split(SLOT_SEP)] + + return hostname, slots + + +def parse_node_config_list(node_config_list: List[str]) -> Dict[str, List[int]]: + NODE_SEP = '@' + + node_configs = defaultdict(list) + + for node_config in node_config_list.split(NODE_SEP): + hostname, slots = parse_node_config(node_config) + node_configs[hostname] += slots + + return {k: sorted(list(set(v))) for k, v in node_configs.items()} + + +def parse_resource_filter(host_info, include_str="", exclude_str=""): + '''Parse an inclusion or exclusion string and filter a hostfile dictionary. + + String format is NODE_SPEC[@NODE_SPEC ...], where + NODE_SPEC = NAME[:SLOT[,SLOT ...]]. + If :SLOT is omitted, include/exclude all slots on that host. + + Examples: + include_str="worker-0@worker-1:0,2" will use all slots on worker-0 and + slots [0, 2] on worker-1. + exclude_str="worker-1:0" will use all available resources except + slot 0 on worker-1. + ''' + + # Ensure include/exclude are mutually exclusive + if (include_str != "") and (exclude_str != ""): + raise ValueError('include_str and exclude_str are mutually exclusive.') + + # no-op + if (include_str == "") and (exclude_str == ""): + return host_info + + # Either build from scratch or remove items + filtered_hosts = dict() + if include_str: + parse_str = include_str + if exclude_str != "": + filtered_hosts = deepcopy(host_info) + parse_str = exclude_str + + # foreach node in the list + for hostname, slots in parse_node_config_list(parse_str).items(): + # Node can either be alone or node:slot,slot,slot + if len(slots) > 0: + # sanity checks + if hostname not in host_info: + raise ValueError(f"Hostname '{hostname}' not found in hostfile") + for slot in slots: + if slot not in host_info[hostname]: + raise ValueError(f"No slot '{slot}' specified on host '{hostname}'") + + # If include string, build the list from here + if include_str: + filtered_hosts[hostname] = slots + elif exclude_str: + for slot in slots: + logger.info(f'removing {slot} from {hostname}') + filtered_hosts[hostname].remove(slot) + + # User just specified the whole node + else: + # sanity check hostname + if hostname not in host_info: + raise ValueError(f"Hostname '{hostname}' not found in hostfile") + + if include_str: + filtered_hosts[hostname] = host_info[hostname] + elif exclude_str: + filtered_hosts[hostname] = [] + + # Post-processing to remove duplicates and empty nodes + del_keys = [] + for hostname in filtered_hosts: + # Remove duplicates + filtered_hosts[hostname] = _stable_remove_duplicates(filtered_hosts[hostname]) + # Remove empty hosts + if len(filtered_hosts[hostname]) == 0: + del_keys.append(hostname) + for name in del_keys: + del filtered_hosts[name] + + # Lastly, go over filtered_hosts and convert to a OrderedDict() to ensure + # we map ranks to nodes correctly by maintaining host_info ordering. + ordered_hosts = collections.OrderedDict() + for host in host_info: + if host in filtered_hosts: + ordered_hosts[host] = filtered_hosts[host] + + return ordered_hosts + + +def parse_inclusion_exclusion(resource_pool, inclusion, exclusion): + active_resources = collections.OrderedDict() + node_configs = parse_node_config_list(inclusion) + + for hostname, slots in resource_pool.items(): + active_resources[hostname] = node_configs[hostname] if hostname in node_configs else list(range(slots)) + + return parse_resource_filter(active_resources, include_str=inclusion, exclude_str=exclusion) + + +def encode_world_info(world_info): + world_info_json = json.dumps(world_info).encode('utf-8') + world_info_base64 = base64.urlsafe_b64encode(world_info_json).decode('utf-8') + return world_info_base64 + + +def run_autotuning(args, active_resources): + tuner = Autotuner(args, active_resources) + logger.info("[Start] Running autotuning") + + tuner.tune() + tuner.print_tuning_results() + + logger.info("[End] Running autotuning") + tuner.write_optimal_config() + + if args.autotuning == "run": + tuner.run_after_tuning() + + +def parse_num_nodes(str_num_nodes: str, elastic_training: bool): + node_list = str_num_nodes.split(":") + + if len(node_list) == 1: + min_nodes, max_nodes = int(node_list[0]), -1 + elif len(node_list) == 2 and elastic_training: + min_nodes, max_nodes = int(node_list[0]), int(node_list[1]) + elif len(node_list) == 2 and not elastic_training: + raise RuntimeError("MIN:MAX format is only supported in elastic training") + else: + raise RuntimeError("num_nodes {} is not in MIN:MAX format".format(str_num_nodes)) + + return min_nodes, max_nodes + + +def main(args=None): + args = parse_args(args) + + if args.elastic_training: + assert args.master_addr != "", "Master Addr is required when elastic training is enabled" + + resource_pool = fetch_hostfile(args.hostfile) + + # respect VISIBLE_DEVICES for a single node and no explicit resource filters + visible_devices_env = get_accelerator().visible_devices_envs()[0] + visible_devices = os.environ.get(visible_devices_env, "") + if not resource_pool and len(visible_devices): + detected_str = f"Detected VISIBLE_DEVICES={visible_devices}" + if len(args.include) or len(args.exclude) or args.num_nodes > 1 or args.num_gpus > 0: + print( + f"{detected_str} but ignoring it because one or several of --include/--exclude/--num_gpus/--num_nodes cl args were used. If you want to use CUDA_VISIBLE_DEVICES don't pass any of these arguments to deepspeed." + ) + else: + args.include = f"localhost:{visible_devices}" + print(f"{detected_str}: setting --include={args.include}") + del os.environ[visible_devices_env] + + if args.num_nodes >= 0 or args.num_gpus >= 0: + if args.include != "" or args.exclude != "": + raise ValueError("Cannot specify num_nodes/gpus with include/exclude") + + multi_node_exec = True + if not resource_pool: + resource_pool = {} + device_count = get_accelerator().device_count() + if device_count == 0: + raise RuntimeError("Unable to proceed, no GPU resources available") + resource_pool['localhost'] = device_count + args.master_addr = "127.0.0.1" + multi_node_exec = False + + if not multi_node_exec and args.num_nodes > 1: + raise ValueError("Num nodes is >1 but no extra nodes available via hostfile") + + active_resources = parse_inclusion_exclusion(resource_pool, args.include, args.exclude) + env = os.environ.copy() + + # validate that passwordless-ssh is workly properly with this hostfile + if multi_node_exec and not args.no_ssh_check and not args.no_ssh: + first_host = list(active_resources.keys())[0] + try: + ssh_check_cmd = ("ssh -o PasswordAuthentication=no " + + (f"-p {args.ssh_port} " if args.ssh_port is not None else "") + f"{first_host} hostname") + safe_ssh_cmd = shlex.split(ssh_check_cmd) + subprocess.check_call(safe_ssh_cmd, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL) + except subprocess.CalledProcessError: + raise RuntimeError( + f"Using hostfile at {args.hostfile} but host={first_host} was not reachable via ssh. If you are running with a single node please remove {args.hostfile} or setup passwordless ssh." + ) + + if not args.master_addr: + assert multi_node_exec + first_host = list(active_resources.keys())[0] + ssh_check_cmd = "ssh " + if args.ssh_port is not None: + ssh_check_cmd += f" -p {args.ssh_port}" + ssh_check_cmd += f" {first_host} hostname -I" + hostname_cmd = shlex.split(ssh_check_cmd) + try: + result = subprocess.check_output(hostname_cmd) + except subprocess.CalledProcessError as err: + logger.error( + "Unable to detect suitable master address via 'hostname -I', please manually specify one via --master_addr" + ) + raise err + args.master_addr = result.decode('utf-8').split()[0] + if not args.master_addr: + raise RuntimeError( + "Unable to detect suitable master address via `hostname -I`, please manually specify one via --master_addr" + ) + logger.info(f"Using IP address of {args.master_addr} for node {first_host}") + + if args.autotuning != "": + run_autotuning(args, active_resources) + return + + if args.num_nodes > 0: + updated_active_resources = collections.OrderedDict() + for count, hostname in enumerate(active_resources.keys()): + if args.num_nodes == count: + break + updated_active_resources[hostname] = active_resources[hostname] + active_resources = updated_active_resources + + if args.num_gpus > 0: + updated_active_resources = collections.OrderedDict() + for hostname in active_resources.keys(): + updated_active_resources[hostname] = list(range(args.num_gpus)) + active_resources = updated_active_resources + + if args.elastic_training: + assert not args.no_local_rank, "--no_local_rank argument is not supported in Elastic training" + + if args.no_ssh: + assert (0 <= args.node_rank < + len(active_resources)), "Launching training without ssh, but --node_rank is not set correctly." + + # encode world info as base64 to make it easier to pass via command line + world_info_base64 = encode_world_info(active_resources) + + multi_node_exec = (args.force_multi or len(active_resources) > 1) and not args.no_ssh + + if not multi_node_exec: + deepspeed_launch = [ + sys.executable, "-u", "-m", "deepspeed.launcher.launch", f"--world_info={world_info_base64}", + f"--master_addr={args.master_addr}", f"--master_port={args.master_port}" + ] + if args.no_ssh: + deepspeed_launch.append(f"--node_rank={args.node_rank}") + if args.no_python: + deepspeed_launch.append("--no_python") + if args.module: + deepspeed_launch.append("--module") + if args.no_local_rank: + deepspeed_launch.append("--no_local_rank") + if args.save_pid: + deepspeed_launch += ["--save_pid", f"{os.getpid()}"] + if args.enable_each_rank_log: + deepspeed_launch.append(f"--enable_each_rank_log={args.enable_each_rank_log}") + if args.elastic_training: + deepspeed_launch.append("--enable_elastic_training") + deepspeed_launch.append(f"--max_elastic_nodes={args.max_elastic_nodes}") + deepspeed_launch.append(f"--min_elastic_nodes={args.min_elastic_nodes}") + if args.bind_cores_to_rank: + deepspeed_launch.append("--bind_cores_to_rank") + if args.bind_core_list is not None: + deepspeed_launch.append(f"--bind_core_list={args.bind_core_list}") + cmd = deepspeed_launch + [args.user_script] + args.user_args + else: + args.launcher = args.launcher.lower() + if args.launcher == PDSH_LAUNCHER: + runner = PDSHRunner(args, world_info_base64) + elif args.launcher == OPENMPI_LAUNCHER: + runner = OpenMPIRunner(args, world_info_base64, resource_pool) + elif args.launcher == MPICH_LAUNCHER: + runner = MPICHRunner(args, world_info_base64, resource_pool) + elif args.launcher == IMPI_LAUNCHER: + runner = IMPIRunner(args, world_info_base64, resource_pool) + elif args.launcher == MVAPICH_LAUNCHER: + runner = MVAPICHRunner(args, world_info_base64, resource_pool) + elif args.launcher == SLURM_LAUNCHER: + runner = SlurmRunner(args, world_info_base64, resource_pool) + else: + raise NotImplementedError(f"Unknown launcher {args.launcher}") + + if not runner.backend_exists(): + raise RuntimeError(f"launcher '{args.launcher}' not installed.") + + curr_path = os.path.abspath('.') + if 'PYTHONPATH' in env: + env['PYTHONPATH'] = curr_path + ":" + env['PYTHONPATH'] + else: + env['PYTHONPATH'] = curr_path + + excluded_vars = [] + for exclude_key, var_list in EXCLUDE_ENVS.items(): + if exclude_key in env.keys(): + # key exists in launcher env -> var list should be used + excluded_vars += var_list + + # load envs from accelerator + exports = EXPORT_ENVS + get_accelerator().export_envs() + for var in env.keys(): + if any([var.startswith(name) for name in exports]): + if not any([var == name for name in excluded_vars]): + runner.add_export(var, env[var]) + + for environ_path in DEEPSPEED_ENVIRONMENT_PATHS: + environ_file = os.path.join(environ_path, DEEPSPEED_ENVIRONMENT_NAME) + if os.path.isfile(environ_file): + logger.info(f"deepspeed_env file = {environ_file}") + with open(environ_file, 'r') as fd: + for var in fd.readlines(): + key, val = var.split('=', maxsplit=1) + runner.add_export(key, val) + + if args.launcher == PDSH_LAUNCHER: + cmd, kill_cmd, env = runner.get_cmd(env, active_resources) + else: + cmd = runner.get_cmd(env, active_resources) + + logger.info(f"cmd = {' '.join(cmd)}") + result = subprocess.Popen(cmd, env=env) + + def sigkill_handler(signum, frame): + result.send_signal(signal.SIGINT) + time.sleep(0.1) + result.send_signal(signal.SIGTERM) + result_kill = subprocess.Popen(kill_cmd, env=env) + result_kill.wait() + time.sleep(1) + sys.exit(1) + + if args.launcher == PDSH_LAUNCHER and multi_node_exec: + signal.signal(signal.SIGINT, sigkill_handler) + signal.signal(signal.SIGTERM, sigkill_handler) + + result.wait() + + # In case of failure must propagate the error-condition back to the caller (usually shell). The + # actual error and traceback should have been printed in the subprocess, so in order to avoid + # unnecessary noise we just quietly exit here with the same code as the subprocess + if result.returncode > 0: + sys.exit(result.returncode) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/deepspeed/linear/__init__.py b/lib/python3.12/site-packages/deepspeed/linear/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9931a95a0a40c449b40b9152a5b8419f66d60a10 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/linear/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .optimized_linear import OptimizedLinear +from .config import LoRAConfig, QuantizationConfig +from .context_manager import Init, init_lora diff --git a/lib/python3.12/site-packages/deepspeed/linear/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dac52f5a6a8dcc9a0a1de41e55ba0c0229aa310e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/linear/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a458ed13dce862d8376761a9643f2369a58a37e5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/linear/__pycache__/context_manager.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/context_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfbfbcbf4eb2028c2020bb777e3dfbc63d16884a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/context_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/linear/__pycache__/optimized_linear.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/optimized_linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4499f777b9c478ded4812edbdc73a5f2571e6fa Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/optimized_linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/linear/__pycache__/quantization.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/quantization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77b60a3abcb08ca255f2ed26c09ec282eaff3525 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/linear/__pycache__/quantization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/linear/config.py b/lib/python3.12/site-packages/deepspeed/linear/config.py new file mode 100644 index 0000000000000000000000000000000000000000..1459704a32c5f36a827d4427de96942e3b6fa06b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/linear/config.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from dataclasses import dataclass, field +from typing import List + +import torch + + +@dataclass +class LoRAConfig: + """ + Configuration settings for LoRAOptimizedLinear. + + Attributes: + lora_r (int): LoRA attention dimension, also known as the rank. Defaults is 64. + lora_alpha (float): LoRA scaling factor, default is 16. + base_weight_sharding (int): The degree to which the base weights are sharded, + should typically be set to the data-parallel world size to maximize the memory + reduction benefits. Defaults to 1, which means this feature is disabled. + offload (bool): offload frozen parameters to cpu when not in use + offload_ratio (float): ratio of parameters to offload to cpu when not in use + delay_lora_init (bool): initialize lora parameters at time of model init or allow manual init later + target_mods (str): target module names to apply LoRA to, defaults to llama-3.1 arch + """ + lora_r: int = 64 + lora_alpha: float = 16. + base_weight_sharding: int = 1 + offload: bool = False + offload_ratio: float = 0.0 + delay_lora_init: bool = False + target_mods: List[str] = field( + default_factory=lambda: ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj']) + + +@dataclass +class QuantizationConfig: + """ + Configuration settings for quantization for LoRAOptimizedLinear, QuantizedLinear, + and QuantizedParameter + + Attributes: + q_bits (int): The number of bits used for quantization. Default is 8. + mantissa_bits (int): The number of bits reserved for the mantissa in fixed-point quantization. Default is 3. + group_size (int): The number of elements used for quantization. Default is 512. + q_dtype (torch.dtype): The data type to quantize to. Default is uint8. (in CUDA, buffers are allocated as + uint8, but inside the kernels the quantization is done to fp8) + """ + q_bits: int = 8 + mantissa_bits: int = 3 + group_size: int = 512 + q_dtype: torch.dtype = torch.uint8 diff --git a/lib/python3.12/site-packages/deepspeed/linear/context_manager.py b/lib/python3.12/site-packages/deepspeed/linear/context_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..204fa0fe9c1da3ed2b1c8ccad2828257da536d57 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/linear/context_manager.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .optimized_linear import LoRAOptimizedLinear, OptimizedLinear + +import torch + +try: + import transformers +except ImportError: + transformers = None + + +def init_lora(model): + model.requires_grad_(False) + for m in model.modules(): + if isinstance(m, LoRAOptimizedLinear): + m.init_lora() + + +class Init(object): + """ + Init context wrapper similar in style to zero.Init. Allows for injecting OptimizedLinear during model + construction which will shard base weights and reduce overall memory usage during model init. Primarily + useful when initializing a model via transformers.AutoModelForCausalLM. + + Example usage: + lora_config = deepspeed.linear.LoRAConfig(..) + quant_config = deepspeed.linear.QuantizationConfig(..) + with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config): + model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-405B") + + """ + + def __init__(self, lora_config=None, quant_config=None): + self._orig_nn_linear = torch.nn.Linear + self._orig_causallm_pretrained = None + if transformers != None: + self._orig_causallm_pretrained = transformers.AutoModelForCausalLM.from_pretrained + self._orig_causallm_config = transformers.AutoModelForCausalLM.from_config + self.lora_config = lora_config + self.quant_config = quant_config + self._post_init_complete = False + + def __enter__(self): + + class OptLinearWrapper: + _orig_nn_linear = self._orig_nn_linear + _lora_config = self.lora_config + _quant_config = self.quant_config + + def __new__(self, *args, **kwargs): + self._lora_config.delay_lora_init = True + kwargs['lora_config'] = self._lora_config + kwargs['quantization_config'] = self._quant_config + kwargs['linear_cls'] = self._orig_nn_linear + return OptimizedLinear(*args, **kwargs) + + def _model_init(model): + if self.lora_config != None: + init_lora(model) + self._post_init_complete = True + return model + + # ensures non-lora params are frozen and lora weights are initialized + def from_pretrained(*args, **kwargs): + model = self._orig_causallm_pretrained(*args, **kwargs) + return _model_init(model) + + def from_config(*args, **kwargs): + model = self._orig_causallm_config(*args, **kwargs) + return _model_init(model) + + torch.nn.Linear = OptLinearWrapper + if transformers != None: + transformers.AutoModelForCausalLM.from_pretrained = from_pretrained + transformers.AutoModelForCausalLM.from_config = from_config + + def __exit__(self, *args, **kwargs): + torch.nn.Linear = self._orig_nn_linear + if not self._post_init_complete: + print('WARNING: For some reason LoRA modules are not initialized, this is usually done automatically ' + 'if using transformers via (AutoModelForCausalLM from_pretrained/from_config). ' + 'You must call `init_lora` on each module in order to use DeepSpeed LoRA, otherwise ' + 'you will error out during runtime.') + else: + transformers.AutoModelForCausalLM.from_pretrained = self._orig_causallm_pretrained + transformers.AutoModelForCausalLM.from_config = self._orig_causallm_config diff --git a/lib/python3.12/site-packages/deepspeed/linear/optimized_linear.py b/lib/python3.12/site-packages/deepspeed/linear/optimized_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..3720196aa2554b27bd235522a73447597ef7f97b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/linear/optimized_linear.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import math +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import is_dataclass +from deepspeed.accelerator import get_accelerator +import deepspeed.comm as dist + +from .config import LoRAConfig, QuantizationConfig +from .quantization import QuantizedParameter, QuantizedLinear + + +class OptimizedLinear(nn.Module): + """ + Optimized version of nn.Linear that adds features such as: + * LoRA w. base weight sharding + * FP [6,8,12] quantization + + Arguments: + input_dim: Required: size of each input sample + output_dim: Required: size of each output sample + bias: Optional: If set to False, the layer will not learn an additive bias. Default: False + lora_config: Optional: LoRAConfig defining lora features and base-weight-sharding degree + quantization_config: Optional: QuantizationConfig defining quantization features + dtype: Optional: parameter dtype, only supports bfloat16 currently + + Returns: + Returns a new nn.Module depending on the input config. Either native + torch.nn.Linear, QuantizedLinear, or the full-featured DSOptimizedLinear. + """ + + def __new__(self, + input_dim: int, + output_dim: int, + bias: bool = False, + lora_config: LoRAConfig = None, + quantization_config: QuantizationConfig = None, + device=None, + dtype=torch.bfloat16, + linear_cls=nn.Linear): + + if quantization_config is not None and not is_dataclass(quantization_config): + raise ValueError(f"Expecting QuantizationConfig but received {type(quantization_config)}") + if lora_config is not None and not is_dataclass(lora_config): + raise ValueError(f"Expecting LoRAConfig but received {type(lora_config)}") + if lora_config is None and quantization_config is None: + # Everything disabled, fall back to normal nn.Linear + self = linear_cls(input_dim, output_dim, bias=bias, dtype=dtype, device=device) + + elif lora_config: + # lora enabled, quantization may or may not be + self = LoRAOptimizedLinear(input_dim=input_dim, + output_dim=output_dim, + bias=bias, + lora_config=lora_config, + quantization_config=quantization_config, + dtype=dtype, + device=device, + linear_cls=linear_cls) + + elif quantization_config: + # only quantization enabled, no lora + self = QuantizedLinear(input_dim=input_dim, + output_dim=output_dim, + bias=bias, + quantization_config=quantization_config, + dtype=dtype) + return self + + +class LoRAOptimizedLinear(nn.Module): + + def __init__(self, + input_dim: int, + output_dim: int, + bias: bool = False, + lora_config: LoRAConfig = None, + quantization_config: QuantizationConfig = None, + device=None, + dtype=torch.bfloat16, + linear_cls=nn.Linear): + super().__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.bias = bias + self.lora_config = lora_config + self.quantization_config = quantization_config + self.device = get_accelerator().current_device_name() if device is None else device + self.linear_cls = linear_cls + self.dtype = dtype + assert self.lora_config is not None, "DSOptimizedLinear requires a LoRA config" + assert not self.bias, "bias=True is not supported by LoRAOptimizedLinear" + self.zero_shards = self.lora_config.base_weight_sharding + self.sharded_weight_size = int(float(self.input_dim) // self.zero_shards) + if self.zero_shards > 1: + assert self.zero_shards == dist.get_world_size( + ), "base weight sharding is only supported across world size" + w = torch.nn.Parameter(torch.empty(self.output_dim * self.sharded_weight_size, dtype=dtype), + requires_grad=False) + else: + w = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=dtype), requires_grad=False) + torch.nn.init.xavier_uniform_(w.reshape(self.sharded_weight_size, self.output_dim)) + + if self.quantization_config is not None: + assert dtype == torch.bfloat16, "only bfloat16 is supported when using quantization" + self.weight = QuantizedParameter(w, quantization_config=quantization_config) + else: + self.weight = w + + self.disabled = False + self._initialized = False + if not self.lora_config.delay_lora_init: + self.init_lora() + + def disable(self): + self.disabled = True + self.weight = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=self.dtype), + requires_grad=False) + + def init_lora(self): + if self.disabled: + return + + if self.quantization_config is not None: + # ensure quant-param wasn't stripped, in some cases transformers will do this during model init + if not isinstance(self.weight, QuantizedParameter): + self.weight = QuantizedParameter(self.weight, quantization_config=self.quantization_config) + + self._initialized = True + self.weight.requires_grad = False + + # Mark base weight to prevent broadcast and ensure proper offload behavior + self.weight.ds_optim_param = True + + self.lora_scaling_factor = self.lora_config.lora_alpha / self.lora_config.lora_r + + # Keeping lora weights in bf16 precision for ease of training. + self.lora_weight_1 = self.linear_cls(self.input_dim, + self.lora_config.lora_r, + bias=self.bias, + device=self.device, + dtype=self.dtype) + self.lora_weight_2 = self.linear_cls(self.lora_config.lora_r, + self.output_dim, + bias=self.bias, + device=self.device, + dtype=self.dtype) + + # initialize "A" with kaiming uniform and "B" with zeros following this + # https://github.com/huggingface/peft/blob/62122b5add8d6892f70c82eaef2147a6ba33b90b/src/peft/tuners/lora/layer.py#L155 + nn.init.kaiming_uniform_(self.lora_weight_1.weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_weight_2.weight) + self.lora_weight_1.weight.requires_grad = True + self.lora_weight_2.weight.requires_grad = True + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, + error_msgs): + if not any([target in prefix for target in self.lora_config.target_mods]): + # module does not match any target_mods, we must revert to normal nn.Linear via disable + self.disable() + return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, + unexpected_keys, error_msgs) + + if self.zero_shards > 1: + if not dist.is_initialized(): + raise RuntimeError( + "attempting to use optimized linear base weight sharding but torch-distributed is not initialized, please init first." + ) + rank = dist.get_rank() + shape_local = self.output_dim * self.sharded_weight_size + base_weight_name = f"{prefix}weight" + incoming_param = state_dict[base_weight_name] + state_dict[base_weight_name] = incoming_param.flatten().narrow(0, rank * shape_local, shape_local) + + return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, + error_msgs) + + def full_weight(self): + base_weight = self.weight + if getattr(base_weight, 'ds_offload', False): + # move to gpu so we can dequant and all-gather + assert base_weight.device == torch.device('cpu'), \ + f"expected base weight on cpu but found {base_weight.device}" + base_weight.offload(revert=True) + local_weight = base_weight.dequantized() if isinstance(base_weight, QuantizedParameter) else base_weight + base_weight.offload() + else: + local_weight = base_weight.dequantized() if isinstance(base_weight, QuantizedParameter) else base_weight + + tensor_out = torch.empty(self.output_dim * self.input_dim, + dtype=local_weight.dtype, + device=local_weight.device) + dist.all_gather_into_tensor(tensor_out, local_weight) + return tensor_out.reshape(self.output_dim, self.input_dim) + + def linear_without_F_linear(self, input, weight): + output = torch.mm(input.reshape(-1, input.shape[-1]), weight) + output = output.view(*input.shape[:-1], weight.shape[1]) + return output + + def forward(self, input_tensor): + if self.disabled: + return F.linear(input_tensor, self.weight) + assert self._initialized, "init_lora was never called, please initialize before proceeding" + + # Gather the sharded base weight + if self.zero_shards > 1: + with torch.no_grad(): + base_weight = self.full_weight() + elif self.quantization_config: + base_weight = self.weight.dequantized() + else: + base_weight = self.weight + + base_weight_output = F.linear(input_tensor, base_weight) + lora_output = self.lora_weight_2(self.lora_weight_1(input_tensor)) + return base_weight_output + self.lora_scaling_factor * lora_output diff --git a/lib/python3.12/site-packages/deepspeed/linear/quantization.py b/lib/python3.12/site-packages/deepspeed/linear/quantization.py new file mode 100644 index 0000000000000000000000000000000000000000..2023601be281734b68c434acda1f652191f7db8d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/linear/quantization.py @@ -0,0 +1,147 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import copy +import torch +import torch.nn as nn +import torch.nn.functional as F + +from typing import Optional + +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.fp_quantizer import Quantizer, FP_Quantize +from .config import QuantizationConfig + + +class QuantizedParameter(nn.Parameter): + """ + Quantized parameter class that implements weight quantization. Weights + are stored in quantized form on GPUs, and can be dequantized on-the-fly when + needed by the model. The weights are actually quantized during any `.to(device)`. + + Arguments: + data (Tensor): parameter tensor. + requires_grad (bool, optional): if the parameter requires gradient. Defaults + to False and is not supported to be True. Argument provided only for interface + compatibility with torch.nn.Parameter. + quantization_config (QuantizationConfig, optional): + quantizer (Quantizer, optional): Defaults to FP_Quantize but can be any quantizer + that implements deepspeed.ops.fp_quantizer.Quantizer. This argument is also + required since the quantizer is stashed in the Parameter itself, some models + may clone the Parameter by passing an attribute __dict__. For an example, see + tests/unit/linear/test_quant_param.py::TestQuantParam::test_hf_clone + """ + + def __new__( + cls, + data: Optional[torch.Tensor] = None, + requires_grad: bool = False, # quantized weights must be frozen + quantization_config: QuantizationConfig = None, + quantizer: Quantizer = None, + ): + if requires_grad: + raise ValueError(f"requires_grad=True is not supported with QuantizedParameter") + if data is None: + data = torch.empty(0) + self = torch.Tensor._make_subclass(cls, data, requires_grad) + self.quantization_config = QuantizationConfig() if quantization_config is None else quantization_config + if quantizer is not None: + self.quantizer = quantizer + else: + # if FPQuantizerBuilder is not compatible in this env this init will fail + self.quantizer = FP_Quantize(quantization_config=self.quantization_config) + self._ensure_quantized(self) + return self + + def _ensure_quantized(self, tensor: torch.Tensor): + # If the tensor is on the accelerator and is not quantized, then quantize it in-place. + if get_accelerator().on_accelerator(tensor) and tensor.dtype != self.quantization_config.q_dtype: + with get_accelerator().stream(get_accelerator().current_stream(tensor.device)): + tensor.data = self.quantizer.quantize(tensor.data, + q_bits=self.quantization_config.q_bits, + q_mantisa_bits=self.quantization_config.mantissa_bits) + assert tensor.dtype == self.quantization_config.q_dtype + + def dequantized(self) -> torch.Tensor: + """ + Return a tensor containing the dequantized weights of this parameter. + """ + if get_accelerator().on_accelerator(self.data) and self.data.dtype == self.quantization_config.q_dtype: + with get_accelerator().stream(get_accelerator().current_stream(self.data.device)): + return self.quantizer.dequantize(self.data, + q_bits=self.quantization_config.q_bits, + q_mantisa_bits=self.quantization_config.mantissa_bits) + return self.data + + def offload(self, revert=False): + if getattr(self, 'ds_offload', False): + if revert: + self.data = self.to(get_accelerator().current_device_name()) + else: + self.data = self.to('cpu') + + def __getstate__(self): + state = self.__dict__ + state["data"] = self.data + state["quantization_config"] = self.quantization_config + state["requires_grad"] = self.requires_grad + return state + + def __setstate__(self, state): + self.quantizer = state["quantizer"] + self.quantization_config = state["quantization_config"] + self.data = state["data"] + self.requires_grad = state["requires_grad"] + + def __deepcopy__(self, memo): + new_instance = type(self).__new__(type(self)) + state = self.__getstate__() + new_instance.__setstate__(state) + new_instance.quantizer = copy.deepcopy(state["quantizer"]) + new_instance.quantization_config = copy.deepcopy(state["quantization_config"]) + new_instance.data = copy.deepcopy(state["data"]) + return new_instance + + def __copy__(self): + new_instance = type(self).__new__(type(self)) + state = self.__getstate__() + new_instance.__setstate__(state) + return new_instance + + def cuda(self, device=None, non_blocking=False): + device = "cuda" if device is None else device + self.quantizer.to(device, non_blocking=non_blocking) + return self.to(device, non_blocking=non_blocking) + + def to(self, *args, **kwargs): + """ + Move the parameter to the given device. Then, if the device is a cuda device, + quantize it. + """ + tensor = super().to(*args, **kwargs) + self.quantizer.to(*args, **kwargs) + self._ensure_quantized(tensor) + return tensor + + +class QuantizedLinear(nn.Linear): + """ + Linear layer that implements weight quantization. Parameters + are stored via `QuantizedParameter` and are dequantized on-the-fly during any + forward pass. + """ + + def __init__(self, + input_dim: int, + output_dim: int, + bias: bool = False, + quantization_config: QuantizationConfig = None, + dtype=torch.bfloat16): + super().__init__(input_dim, output_dim, bias=bias, dtype=dtype) + assert dtype == torch.bfloat16, "currently only supports bfloat16 dtype" + self.weight = QuantizedParameter(self.weight.data, quantization_config=quantization_config) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return F.linear(input, self.weight.dequantized(), self.bias) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/__init__.py b/lib/python3.12/site-packages/deepspeed/model_implementations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ad95c58f76090188b5d9c73e1e3df4eb2b2c678f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .transformers.ds_transformer import DeepSpeedTransformerInference +from .transformers.clip_encoder import DSClipEncoder diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19270e39a971c07b86f1cfa7c8b5f20b1d708eea Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__init__.py b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..99f8d8622a160510d0cad9ef4f61c60d93c572ce Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4896bfcc94efab0d3ba1abeba6da4d544e3dfc9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25736a75bc98ec2a3451c71f957e9b9aad475fdf Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/unet.py b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..8d5ddd95437a42de5cb8552c529d20e422c792f2 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/unet.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.accelerator import get_accelerator +from ..features.cuda_graph import CUDAGraph + + +class DSUNet(CUDAGraph, torch.nn.Module): + + def __init__(self, unet, enable_cuda_graph=True): + super().__init__(enable_cuda_graph=enable_cuda_graph) + self.unet = unet + # SD pipeline accesses this attribute + self.in_channels = unet.in_channels + self.device = self.unet.device + self.dtype = self.unet.dtype + self.config = self.unet.config + self.fwd_count = 0 + self.unet.requires_grad_(requires_grad=False) + self.unet.to(memory_format=torch.channels_last) + self.cuda_graph_created = False + + def _graph_replay(self, *inputs, **kwargs): + for i in range(len(inputs)): + if torch.is_tensor(inputs[i]): + self.static_inputs[i].copy_(inputs[i]) + for k in kwargs: + if torch.is_tensor(kwargs[k]): + self.static_kwargs[k].copy_(kwargs[k]) + get_accelerator().replay_graph(self._cuda_graphs) + return self.static_output + + def forward(self, *inputs, **kwargs): + if self.enable_cuda_graph: + if self.cuda_graph_created: + outputs = self._graph_replay(*inputs, **kwargs) + else: + self._create_cuda_graph(*inputs, **kwargs) + outputs = self._graph_replay(*inputs, **kwargs) + return outputs + else: + return self._forward(*inputs, **kwargs) + + def _create_cuda_graph(self, *inputs, **kwargs): + # warmup to create the workspace and cublas handle + cuda_stream = torch.cuda.Stream() + cuda_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(cuda_stream): + for i in range(3): + ret = self._forward(*inputs, **kwargs) + torch.cuda.current_stream().wait_stream(cuda_stream) + + # create cuda_graph and assign static_inputs and static_outputs + self._cuda_graphs = get_accelerator().create_graph() + self.static_inputs = inputs + self.static_kwargs = kwargs + + with get_accelerator().capture_to_graph(self._cuda_graphs): + self.static_output = self._forward(*self.static_inputs, **self.static_kwargs) + + self.cuda_graph_created = True + + def _forward(self, + sample, + timestamp, + encoder_hidden_states, + return_dict=True, + cross_attention_kwargs=None, + timestep_cond=None, + added_cond_kwargs=None): + if cross_attention_kwargs: + return self.unet(sample, + timestamp, + encoder_hidden_states, + return_dict, + cross_attention_kwargs=cross_attention_kwargs) + else: + return self.unet(sample, timestamp, encoder_hidden_states, return_dict) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/vae.py b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..ce50ade647a8b85c1858e9d47e1f6a1133379467 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/diffusers/vae.py @@ -0,0 +1,151 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.accelerator import get_accelerator +from ..features.cuda_graph import CUDAGraph + + +class DSVAE(CUDAGraph, torch.nn.Module): + + def __init__(self, vae, enable_cuda_graph=True): + super().__init__(enable_cuda_graph=enable_cuda_graph) + self.vae = vae + self.config = vae.config + self.device = self.vae.device + self.dtype = self.vae.dtype + self.vae.requires_grad_(requires_grad=False) + self.decoder_cuda_graph_created = False + self.encoder_cuda_graph_created = False + self.all_cuda_graph_created = False + + def _graph_replay_decoder(self, *inputs, **kwargs): + for i in range(len(inputs)): + if torch.is_tensor(inputs[i]): + self.static_decoder_inputs[i].copy_(inputs[i]) + for k in kwargs: + if torch.is_tensor(kwargs[k]): + self.static_decoder_kwargs[k].copy_(kwargs[k]) + get_accelerator().replay_graph(self._decoder_cuda_graph) + return self.static_decoder_output + + def _decode(self, x, return_dict=True, generator=None): + return self.vae.decode(x, return_dict=return_dict) + + def _create_cuda_graph_decoder(self, *inputs, **kwargs): + # warmup to create the workspace and cublas handle + cuda_stream = torch.cuda.Stream() + cuda_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(cuda_stream): + for i in range(3): + ret = self._decode(*inputs, **kwargs) + torch.cuda.current_stream().wait_stream(cuda_stream) + + # create cuda_graph and assign static_inputs and static_outputs + self._decoder_cuda_graph = get_accelerator().create_graph() + self.static_decoder_inputs = inputs + self.static_decoder_kwargs = kwargs + + with get_accelerator().capture_to_graph(self._decoder_cuda_graph): + self.static_decoder_output = self._decode(*self.static_decoder_inputs, **self.static_decoder_kwargs) + + self.decoder_cuda_graph_created = True + + def decode(self, *inputs, **kwargs): + if self.enable_cuda_graph: + if self.decoder_cuda_graph_created: + outputs = self._graph_replay_decoder(*inputs, **kwargs) + else: + self._create_cuda_graph_decoder(*inputs, **kwargs) + outputs = self._graph_replay_decoder(*inputs, **kwargs) + return outputs + else: + return self._decode(*inputs, **kwargs) + + def _graph_replay_encoder(self, *inputs, **kwargs): + for i in range(len(inputs)): + if torch.is_tensor(inputs[i]): + self.static_encoder_inputs[i].copy_(inputs[i]) + for k in kwargs: + if torch.is_tensor(kwargs[k]): + self.static_encoder_kwargs[k].copy_(kwargs[k]) + get_accelerator().replay_graph(self._encoder_cuda_graph) + return self.static_encoder_output + + def _encode(self, x, return_dict=True): + return self.vae.encode(x, return_dict=return_dict) + + def _create_cuda_graph_encoder(self, *inputs, **kwargs): + # warmup to create the workspace and cublas handle + cuda_stream = torch.cuda.Stream() + cuda_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(cuda_stream): + for i in range(3): + ret = self._encode(*inputs, **kwargs) + torch.cuda.current_stream().wait_stream(cuda_stream) + + # create cuda_graph and assign static_inputs and static_outputs + self._encoder_cuda_graph = get_accelerator().create_graph() + self.static_encoder_inputs = inputs + self.static_encoder_kwargs = kwargs + + with get_accelerator().capture_to_graph(self._encoder_cuda_graph): + self.static_encoder_output = self._encode(*self.static_encoder_inputs, **self.static_encoder_kwargs) + + self.encoder_cuda_graph_created = True + + def encode(self, *inputs, **kwargs): + if self.enable_cuda_graph: + if self.encoder_cuda_graph_created: + outputs = self._graph_replay_encoder(*inputs, **kwargs) + else: + self._create_cuda_graph_encoder(*inputs, **kwargs) + outputs = self._graph_replay_encoder(*inputs, **kwargs) + return outputs + else: + return self._encode(*inputs, **kwargs) + + def _graph_replay(self, *inputs, **kwargs): + for i in range(len(inputs)): + if torch.is_tensor(inputs[i]): + self.static_inputs[i].copy_(inputs[i]) + for k in kwargs: + if torch.is_tensor(kwargs[k]): + self.static_kwargs[k].copy_(kwargs[k]) + get_accelerator().replay_graph(self._all_cuda_graph) + return self.static_output + + def forward(self, *inputs, **kwargs): + if self.enable_cuda_graph: + if self.cuda_graph_created: + outputs = self._graph_replay(*inputs, **kwargs) + else: + self._create_cuda_graph(*inputs, **kwargs) + outputs = self._graph_replay(*inputs, **kwargs) + return outputs + else: + return self._forward(*inputs, **kwargs) + + def _create_cuda_graph(self, *inputs, **kwargs): + # warmup to create the workspace and cublas handle + cuda_stream = torch.cuda.Stream() + cuda_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(cuda_stream): + for i in range(3): + ret = self._forward(*inputs, **kwargs) + torch.cuda.current_stream().wait_stream(cuda_stream) + + # create cuda_graph and assign static_inputs and static_outputs + self._all_cuda_graph = get_accelerator().create_graph() + self.static_inputs = inputs + self.static_kwargs = kwargs + + with get_accelerator().capture_to_graph(self._all_cuda_graph): + self.static_output = self._forward(*self.static_inputs, **self.static_kwargs) + + self.all_cuda_graph_created = True + + def _forward(self, sample, timestamp, encoder_hidden_states, return_dict=True): + return self.vae(sample, timestamp, encoder_hidden_states, return_dict) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/features/__init__.py b/lib/python3.12/site-packages/deepspeed/model_implementations/features/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/features/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/features/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/features/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3dbc7f31bc486d89dd73b3f370ed49ceff2bde89 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/features/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e608a9b815fef23dff6fcf6ec796e61adedeb244 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/features/cuda_graph.py b/lib/python3.12/site-packages/deepspeed/model_implementations/features/cuda_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..fc3a2951665426b5c98112ce307e395993644f8b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/features/cuda_graph.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod + + +class CUDAGraph(ABC): + + def __init__(self, enable_cuda_graph=False): + super().__init__() + self.enable_cuda_graph = enable_cuda_graph + + @abstractmethod + def _create_cuda_graph(self): + """ + Create CUDA graph(s) + """ + raise NotImplementedError + + @abstractmethod + def _graph_replay(self): + """ + Replay CUDA graph(s) + """ + raise NotImplementedError diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__init__.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d750801da09422a5a606ee921c62fb3e3d1ab02c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82cd77b863f2c86a52e9d43c42af4bb39a89ad9f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95bcca42f3ee654b44350eb3bb62fec8d4051e87 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a870706cf385e4e6d4c926e21f66a0ad3b31f9e4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20c1863b22e95452adb2712ac1fef75e4bec925c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d010d142adcedd0e1f7532367b39933e776e5ca Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed03b308a795dbf9596d74d5873ded4cf4d1fc0b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a0a560b4654369306144fbc1776d136586e9de2 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b989fcad991dc528f8fad16e78bc7babb580f781 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dc2a46db599db2c7c39a97a3a661c75932fc3e0 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..848a5b48dcf1336a4c80c20a9e21388cdff9e514 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.accelerator import get_accelerator +from ..features.cuda_graph import CUDAGraph + + +class DSClipEncoder(CUDAGraph, torch.nn.Module): + + def __init__(self, enc, enable_cuda_graph=False): + super().__init__(enable_cuda_graph=enable_cuda_graph) + enc.text_model._build_causal_attention_mask = self._build_causal_attention_mask + self.enc = enc + self.device = self.enc.device + self.dtype = self.enc.dtype + self.cuda_graph_created = [False, False] + self.static_inputs = [None, None] + self.static_kwargs = [None, None] + self.static_output = [None, None] + self._cuda_graphs = [None, None] + self.iter = 0 + self.config = self.enc.config + + def _build_causal_attention_mask(self, bsz, seq_len, dtype): + mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype, device=get_accelerator().current_device_name()) + mask.fill_(torch.tensor(torch.finfo(dtype).min)) + mask.triu_(1) + mask = mask.unsqueeze(1) + return mask + + def _graph_replay(self, *inputs, **kwargs): + for i in range(len(inputs)): + if torch.is_tensor(inputs[i]): + self.static_inputs[self.iter][i].copy_(inputs[i]) + for k in kwargs: + if torch.is_tensor(kwargs[k]): + self.static_kwargs[self.iter][k].copy_(kwargs[k]) + get_accelerator().replay_graph(self._cuda_graphs[self.iter]) + return self.static_output[self.iter] + + def forward(self, *inputs, **kwargs): + if self.enable_cuda_graph: + if self.cuda_graph_created[self.iter]: + outputs = self._graph_replay(*inputs, **kwargs) + else: + self._create_cuda_graph(*inputs, **kwargs) + outputs = self._graph_replay(*inputs, **kwargs) + self.iter = (self.iter + 1) % 2 + return outputs + else: + return self.enc(*inputs, **kwargs) + + def _create_cuda_graph(self, *inputs, **kwargs): + # warmup to create the workspace and cublas handle + cuda_stream = torch.cuda.Stream() + cuda_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(cuda_stream): + for i in range(3): + ret = self._forward(*inputs, **kwargs) + torch.cuda.current_stream().wait_stream(cuda_stream) + + # create cuda_graph and assign static_inputs and static_outputs + self._cuda_graphs[self.iter] = get_accelerator().create_graph() + self.static_inputs[self.iter] = inputs + self.static_kwargs[self.iter] = kwargs + + with get_accelerator().capture_to_graph(self._cuda_graphs[self.iter]): + self.static_output[self.iter] = self._forward(*self.static_inputs[self.iter], + **self.static_kwargs[self.iter]) + + self.cuda_graph_created[self.iter] = True + + def _forward(self, *inputs, **kwargs): + return self.enc(*inputs, **kwargs) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_base.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_base.py new file mode 100644 index 0000000000000000000000000000000000000000..b4b113904997fd124aecdfe51529aab6c9e1b25c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_base.py @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch.nn as nn + + +class DeepSpeedTransformerBase(nn.module): + + def __init__(self): + pass + + # this would be the new clean base class that will replace DeepSpeedTransformerInference. + # we currently don't know how this will look like but keeping it here as a placeholder. diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_bert.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..13075553ec8b6d818744cc5404ef4db31283ac8e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_bert.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedBERTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed BERT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py new file mode 100644 index 0000000000000000000000000000000000000000..c48c7ed58ae5e2746c513a1acea509bbfd129784 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedBloomInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed Bloom Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..3c3baed1f6186a89ed2b852f27ad56a5ab00d9f6 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedGPTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed GPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py new file mode 100644 index 0000000000000000000000000000000000000000..325bfb4f7e181d8b718748fd5ebf2f4f1ca3af64 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py @@ -0,0 +1,58 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedLlama2Inference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed OPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) + + def forward(self, *args, **kwargs): + + input = args[0] + input_mask = None + get_present = True + + self.allocate_workspace(input.size()) + + # We set the prev key/value to None when there is a prompt + if input.shape[1] > 1: + self.layer_past = None + layer_past = self.layer_past + + input_type = input.dtype + + if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \ + and input.dtype == torch.float: + target_dtype = torch.half if self.dtype == torch.int8 else self.dtype + input = input.to(target_dtype) + + with torch.no_grad(): + attention_output, key, value, context_outputtn_ctx, inp_norm = \ + self.attention(input, + input_mask, + None, + layer_past, + get_present, + None, None, None, + self.norm_w, + self.norm_b, + None) + self.layer_past = (key, value) + output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob) + + output = output.to(input_type) + return output diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..055ed6d27d7db4cb39cf868d5c189ef517f4a08a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedMegatronGPTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed Megatron GPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_opt.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_opt.py new file mode 100644 index 0000000000000000000000000000000000000000..7bc5524d71c7155fad0b40eafeca490705f3e324 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_opt.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference + + +class DeepSpeedOPTInference(DeepSpeedTransformerInference): + """Initialize the DeepSpeed OPT Transformer Layer. + """ + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..7e3c81b714c06b84fff1d18daf28ecccbc1db4f6 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py @@ -0,0 +1,191 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn as nn +from deepspeed import comm as dist +from deepspeed.ops.transformer.inference.op_binding.layer_norm import LayerNormOp +from deepspeed.utils.logging import log_dist + +from deepspeed.ops.transformer.inference.ds_mlp import DeepSpeedMLP +from deepspeed.ops.transformer.inference.ds_attention import DeepSpeedSelfAttention, BloomSelfAttention +from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp +from deepspeed.accelerator import get_accelerator +import deepspeed +if deepspeed.HAS_TRITON and get_accelerator().is_triton_supported(): + from deepspeed.ops.transformer.inference.triton.mlp import TritonMLP + from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention + + +class DeepSpeedTransformerInference(nn.Module): + """Initialize the DeepSpeed Transformer Layer. + Arguments: + layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers, + layer_id will be 0,1,2...23 when each layer object is instantiated + config: An object of DeepSpeedInferenceConfig + mp_group: Model parallelism group initialized on the modeling side. + quantize_scales: This argument groups all the layers' scales used for quantization + quantize_groups: Number of groups used for quantizing the model + merge_count: Shows the number of model-parallel checkpoints merged before running inference. + We use this argument to control the quantization scale for the model parameters if a bigger + quantize-grouping than 1 is used. + mlp_extra_grouping: This flag is used to show a 2x higher number of groups used for the MLP part + of a Transformer layer. We use this feature for quantization to reduce the convergence impact + for specific downstream tasks. + """ + layer_id = 0 + workspace = None + + def __init__(self, + config, + mp_group=None, + quantize_scales=None, + quantize_groups=1, + merge_count=1, + mlp_extra_grouping=False): + super(DeepSpeedTransformerInference, self).__init__() + + self.config = config + self.config.layer_id = DeepSpeedTransformerInference.layer_id + DeepSpeedTransformerInference.layer_id += 1 + + data_type = torch.half if self.config.dtype == torch.int8 else self.config.dtype + + if DeepSpeedTransformerInference.layer_id == 1: + log_dist(f"DeepSpeed-Inference config: {self.config.__dict__}", [0]) + if deepspeed.HAS_TRITON and self.config.use_triton: + log_dist(f"Injecting Triton kernels ...", [0]) + + if self.config.bigscience_bloom: + self.attention = BloomSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, merge_count) + assert not self.config.use_triton + else: + if deepspeed.HAS_TRITON and self.config.use_triton: + self.attention = TritonSelfAttention(self.config) + else: + self.attention = DeepSpeedSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, + merge_count) + + if deepspeed.HAS_TRITON and self.config.use_triton: + self.mlp = TritonMLP(self.config) + else: + self.mlp = DeepSpeedMLP(self.config, mp_group, quantize_scales, quantize_groups, merge_count, + mlp_extra_grouping) + + device = get_accelerator().current_device_name() # if config.bigscience_bloom else 'cpu' + if self.config.set_empty_params: + self.norm_w = None + self.norm_b = None + else: + self.norm_w = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device), + requires_grad=False) + self.norm_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device), + requires_grad=False) + self.layer_past = None + self.layer_norm = LayerNormOp() + if DeepSpeedTransformerInference.workspace is None: + DeepSpeedTransformerInference.workspace = WorkspaceOp(self.config) + self._should_allocate_workspace = True + + def allocate_workspace(self, size): + # Allocate memory only on first layer forward + if self.config.layer_id == 0 and self._should_allocate_workspace: + DeepSpeedTransformerInference.workspace.allocate_workspace( + self.config.hidden_size, self.config.heads, size[1], size[0], DeepSpeedTransformerInference.layer_id, + self.config.mp_size, self.config.bigscience_bloom, + dist.get_rank() if dist.is_initialized() else 0, self.config.max_out_tokens, + self.config.min_out_tokens) + self._should_allocate_workspace = False + + @classmethod + def reset_cache(cls): + if cls.workspace is not None: + cls.workspace.reset_cache() + + def forward( + self, + input=None, + input_mask=None, + attention_mask=None, + attn_mask=None, + head_mask=None, + layer_past=None, + get_key_value=False, + get_present=False, + encoder_output=None, + enc_dec_attn_mask=None, + x=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + use_cache=False, + alibi=None, + output_attentions=False, + # TODO(arashb): 'layer_head_mask' and 'past_key_value' are only added to satisfy the OPT models API. + # This needs to be redesigned later! + layer_head_mask=None, + past_key_value=None, + **kwargs): + + if x is not None: + input = x + if "hidden_states" in kwargs: + input = kwargs["hidden_states"] + + input_mask = (input_mask if attn_mask is None else attn_mask) if attention_mask is None else attention_mask + + self.allocate_workspace(input.size()) + + get_present = (get_present or get_key_value or use_cache) + input_mask = input_mask if attention_mask is None else attention_mask + + # We set the prev key/value to None when there is a prompt + if input.shape[1] > 1: + self.layer_past = None + layer_past = layer_past if layer_past is not None else self.layer_past + head_mask = layer_head_mask if layer_head_mask is not None else head_mask + + attn_mask = None + if isinstance(input, tuple): + attn_mask = input[1] + input = input[0] + input_type = input.dtype + + if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \ + and input.dtype == torch.float: + target_dtype = torch.half if self.config.dtype == torch.int8 else self.config.dtype + input = input.to(target_dtype) + + with torch.no_grad(): + attention_output, key, value, context_outputtn_ctx, inp_norm = \ + self.attention(input, + input_mask, + head_mask, + layer_past, + get_present, + encoder_hidden_states, + encoder_attention_mask, + output_attentions, + self.norm_w, + self.norm_b, + alibi, + **kwargs) + + presents = (key, value) + self.layer_past = presents if layer_past is None else None + output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob) + + if not self.config.pre_layer_norm: + output = self.layer_norm(output, self.norm_w, self.norm_b, self.config.epsilon) + + output = output.to(input_type) + if get_present: + output = (output, presents) + + if self.config.return_single_tuple: + return (output, ) + elif self.config.return_tuple: + return output if type(output) is tuple else (output, attn_mask) + else: + return output diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__init__.py b/lib/python3.12/site-packages/deepspeed/module_inject/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fc2f979a04b7adb8aaa638d8e2db3d6b2946d98 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .replace_module import replace_transformer_layer, revert_transformer_layer, ReplaceWithTensorSlicing, GroupQuantizer, generic_injection +from .module_quantize import quantize_transformer_layer +from .replace_policy import HFBertLayerPolicy +from .layers import LinearAllreduce, LinearLayer, EmbeddingLayer, Normalize, set_autotp_mode +from .policy import DSPolicy diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d2571f93b0f6b4a9795164b660139f87d2b5fb6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aea80fe5779b9cb36c90c9daccd4e6dcab962c87 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00bfec7621ce2f53c5d8ef5a6439a023632a5128 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc0c36827d68c59125f1b6ab75e065122baa85e4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..882d646340cbfa06ef25331bbac64797db3e9d2b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47e94b4d64bca162a7238a4e29eaab8c5b9b2abc Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b57b82840c1c9d93be6311e8798a63fdac1fd90 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9225d6f86f234883c15724e9d37d7e64cd66504f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8864edb87e5afae0e4d031f6aaf9d3dfb895f935 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f3a841d2c2c749fb9c9db8640f2921982440fc4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5bd8164cb99a9c0eca9b1e738e0c4934e65a273c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20574395d69f103cfe13b0989696a4700abb3a41 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6fd1aa3b758cbba0f0b5d7263b38c1223b284bb1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/auto_tp.py b/lib/python3.12/site-packages/deepspeed/module_inject/auto_tp.py new file mode 100644 index 0000000000000000000000000000000000000000..82cd9042071e7de85cffdcbdba8dad6aa94568bc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/auto_tp.py @@ -0,0 +1,509 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Automatic Tensor Parallelism +import re + +from torch import nn +from .replace_policy import replace_policies +from typing import Optional +import torch +from deepspeed import comm as dist +from .layers import * +from deepspeed.accelerator import get_accelerator +from .fusedqkv_utils import require_tp_fused_qkvw +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list +from deepspeed.utils import groups +from deepspeed.module_inject.layers import is_autotp_training_mode + + +def move(tensor, device, copy=True): + if tensor.is_meta: + return torch.empty_like(tensor, device=device) + else: + # Using new tensors help in freeing memory (after split for example) was done before by calling clone(). + # Using copy=True instead of clone() will help in case of cpu --> cpu. + # Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced. + return tensor.to(device, copy=copy) + + +class ReplaceWithTensorSlicing: + + def __init__(self, mp_group=None, mp_size=1, out_dim=1, in_dim=0): + if mp_group is not None: + self.gpu_index = dist.get_rank(group=mp_group) + else: + self.gpu_index = 0 + self.out_dim = out_dim + self.in_dim = in_dim + self.mp_size = mp_size + + def merge_assert(self, dim1, dim2): + assert dim1 > dim2, \ + 'Merging tensors is not allowed here! Please use deepspeed load_checkpoint\ + for merging your checkpoints before replacing the transformer layer with\ + inference-kernels' + + def strided_copy(self, + dst: Optional[torch.Tensor], + src: Optional[torch.Tensor], + num_splits: int, + int8: bool = False, + allocate_tensor: bool = False): + if src is None: + return src + src_shape = src.shape + dst_shape = dst.shape + + outer_dim = 0 if int8 else -1 + + if allocate_tensor: + dst = torch.empty_like(dst) + + src_split = torch.split(src.data, src.shape[outer_dim] // num_splits, dim=outer_dim) + if (len(src_shape) == 2 and len(dst_shape) == 2): + if src_shape[outer_dim] == dst_shape[self.out_dim]: + try: + dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape) + except: + print(dst.shape, src.shape) + exit() + dst = torch.nn.parameter.Parameter(dst, requires_grad=False) + if hasattr(src, 'scale'): + dst.scale = src.scale + return dst + self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim]) + qkv_size = dst_shape[self.out_dim] // num_splits + qkv_split = [torch.split(src_s, qkv_size, dim=outer_dim) for src_s in src_split] + weight_split = [ + torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=outer_dim) for i in range(len(qkv_split[0])) + ] + dst = dst.reshape(-1).data.copy_(weight_split[self.gpu_index].contiguous().reshape(-1)).reshape( + weight_split[self.gpu_index].shape) + else: + if src_shape[0] == dst_shape[0]: + return torch.nn.parameter.Parameter(src) + qkv_size = dst_shape[0] // num_splits + qkv_split = [torch.split(src_s, qkv_size, dim=0) for src_s in src_split] + bias_split = [torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=0) for i in range(len(qkv_split[0]))] + dst.data.copy_(bias_split[self.gpu_index].contiguous()) + + dst = torch.nn.parameter.Parameter(dst, requires_grad=False) + if hasattr(src, 'scale'): + dst.scale = src.scale + return dst + + def copy(self, dst, src, int8=False, allocate_tensor=False): + if src is None: + return src + assert not dst.data.is_meta # the torch.Tensor.copy_ method used below will silently fail on meta tensors + if allocate_tensor: + dst = torch.empty_like(dst) + outer_dim = 0 if int8 else 1 + inner_dim = 1 if int8 else 0 + src_shape = src.shape + dst_shape = dst.shape + if (len(src_shape) == 2 and len(dst_shape) == 2): + + if src_shape[inner_dim] == dst_shape[self.in_dim] and src_shape[outer_dim] == dst_shape[self.out_dim]: + dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape) + else: + if src_shape[inner_dim] != dst_shape[self.in_dim]: + self.merge_assert(src_shape[inner_dim], dst_shape[self.in_dim]) + dst.data.copy_(src[:, self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim]] if inner_dim == 1 else \ + src[self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim], :]) + else: + self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim]) + dst.data.copy_(src[:, self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim]] if outer_dim == 1 else \ + src[self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim], :]) + else: + if src_shape[0] == dst_shape[0]: + dst = src if src.dtype == dst.dtype else dst.data.copy_(src) + else: + dst.data.copy_(src[self.gpu_index * dst_shape[-1]:(self.gpu_index + 1) * dst_shape[-1]]) + dst = torch.nn.parameter.Parameter(dst, requires_grad=False) + if hasattr(src, 'scale'): + dst.scale = src.scale + return dst + + +class Loading(): + + def is_load_module(module): + load_layers = [nn.Linear, nn.Embedding, nn.LayerNorm] + load_layer_names = [ + "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "FalconLinear", + "MistralRMSNorm", "T5LayerNorm", "MixtralRMSNorm", "Phi3RotaryEmbedding", "Phi3SuScaledRotaryEmbedding", + "Phi3RMSNorm", "YuanRMSNorm", "YuanRotaryEmbedding", "Phi3LongRoPEScaledRotaryEmbedding", "Qwen2RMSNorm", + "Qwen3RMSNorm", "Qwen3MoeRMSNorm", "DeepseekV2RMSNorm", "DeepseekV3RMSNorm", + "DeepseekV2YarnRotaryEmbedding", "DeepseekV3YarnRotaryEmbedding", "MoEGate" + ] + return module.__class__ in load_layers or module._get_name() in load_layer_names + + def load_buffer(module, state_dict, prefix): + for name in module._buffers.keys(): + if module._buffers[name].data.is_meta: + module._buffers[name] = torch.nn.parameter.Parameter( + data=torch.empty_like(module._buffers[name].data, device="cpu"), + requires_grad=module._buffers[name].data.requires_grad) + if prefix + name in state_dict.keys(): + module._buffers[name].data.copy_(state_dict[prefix + name]) + + def load(module, state_dict, prefix, mp_group=None): + mp_replace = ReplaceWithTensorSlicing(mp_group=mp_group) + if hasattr(module, 'weight'): + if module.weight.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.weight = torch.nn.parameter.Parameter(data=torch.empty_like(module.weight.data, device="cpu"), + requires_grad=module.weight.data.requires_grad) + if 'query_key_value' in prefix: + module.weight = mp_replace.strided_copy(module.weight.data, + state_dict[prefix + 'weight'], + num_splits=3) + else: + module.weight = mp_replace.copy(module.weight.data, state_dict[prefix + 'weight']) + else: + if hasattr(module, 'norm') and hasattr(module.norm, 'weight'): + if module.norm.weight.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.norm.weight = torch.nn.parameter.Parameter( + data=torch.empty_like(module.norm.weight.data, device="cpu"), + requires_grad=module.norm.weight.data.requires_grad) + module.norm.weight = mp_replace.copy(module.norm.weight.data, state_dict[prefix + 'weight']) + + if prefix + 'bias' in state_dict.keys(): + if hasattr(module, 'bias'): + if module.bias.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"), + requires_grad=module.bias.data.requires_grad) + module.bias = mp_replace.copy(module.bias, state_dict[prefix + 'bias']) + else: + if hasattr(module, 'norm') and hasattr(module.norm, 'bias'): + if module.norm.bias.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.norm.bias = torch.nn.parameter.Parameter( + data=torch.empty_like(module.norm.bias.data, device="cpu"), + requires_grad=module.norm.bias.data.requires_grad) + module.norm.bias = mp_replace.copy(module.norm.bias, state_dict[prefix + 'bias']) + + +class AutoTP(): + + def __init__(self, + module, + all_reduce_linears, + prefix, + state_dict, + linear_layer_setting, + orig_layer_impl, + keep_module_on_host=False): + self.module = module + self.all_reduce_linears = all_reduce_linears + self.prefix = prefix + self.state_dict = state_dict + + self.mp_size = None + self.mp_group = None + self.linear_layer_setting = linear_layer_setting + self.orig_layer_impl = orig_layer_impl + self.linear_policies = None + self.conv_linear_layer = False + TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host) + + def in_module_list(module, module_list): + for item in module_list: + if type(item).__name__ == type(module).__name__: + return True + return False + + def get_module_list(model): + mlist = [] + for child in model.children(): + if isinstance(child, nn.ModuleList): + for module in child.children(): + if not mlist: + mlist = [module] + elif not AutoTP.in_module_list(module, mlist): + mlist = mlist + [module] + else: + mlist = mlist + AutoTP.get_module_list(child) + return mlist + + def supported(model): + unsupported = ['deberta', 'flaubert', 'fsmt', 'gpt2', 'led', 'longformer', 'xlm', 'xlnet'] + model = str(model) + key = re.search(r": (.*?)Model", model) + if key is None: + key = re.search(r": (.*?)Stack", model) + if key is None: + key = re.match(r"(.*?)Model", model) + assert key is not None, "Not able to determine model policy automatically. Please provide policy." + if key.group(1).lower() in unsupported: + return False + return True + + def get_layers(parent, module): + layer_list = [] + for key, submodule in module._modules.items(): + if isinstance(submodule, nn.Linear): + layer_list = layer_list + [parent + "." + key] + elif isinstance(submodule, nn.LayerNorm) or key == 'LayerNorm' or key == 'layer_norm': + layer_list = layer_list + ["ln"] + else: + layer_list = layer_list + AutoTP.get_layers(key, submodule) + return layer_list + + def update_policy_list(policy_list, new_module, new_gems): + if len(policy_list): + for i, policy in enumerate(policy_list): + # if module already exists in policy, combine gems and remove duplicates + if policy[0] == type(new_module): + new_gems = set(new_gems + policy[1]) + policy_list[i] = tuple([type(new_module), new_gems]) + return policy_list + policy_list.append(tuple([type(new_module), new_gems])) + return policy_list + + def kernel_supported(module_list): + policy = [] + for plcy in replace_policies: + # instantiate a throw-away policy in order to populate the _orig_layer_class + _ = plcy(None) + if isinstance(plcy._orig_layer_class, list): + for orig_layer_class in plcy._orig_layer_class: + policy.append(orig_layer_class) + elif plcy._orig_layer_class is not None: + policy.append(plcy._orig_layer_class) + for child in module_list: + if child.__class__ in policy: + return True + return False + + def tp_parser(model): + policy_list = [] + module_list = [] + layer_list = [] + gem_list = [] + + module_list = AutoTP.get_module_list(model) + assert AutoTP.supported(model), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \ + if AutoTP.kernel_supported(module_list) else "AutoTP not supported for model. Please provide policy." + norm_layer_name_list = ['LayerNorm', 'layer_norm', 'ln_1', 'ln_2'] + #ln_1 , ln_2 for Qwen + for module in module_list: + for key, submodule in module._modules.items(): + if isinstance(submodule, nn.Linear): + layer_list = layer_list + ["." + key] + elif isinstance(submodule, nn.LayerNorm) or key in norm_layer_name_list: + layer_list = layer_list + ["ln"] + else: + layer_list = layer_list + AutoTP.get_layers(key, submodule) + for i, layer in enumerate(layer_list): + if layer == 'ln': + if layer_list[i - 1] != 'ln': + gem_list = gem_list + [layer_list[i - 1]] + elif 'out_proj' in layer: + gem_list = gem_list + [layer] + elif 'o_proj' in layer: + gem_list = gem_list + [layer] + elif 'down_proj' in layer: + gem_list = gem_list + [layer] + elif 'attention.dense' in layer and 'GPTNeoX' in str(model): + gem_list = gem_list + [layer] + elif 'self_attention.dense' in layer and 'falcon' in str( + type(module)): # this is a hack to get the right linear layer for this model! + gem_list = gem_list + [layer] + # Mixtral-7x8b used w2*act(w1*w3) linear. need to replace w2 to linearallreduce. + elif 'w2' in layer and 'Mixtral' in str(type(module)): + gem_list = gem_list + [layer] + elif 'self_attn.dense' in layer and 'Phi' in str(type(module)): + gem_list = gem_list + [layer] + elif 'self_attention.dense' in layer and 'ChatGLM' in str(model): + gem_list = gem_list + [layer] + elif 'dense_4h_to_h' in layer and 'ChatGLM' in str(model): + gem_list = gem_list + [layer] + + layer_list = [] + if gem_list != []: + gem_list = list(set(gem_list)) + policy_list = AutoTP.update_policy_list(policy_list, module, gem_list) + gem_list = [] + assert len(policy_list), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \ + if AutoTP.kernel_supported(module_list) else "Not able to determine model policy automatically. Please provide policy." + return policy_list + + def set_tensor_parallel_config(self, mp_size, mp_group): + + if is_autotp_training_mode(): + self.mp_group = groups.get_tensor_model_parallel_group() + self.mp_size = groups.get_tensor_model_parallel_world_size() + return + + self.mp_size = mp_size + self.mp_group = mp_group + + def _replace(self, child, name, conv_linear_layer): + # This function should clearly define the routing rules for specific layers + # and avoid any complex shard-related logic. + if getattr(child, "replaced", False) == True: + return + + weight_shape = child.weight.shape + mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group) + # For TP layer skip, e.g., MoE gate, deepseek low rank layer skip + if "mlp.gate" == name or "q_a_proj" in name or "kv_a_proj_with_mqa" in name or name == "block_sparse_moe.gate" or ( + ('mlp.shared_expert_gate' == name or 'mlp.gate' == name) and 'qwen2_moe' in str(type(self.module))): + return child + # For Yuan model + if 'Yuan' in str(self.module): + if 'v_proj' in name: + return Yuan_LinearLayer(child, self.mp_group) + + elif 'o_proj' in name: + return Yuan_LinearAllreduce(child, self.mp_group) + + # For MLP including chunk layer. + if 'gate_up_proj' in name or ('dense_h_to_4h' in name and 'GLM' in str(self.module)): + return GateUpPack_LinearLayer(child, self.mp_group) + # For Arctic model, bypass to all_reduce replacement for w2 weights + arctic_w2_all_reduce_linear = False + if 'Arctic' in str(self.module) and 'w2' in name: + arctic_w2_all_reduce_linear = True + # For MoE MLP model, e.g., deepseek and jamba + down_proj = False + if 'down_proj' in name: + down_proj = True + if name in self.all_reduce_linears or arctic_w2_all_reduce_linear or down_proj: + + setattr(child, "replaced", True) + if self.conv_linear_layer: + return Conv_LinearALlreduce(child, self.mp_group, name=name) + elif name == "lm_head" or name == 'embed_out': + return LmHeadLinearAllreduce(child, self.mp_group) + + return LinearAllreduce(child, self.mp_group, name=name) + else: + + setattr(child, "replaced", True) + if self.conv_linear_layer: + conv_LinearLayer(child, self.mp_group) + elif require_tp_fused_qkvw(name, self.mp_size): + #Check and handle fused qkv for TP + return fused_LinearLayer(child, self.mp_group, fused_module=self.module) + + return LinearLayer(child, self.mp_group, name=name) + + def _slice_embedding(self, child, name, conv_linear_layer): + if getattr(child, "replaced", False) == True: + return + mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group) + + if hasattr(child.weight, 'ds_tensor'): + data = child.weight.ds_tensor.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size), dim=1) + else: + data = child.weight.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size, name), dim=1) + data = data[mp_replace.gpu_index].to(get_accelerator().current_device_name()) + data = torch.nn.parameter.Parameter(data, requires_grad=False) + + new_embedding = nn.Embedding(child.weight.shape[0], get_shard_size(child.weight.shape[1], self.mp_size, name)) + new_embedding.weight.data.copy_(data) + setattr(child, "replaced", True) + return new_embedding + + def update_mp_params(self, child): + if getattr(child, "replaced", False) == True: + return + param_list = [ + "n_heads", "inner_dim", "num_heads", "num_kv", "num_attention_heads", "num_attn_heads", "all_head_size", + "embed_dim", "hidden_size", "num_key_value_heads", "num_kv_heads", "kv_n_heads", "d_model", + "num_attention_heads_per_partition", "num_multi_query_groups_per_partition", "hidden_size_per_partition" + ] + for param in param_list: + if "Yuan" in str(child) and 'embed_dim' in param_list: + param_list.remove('embed_dim') + if hasattr(child, param): + param_val = getattr(child, param) + setattr(child, param, get_shard_size(param_val, self.mp_size)) + setattr(child, "replaced", True) + + def update_linear_policies(self): + self.conv_linear_layer = False + if self.linear_layer_setting is not None: + self.linear_policies = {self.linear_layer_setting[0]: self._replace} + if len(self.linear_layer_setting) == 2: + self.linear_policies.update({self.linear_layer_setting[1]: self._slice_embedding}) + else: + import transformers + if self.orig_layer_impl is transformers.models.gpt2.modeling_gpt2.GPT2Block: + try: + self.conv_linear_layer = True + self.linear_policies = {transformers.pytorch_utils.Conv1D: self._replace} + except ImportError: + self.linear_policies = {nn.Linear: self._replace} + else: + self.linear_policies = {nn.Linear: self._replace, nn.Embedding: self._slice_embedding} + + def _replace_module(self, r_module, prev_name='', prev_class_name=''): + for name, child in r_module.named_children(): + if prev_class_name == "": + class_name = prev_name + elif prev_name == "": + class_name = prev_class_name + else: + class_name = prev_class_name + '.' + prev_name + checking_key = self.prefix + '.' + class_name + '.' + name + '.' if class_name != "" else self.prefix + '.' + name + '.' + if Loading.is_load_module(child) and self.state_dict is not None: + if any(checking_key in item for item in self.state_dict): + Loading.load(child, self.state_dict, checking_key, self.mp_group) + else: + continue + if len(child._buffers) != 0 and self.state_dict is not None: + Loading.load_buffer(child, self.state_dict, checking_key) + if child.__class__ in self.linear_policies: + setattr(r_module, name, self.linear_policies[child.__class__](child, prev_name + '.' + name, + self.conv_linear_layer)) + elif any(isinstance(child, lp) for lp in self.linear_policies): + # Added for falcon model support + # Note: isinstance will account for class inheritance, child.__class__ does not + key = None + for lp in self.linear_policies: + if isinstance(child, lp): + key = lp + break + assert key is not None + setattr(r_module, name, self.linear_policies[key](child, prev_name + '.' + name, + self.conv_linear_layer)) + else: + self.update_mp_params(child) + self._replace_module(child, name, class_name) + return r_module + + def get_model_num_kv_heads(self, config): + num_kv_heads = None + # multi_query_group_num is for chatglm2 & chatglm3 + kv_head_names = [ + 'multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads', + 'attention_heads' + ] + for name in kv_head_names: + if hasattr(config, name): + num_kv_heads = getattr(config, name) + if num_kv_heads is not None: + break + return num_kv_heads + + def _replace_last_linear_module(self, r_module): + if hasattr(r_module, "lm_head"): + name = "lm_head" + child = r_module.lm_head + elif hasattr(r_module, "embed_out"): + name = "embed_out" + child = r_module.embed_out + else: + return r_module + if child.__class__ in self.linear_policies: + setattr(r_module, name, self.linear_policies[child.__class__](child, name, self.conv_linear_layer)) + return r_module diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/auto_tp_model_utils.py b/lib/python3.12/site-packages/deepspeed/module_inject/auto_tp_model_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a71b1a54d6f6031c18899b1a5294dd5dd963e92d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/auto_tp_model_utils.py @@ -0,0 +1,104 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed import comm as dist +import torch +from typing import Optional +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list + + +def build_bloom_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor: + """ + Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it + relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value + `softmax(l+a) = softmax(l)`. Based on + https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 + TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly. + + Args: + Returns tensor shaped (batch_size * num_heads, 1, max_seq_len) + attention_mask (`torch.Tensor`): + Token-wise attention mask, this should be of shape (batch_size, max_seq_len). + num_heads (`int`, *required*): + number of heads + dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`): + dtype of the output tensor + """ + import math + batch_size, seq_length = attention_mask.shape + closest_power_of_2 = 2**math.floor(math.log2(num_heads)) + base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))), + device=attention_mask.device, + dtype=torch.float32) + powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32) + slopes = torch.pow(base, powers) + + if closest_power_of_2 != num_heads: + extra_base = torch.tensor(2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))), + device=attention_mask.device, + dtype=torch.float32) + num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2) + extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32) + slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0) + + # Note: alibi will added to the attention bias that will be applied to the query, key product of attention + # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length) + # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length) + # => the query_length dimension will then be broadcasted correctly + # This is more or less identical to T5's relative position bias: + # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527 + arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :] + alibi = slopes[..., None] * arange_tensor + if dist.is_initialized(): + num_heads_per_rank = get_shard_size(num_heads, dist.get_world_size()) + offset = sum(get_shard_size_list(num_heads, dist.get_world_size())[0:dist.get_rank()]) + alibi = alibi.view(batch_size, num_heads, 1, seq_length) + alibi = alibi[:, offset:num_heads_per_rank + offset, :, :] + return alibi.reshape(batch_size * num_heads_per_rank, 1, seq_length).to(dtype) + else: + return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype) + + +def get_alibi_mask(self, tensor, seq_length_with_past): + mask = self.get_alibi_mask_orig(tensor, seq_length_with_past) + if not self.training and dist.is_initialized(): + num_heads_per_rank = get_shard_size(self.n_head, dist.get_world_size()) + offset = sum(get_shard_size_list(self.n_head, dist.get_world_size())[0:dist.get_rank()]) + mask = mask[offset:num_heads_per_rank + offset, :seq_length_with_past, :seq_length_with_past] + + return mask + + +def build_mpt_atten_bias_tensor(self, + device, + dtype, + attention_mask: Optional[torch.ByteTensor] = None, + prefix_mask: Optional[torch.ByteTensor] = None, + sequence_id: Optional[torch.LongTensor] = None): + (attn_bias, attention_mask) = self._attn_bias_orig(device, + dtype, + attention_mask=attention_mask, + prefix_mask=prefix_mask, + sequence_id=sequence_id) + if dist.is_initialized(): + num_heads_per_rank = get_shard_size(self.config.n_heads, dist.get_world_size()) + offset = sum(get_shard_size_list(self.config.n_heads, dist.get_world_size())[0:dist.get_rank()]) + attn_bias = attn_bias[:, offset:num_heads_per_rank + offset, :, :] + return attn_bias, attention_mask + + +def build_mpt_alibi_tensor(self, num_heads, sequence_length, alibi_bias_max=8, device=None) -> torch.Tensor: + r""" + Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it + relies on a translation invariance of softmax for quick implementation. This implementation has been copied from + the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi: + https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292 + """ + alibi = self.build_mpt_alibi_tensor_orig(num_heads, sequence_length, alibi_bias_max, device) + if dist.is_initialized(): + num_heads_per_rank = int(num_heads / dist.get_world_size()) + offset = dist.get_rank() * num_heads_per_rank + alibi = alibi[offset:num_heads_per_rank + offset, :, :] + return alibi diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__init__.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..993d14071659303ea6702b0d5182c6d04c931905 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .bert import DS_BERTContainer, HFBertLayerPolicy +from .bloom import DS_BloomContainer, BLOOMLayerPolicy, supported_models +from .distil_bert import DS_DistilBERTContainer, HFDistilBertLayerPolicy +from .gpt2 import DS_GPT2Container, HFGPT2LayerPolicy +from .gptj import DS_GPTJContainer, HFGPTJLayerPolicy +from .gptneo import DS_GPTNEOContainer, HFGPTNEOLayerPolicy +from .gptneox import DS_GPTNEOXContainer, GPTNEOXLayerPolicy +from .llama import DS_LLAMAContainer, LLAMALayerPolicy +from .llama2 import LLAMA2LayerPolicy, DS_LLAMA2Container +from .internlm import DS_InternLMContainer, InternLMLayerPolicy +from .megatron_gpt import DS_MegatronGPTContainer, MegatronLayerPolicy +from .megatron_gpt_moe import DS_MegatronGPTMoEContainer, MegatronMoELayerPolicy +from .opt import DS_OPTContainer, HFOPTLayerPolicy +from .clip import DS_CLIPContainer, HFCLIPLayerPolicy +from .unet import UNetPolicy +from .vae import VAEPolicy diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..546e8a293d48bc9fc36839b6dfe29c57fa7d6918 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f64134b3a84c8f546e1589b58e93f838595a579e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..48a85f2464e5017998a9a11e61053c3f4d305ee1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3039f351543dc40b65cc1618c776d793d7276aa1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9e6d2317cf1ead67b54341339c541e33f5e60c3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ef41f73164e71df25e38037028892f9120c4ee5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09dc3bb756942223ed7a2f021933eb23f46189a5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6dbd97461094307ffb5cd9a463587b362e0b45b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7d1404279ef0386d614242a53bb193f8ab80fff Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f8cc04443d52a71284cfb04a250685f638ee980 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6893d0b25255df36c4e18c630a27a4809506ab5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f8a285bb11af5dffab5f5606ae74ba41cf1a5e5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f90d55bcac302837ab4e9fce4e9dfa03c0ee88ab Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e472ca9e18f12932cf6ffe267db540a7299b4e07 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf72ad2e5727b07db8bba902b55d8a7f5abc6bc1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3930bc525240611c25df9f709b3c64c1c192a119 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edd1923f24c85517c87fe0a6ade05e0d6b9c3c56 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af0dec01ed0ccb1641431c5ad02a6e9108fcde6b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5335e908c51a9881d875f1c1ef4203810762c7c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/base.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..83e109167ffec669b6b1d77c2996010a70678f27 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/base.py @@ -0,0 +1,322 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Create a container object to save model-specific tensors using the policy file above. +from abc import ABC + +import torch + +import deepspeed +from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig +from deepspeed.accelerator import get_accelerator + +# If the intermediate size attribute is set DEFAULT_INTERMEDIATE_SIZE +# it is assumed the intermediate size is 4x the embedding dimension +DEFAULT_INTERMEDIATE_SIZE = -1 + + +class BaseConvolutionContainer(ABC): + # not implemented + def __init__(self): + pass + + +class BaseTransformerContainer(ABC): + + def __init__(self, policy, config, model_config, layer_id, child): + self.policy = policy + self.config = config + self.model_config = model_config + self.layer_id = layer_id + self.child = child + + self.megatron_v2 = self.policy.is_megatron_v2 + self.scale_attention = self.policy.scale_attention + self.ckpt_load_enabled = False + + # configuration for models. todo: can this be moved to a pydantic model config? + self.hidden_size = None + self.intermediate_size = None + self.num_attention_heads = None + self.mp_size = self.config.tensor_parallel.tp_size + self.pre_layer_norm = self.model_config.do_layer_norm_before if \ + hasattr(self.model_config, 'do_layer_norm_before') else self.policy.pre_attn_norm + self.dtype = self.config.dtype + self.attn_linear_layer = self.policy.linear_layer + self.mlp_linear_layer = self.policy.linear_layer + self.return_tuple = self.config.return_tuple + self.triangular_masking = True + self.local_attention = ((self.model_config.attention_layers[self.layer_id] == "local") if hasattr( + self.model_config, 'attention_layers') else False) + self.window_size = getattr(self.model_config, "window_size", 1) + self.mlp_act_func_type = self.policy.mlp_act_func_type + self.norm_type = self.policy.norm_type + self.training_mp_size = self.config.training_mp_size + self.bigscience_bloom = False + self.max_out_tokens = self.config.max_out_tokens + self.min_out_tokens = self.config.min_out_tokens + self.scale_attn_by_inverse_layer_idx = getattr(self.config, "scale_attn_by_inverse_layer_idx", False) + self.use_mup = self.policy.use_mup + self.return_single_tuple = False + self.rotary_dim = self.get_rotary_dim() + self.mlp_after_attn = (self.rotary_dim is None or self.rotary_dim < 0) + + # Attention tensors + self.qkvw = None + self.qkvb = None + self.dense_w = None + self.dense_b = None + # MLP tensors + self._h4h_w = None + self._h4h_b = None + self._4hh_w = None + self._4hh_b = None + # LayerNorm tensors + self.attn_nw = None + self.attn_nb = None + self.input_nw = None + self.input_nb = None + + self.mp_group = None + self.use_triton = False + + # Triton + self.use_triton = config.use_triton and deepspeed.HAS_TRITON + + def create_ds_model_config(self): + self.set_hidden_heads(*self.policy.get_hidden_heads()) + assert self.num_attention_heads % self.mp_size == 0,\ + "To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\ + "This is because the attention computation is partitioned evenly among the parallel GPUs." + + self.ds_model_config = DeepSpeedInferenceConfig( + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + heads=self.num_attention_heads, + layer_norm_eps=self.layernorm_epsilon, + dtype=self.dtype, + pre_layer_norm=self.pre_layer_norm, + norm_type=self.norm_type, + mp_size=self.mp_size, + return_tuple=self.return_tuple, + triangular_masking=self.triangular_masking, + local_attention=self.local_attention, + window_size=self.window_size, + rotary_dim=self.rotary_dim, + mlp_after_attn=self.mlp_after_attn, + mlp_act_func_type=self.mlp_act_func_type, + training_mp_size=self.training_mp_size, + bigscience_bloom=self.bigscience_bloom, + max_out_tokens=self.max_out_tokens, + min_out_tokens=self.min_out_tokens, + scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx, + use_mup=self.use_mup, + return_single_tuple=self.return_single_tuple, + set_empty_params=self.config.set_empty_params, + transposed_mode=self.config.transposed_mode, + use_triton=self.use_triton, + triton_autotune=self.config.triton_autotune) + + if self.use_triton and deepspeed.HAS_TRITON: + from .bert import DS_BERTContainer + if not isinstance(self, DS_BERTContainer): + raise NotImplementedError("Triton kernels are only for BERT-like models yet") + + if not self.config.triton_autotune: + from deepspeed.ops.transformer.inference.triton.matmul_ext import fp16_matmul + fp16_matmul.skip_autotune() + + return self.ds_model_config + + def check_meta_tensor_support(self): + if hasattr(self.qkvw, 'is_meta'): + if self.qkvw.is_meta: + assert self.ckpt_load_enabled, "Meta tensors are not supported for this model currently." + else: + raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") + + def initialize_tensors(self, enable_training=False): + # Set the tensors from policy (user module) to container (DS module) + self.set_attention(*self.policy.attention(enable_training=enable_training)) + self.set_mlp(*self.policy.mlp(enable_training=enable_training)) + self.set_layernorm(*self.policy.layernorm()) + #self.check_meta_tensor_support() + + def convert_to_required_dtype(self): + # Note: converting tensors to fp16 requires that we do it in-place using self.__dict__ and not make a list/dict copy + if self.dtype in [torch.half, torch.bfloat16]: + for k, v in self.__dict__.items(): + # The list comprehension is used for MoE tensor lists + if isinstance(v, list) and all((isinstance(tensor, torch.Tensor) \ + or isinstance(tensor, torch.nn.Parameter)) for tensor in v): + self.__dict__[k] = [moe_tensor.to(self.dtype) for moe_tensor in v] + + if isinstance(v, torch.Tensor) or isinstance(v, torch.nn.Parameter): + self.__dict__[k] = v.to(self.dtype) + + def get_rotary_dim(self): + if hasattr(self.model_config, 'rotary_dim'): + return self.model_config.rotary_dim + if hasattr(self.child, 'attention') and hasattr(self.child.attention, 'rotary_ndims'): + return self.child.attention.rotary_ndims + return -1 + + def set_moe(self, moe=False): + self.moe = moe + + def set_tensor_parallel_config(self, mp_size, mp_group): + self.mp_size = mp_size + self.mp_group = mp_group + + def set_quantization_config(self, quantizer): + self.quantizer = quantizer + + def set_hidden_heads(self, hidden_size, num_attention_heads, epsilon, intermediate_size): + """ + Args: + hidden_size: embedding dimension of the model + num_attention_heads: number of attention heads in the model + epsilon: epsilon value for layer norm (same value used for all norms) + intermediate_size: Size of MLP projection. If `DEFAULT_INTERMEDIATE_SIZE` is passed + it is assumed to be `4 * hidden_size` + """ + self.hidden_size = hidden_size + if intermediate_size == DEFAULT_INTERMEDIATE_SIZE: + self.intermediate_size = 4 * hidden_size + else: + self.intermediate_size = intermediate_size + self.num_attention_heads = num_attention_heads + self.layernorm_epsilon = epsilon + + def set_attention(self, qkvw, qkvb, dense_w, dense_b): + self.qkvw = qkvw + self.qkvb = qkvb + self.dense_w = dense_w + self.dense_b = dense_b + + def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): + self._h4h_w = _h4h_w + self._h4h_b = _h4h_b + self._4hh_w = _4hh_w + self._4hh_b = _4hh_b + + def set_layernorm(self, attn_nw, attn_nb, input_nw, input_nb): + self.attn_nw = attn_nw + self.attn_nb = attn_nb + self.input_nw = input_nw + self.input_nb = input_nb + + def apply_weight_quantization(self): + # quantize attention weights + self.attention_quantization() + + # quantize mlp weights + self.mlp_quantization() + + def attention_quantization(self): + self.module.attention.attn_qkvw = self.quantizer.quantize(self.module.attention.attn_qkvw) + self.module.attention.attn_ow = self.quantizer.quantize(self.module.attention.attn_ow) + + def mlp_quantization(self): + self.module.mlp.inter_w = self.quantizer.quantize(self.module.mlp.inter_w) + self.module.mlp.output_w = self.quantizer.quantize(self.module.mlp.output_w) + + def apply_tensor_parallelism(self, mp_replace): + # setup the new Attention module + self.attention_qkv_mp(mp_replace) + self.attention_o_mp(mp_replace) + + # setup the new MLP module + self.mlp_inter_mp(mp_replace) + self.mlp_output_mp(mp_replace) + + # Apply weight quantization + # TODO(cmikeh2): Re-enable this once verified + #self.apply_weight_quantization() + + def attention_qkv_mp(self, mp_replace, reversed_dim=False): + self.module.attention.attn_qkvw = mp_replace.strided_copy(self.module.attention.attn_qkvw, + self.qkvw, + num_splits=3, + int8=reversed_dim) + self.module.attention.attn_qkvb = mp_replace.strided_copy(self.module.attention.attn_qkvb, + self.qkvb, + num_splits=3, + int8=reversed_dim) + + def attention_o_mp(self, mp_replace, reversed_dim=False): + self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow, self.dense_w, int8=reversed_dim) + self.module.attention.attn_ob = mp_replace.copy(self.module.attention.attn_ob, + self.dense_b, + int8=reversed_dim, + allocate_tensor=reversed_dim) + + def mlp_inter_mp(self, mp_replace, reversed_dim=False): + self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w, self._h4h_w, int8=reversed_dim) + self.module.mlp.inter_b = mp_replace.copy(self.module.mlp.inter_b, self._h4h_b, int8=reversed_dim) + + def mlp_output_mp(self, mp_replace, reversed_dim=False): + self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w, self._4hh_w, int8=reversed_dim) + self.module.mlp.output_b = mp_replace.copy(self.module.mlp.output_b, + self._4hh_b, + int8=reversed_dim, + allocate_tensor=reversed_dim) + + def copy_data_to_new_module(self): + params = {'attn_nw': self.attn_nw, 'attn_nb': self.attn_nb} + for key in params: + if params[key] is None: + setattr(self.module.mlp, key, None) + else: + setattr(self.module.mlp, key, + torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) + + params = {'norm_w': self.input_nw, 'norm_b': self.input_nb} + for key in params: + if params[key] is None: + setattr(self.module, key, None) + else: + setattr(self.module, key, + torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) + + def transpose(self): + self.transpose_attention() + self.transpose_mlp() + + def transpose_attention(self): + if self.attn_linear_layer: + self.qkvw = self.transpose_impl(self.qkvw.data) + self.dense_w = self.transpose_impl(self.dense_w.data) + + def transpose_mlp(self): + if self.mlp_linear_layer: + self._h4h_w = self.transpose_impl(self._h4h_w.data) + self._4hh_w = self.transpose_impl(self._4hh_w.data) + + def transpose_impl(self, data): + data = data.contiguous() + data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1)) + data = data.reshape(data.shape[-1], data.shape[-2]) + data.to(get_accelerator().current_device_name()) + return data + + def get_all_params(self): + params = [ + self.attn_nw, + self.attn_nb, + self.input_nw, + self.input_nb, + ] + + params.extend(self.get_attn_params()) + params.extend(self.get_mlp_params()) + + return params + + def get_attn_params(self): + return [self.qkvw, self.qkvb, self.dense_w, self.dense_b] + + def get_mlp_params(self): + return [self._h4h_w, self._h4h_b, self._4hh_w, self._4hh_b] diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/base_moe.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/base_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..4be1b849ba70da04b6b08ea011c27bbbcf96b8bf --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/base_moe.py @@ -0,0 +1,130 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Create a container object to save model-specific tensors using the policy file above. +from .base import * +from deepspeed import comm as dist +import deepspeed.ops.transformer as transformer_inference +from deepspeed.accelerator import get_accelerator + + +class BaseTransformerMoEContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + # Call the init function of the parent class to initialize the tensors and configs from parent class + super().__init__(**kwargs) + + self.num_experts = self.policy.get_num_experts() + self.ep_world_size = dist.get_world_size() + self.local_ep_size = 1 if self.num_experts < self.ep_world_size else self.num_experts // self.ep_world_size + + self.layer_norm_eps = self.config.layer_norm_eps if hasattr(self.config, 'layer_norm_eps') else 1e-12, + + # MoE models will have a list of mlp related tensors + self._h4h_w = [] + self._h4h_b = [] + self._4hh_w = [] + self._4hh_b = [] + + # Residual MoE needs extra parameters + self._res_h4h_w = None + self._res_h4h_b = None + self._res_4hh_w = None + self._res_4hh_b = None + self._res_coef = None + + def create_ds_model_config(self): + self.set_hidden_heads(*self.policy.get_hidden_heads()) + assert self.num_attention_heads % self.mp_size == 0,\ + "To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\ + "This is because the attention computation is partitioned evenly among the parallel GPUs." + + self.ds_model_config = transformer_inference.DeepSpeedMoEInferenceConfig( + hidden_size=self.hidden_size, + heads=self.num_attention_heads, + layer_norm_eps=self.layer_norm_eps, + fp16=self.fp16, + pre_layer_norm=self.pre_layer_norm, + mp_size=self.mp_size, + q_int8=self.quantize, + moe_experts=self.local_ep_size, + global_experts=self.num_experts, + mlp_type=self.config.moe.type, + scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx, + ) + + return self.ds_model_config + + def initialize_tensors(self): + # Set the tensors from policy (user module) to container (DS module) + self.set_attention(*self.policy.attention()) + self.set_mlp(self.config.moe.type) + self.set_layernorm(*self.policy.layernorm()) + + def set_mlp(self, config_moe_type): + if config_moe_type == 'standard': + self._h4h_w, self._h4h_b, \ + self._4hh_w, self._4hh_b = self.policy.mlp() + else: + self._h4h_w, self._h4h_b, self._4hh_w, \ + self._4hh_b, self._res_h4h_w, self._res_h4h_b, \ + self._res_4hh_w, self._res_4hh_b, \ + self._res_coef = self.policy.mlp(config_moe_type) + + def transpose(self): + self.transpose_attention() + self.transpose_mlp() + + if self.config.moe.type == 'residual': + self.transpose_residual() + + def transpose_mlp(self): + self._h4h_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._h4h_w] + self._4hh_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._4hh_w] + + def transpose_residual(self): + self._res_h4h_w.data = self.transpose_impl(self._res_h4h_w.data) + self._res_4hh_w.data = self.transpose_impl(self._res_4hh_w.data) + self._res_coef.data = self.transpose_impl(self._res_coef.data) + + def apply_tensor_parallelism(self, mp_replace): + # setup the new Attention module + self.attention_qkv_mp(mp_replace) + self.attention_o_mp(mp_replace) + + # quantize attention weights + self.attention_quantization() + + # setup the new MLP module + self.mlp_mp() + + def mlp_mp(self): + gpu_index = dist.get_rank() + for ep_index in range(self.local_ep_size): + # mlp inter + self.module.mlp[ep_index].inter_w.data = self._h4h_w[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + self.module.mlp[ep_index].inter_b.data = self._h4h_b[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + + # mlp output + self.module.mlp[ep_index].output_w.data = self._4hh_w[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + self.module.mlp[ep_index].output_b.data = self._4hh_b[gpu_index * self.local_ep_size + ep_index].to( + get_accelerator().current_device_name()) + + def copy_data_to_new_module(self): + self.module.attn_nw.data = self.attn_nw.to(get_accelerator().current_device_name()) + self.module.attn_nb.data = self.attn_nb.to(get_accelerator().current_device_name()) + + self.module.norm_w.data.copy_(self.input_nw.to(get_accelerator().current_device_name())) + self.module.norm_b.data.copy_(self.input_nb.to(get_accelerator().current_device_name())) + + if self.config.moe.type == 'residual': + self.module.res_mlp.inter_w.data = self._res_h4h_w.to(get_accelerator().current_device_name()) + self.module.res_mlp.inter_b.data = self._res_h4h_b.to(get_accelerator().current_device_name()) + self.module.res_mlp.output_w.data = self._res_4hh_w.to(get_accelerator().current_device_name()) + self.module.res_mlp.output_b.data = self._res_4hh_b.to(get_accelerator().current_device_name()) + self.module.res_coef.data = self._res_coef.to(get_accelerator().current_device_name()) diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/bert.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/bert.py new file mode 100644 index 0000000000000000000000000000000000000000..20ae575f45144733a82b609eed21284746cb95d0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/bert.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy + + +class DS_BERTContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + self.return_tuple = True + self.triangular_masking = False + self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFBertLayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=False): + super().__init__(inference, pre_attn_norm=False) + self.client_module = client_module + self.cuda_graph_supported = True + + if HFBertLayerPolicy._orig_layer_class is None: + try: + import transformers + HFBertLayerPolicy._orig_layer_class = [ + transformers.models.bert.modeling_bert.BertLayer, + transformers.models.roberta.modeling_roberta.RobertaLayer + ] + except: + HFBertLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + if self.pre_attn_norm: + attention_layernorm = self.client_module.PostAttentionLayerNorm + else: + attention_layernorm = self.client_module.attention.output.LayerNorm + return self.client_module.attention.self.query.weight.shape[1], \ + self.client_module.attention.self.num_attention_heads, \ + attention_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.attention.self.query.weight + qb = self.client_module.attention.self.query.bias + kw = self.client_module.attention.self.key.weight + kb = self.client_module.attention.self.key.bias + vw = self.client_module.attention.self.value.weight + vb = self.client_module.attention.self.value.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.attention.output.dense.weight, \ + self.client_module.attention.output.dense.bias, \ + + def mlp(self, enable_training=False): + if self.pre_attn_norm: + intermediate_ff = self.client_module.intermediate.dense_act + else: + intermediate_ff = self.client_module.intermediate.dense + + return intermediate_ff.weight, intermediate_ff.bias, \ + self.client_module.output.dense.weight, \ + self.client_module.output.dense.bias + + def layernorm(self): + if self.pre_attn_norm: + attention_layernorm = self.client_module.PostAttentionLayerNorm + transformer_layernorm = self.client_module.PreAttentionLayerNorm + else: + attention_layernorm = self.client_module.attention.output.LayerNorm + transformer_layernorm = self.client_module.output.LayerNorm + return attention_layernorm.weight, \ + attention_layernorm.bias, \ + transformer_layernorm.weight, \ + transformer_layernorm.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/bloom.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/bloom.py new file mode 100644 index 0000000000000000000000000000000000000000..7a9b9ca2065bf969efba949588bc9ac4c95ec8dc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/bloom.py @@ -0,0 +1,143 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.hybrid_engine import HybridEngineContainer +from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy + +from ..policy import maybe_get_lora + +supported_models = {None} + + +class DS_BloomContainer(MetaTensorContainer, HybridEngineContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + # Check transformers version, error if > 4.43.4 (breaks at 4.44.0) + from importlib.metadata import version + v_transformers = version('transformers') + vers = v_transformers.split('.') + major = int(vers[0]) + minor = int(vers[1]) + if major > 4 or (major == 4 and minor > 43): + import sys + sys.exit( + f"Transformers version {v_transformers} exceeds version 4.43.4! After transformers version 4.43.4, BLOOM inference with DeepSpeed is no longer supported." + ) + + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + self.bigscience_bloom = True + self.triangular_masking = False + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + self.module = DeepSpeedBloomInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + self.module.config.invert_mask = False + return self.module + + def attention_qkv_mp(self, mp_replace, reversed_dim=False): + self.module.attention.attn_qkvw = mp_replace.copy(self.module.attention.attn_qkvw, self.qkvw) + self.module.attention.attn_qkvb = mp_replace.copy(self.module.attention.attn_qkvb, self.qkvb) + + def get_lora_matched_pair(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)] + return ret + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, self.policy. + client_module.self_attention.query_key_value, self.policy.client_module.self_attention.dense + ] + ] + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attention.query_key_value.weight', \ + 'self_attention.query_key_value.bias', \ + 'self_attention.dense.weight', \ + 'self_attention.dense.bias', \ + 'mlp.dense_h_to_4h.weight', \ + 'mlp.dense_h_to_4h.bias', \ + 'mlp.dense_4h_to_h.weight', \ + 'mlp.dense_4h_to_h.bias', \ + 'post_attention_layernorm.weight', \ + 'post_attention_layernorm.bias', \ + 'input_layernorm.weight', \ + 'input_layernorm.bias' + ) + for i in range(0, 2): + maybe_copy(module.attention, + sd, + weight_quantizer, + mp_replace, + transformer_param_names[i], + prefix + param_names[i], + qkv=True, + megatron_v2=self.policy.is_megatron_v2, + split_qkv=self.policy.split_qkv) + for i in range(2, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(4, 10): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(10, 12): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i]) + + +class BLOOMLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True, use_load_prefix=True, split_qkv=False): + super().__init__(inference, linear_layer=True, use_load_prefix=use_load_prefix, split_qkv=split_qkv) + self.client_module = client_module + try: + import transformers + BLOOMLayerPolicy._orig_layer_class = transformers.models.bloom.modeling_bloom.BloomBlock + global supported_models + supported_models.update({transformers.models.bloom.modeling_bloom.BloomModel}) + except Exception as e: + print(f"WARNING! Setting BLOOMLayerPolicy._orig_layer_class to None due to Exception: {e}") + BLOOMLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.self_attention.hidden_size, \ + self.client_module.self_attention.num_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + return self.client_module.self_attention.query_key_value.weight, \ + self.client_module.self_attention.query_key_value.bias, \ + self.client_module.self_attention.dense.weight, \ + self.client_module.self_attention.dense.bias, + + def mlp(self, enable_training=False): + return self.client_module.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.dense_4h_to_h.bias + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + self.client_module.post_attention_layernorm.bias, \ + self.client_module.input_layernorm.weight, \ + self.client_module.input_layernorm.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/clip.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..afe4a76086d80ceec7dce15c9cfb2a9b8718b0c4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/clip.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy + + +class DS_CLIPContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFCLIPLayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=False): + super().__init__(inference, pre_attn_norm=True, scale_attention=True) + self.client_module = client_module + self.cuda_graph_supported = True + + if HFCLIPLayerPolicy._orig_layer_class is None: + try: + import transformers + HFCLIPLayerPolicy._orig_layer_class = transformers.models.clip.modeling_clip.CLIPEncoderLayer + except: + HFCLIPLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.self_attn.q_proj.weight.shape[1], \ + self.client_module.self_attn.num_heads, \ + self.client_module.layer_norm1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + qb = self.client_module.self_attn.q_proj.bias + kw = self.client_module.self_attn.k_proj.weight + kb = self.client_module.self_attn.k_proj.bias + vw = self.client_module.self_attn.v_proj.weight + vb = self.client_module.self_attn.v_proj.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.self_attn.out_proj.weight, \ + self.client_module.self_attn.out_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.fc1.weight, \ + self.client_module.mlp.fc1.bias, \ + self.client_module.mlp.fc2.weight, \ + self.client_module.mlp.fc2.bias + + def layernorm(self): + return self.client_module.layer_norm2.weight, \ + self.client_module.layer_norm2.bias, \ + self.client_module.layer_norm1.weight, \ + self.client_module.layer_norm1.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/distil_bert.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/distil_bert.py new file mode 100644 index 0000000000000000000000000000000000000000..ecd0562438b5ac634cf8b4536fd3413d0f9ed9d8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/distil_bert.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy + + +class DS_DistilBERTContainer(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + self.triangular_masking = False + self.return_single_tuple = True + self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFDistilBertLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=False, preln=False): + super().__init__(inference) + self.client_module = client_module + self.preln = preln + self.cuda_graph_supported = True + if HFDistilBertLayerPolicy._orig_layer_class is None: + try: + import transformers + HFDistilBertLayerPolicy._orig_layer_class = [ + transformers.models.distilbert.modeling_distilbert.TransformerBlock, + ] + except: + HFDistilBertLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attention.q_lin.weight.shape[1], \ + self.client_module.attention.n_heads, \ + self.client_module.sa_layer_norm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.attention.q_lin.weight + qb = self.client_module.attention.q_lin.bias + kw = self.client_module.attention.k_lin.weight + kb = self.client_module.attention.k_lin.bias + vw = self.client_module.attention.v_lin.weight + vb = self.client_module.attention.v_lin.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.attention.out_lin.weight, \ + self.client_module.attention.out_lin.bias + + def mlp(self, enable_training=False): + intermediate_ff = self.client_module.ffn.lin1 + + return intermediate_ff.weight, intermediate_ff.bias, \ + self.client_module.ffn.lin2.weight, \ + self.client_module.ffn.lin2.bias + + def layernorm(self): + attention_layernorm = self.client_module.sa_layer_norm + transformer_layernorm = self.client_module.output_layer_norm + return attention_layernorm.weight, \ + attention_layernorm.bias, \ + transformer_layernorm.weight, \ + transformer_layernorm.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__init__.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fc2eb2a65531e61bcea078c0035bd2ee6ee861a6 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .gated_mlp import HybridGatedMLPContainer +from .megatron import MegatronContainer +from .meta_tensor import MetaTensorContainer +from .split_qkv import HybridSplitQKVContainer diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8aa9f7147995a241771557c957e7677bd9b6f28a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5dffa866f3e06a3c882a28678e1ad8feb0716ada Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6041ae58a91a706b6d3de03e4aa6c0243ac2d6a5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd8c1a50a9f33cec546b62d598e3653116bfc3ce Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35b636ddc86903a4272beb628d3faf18a4f2b9d4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e89264ca95b5c6ded17e0506c0a94dc537c7b53 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0321832807836a6345c4fcce5f5340c0aa5739c3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..24f0826db14ed08f373a2824f845ac6f9d7d9508 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import abstractmethod + +from .hybrid_engine import HybridEngineContainer + + +class HybridGatedMLPContainer(HybridEngineContainer): + """ + The HybridGatedMLPContainer supports models for which the first MLP layer + is represented with two separate weights, one for the activation function + and one for the gating function. + """ + + def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): + super().set_mlp(_h4h_w, _h4h_b, _4hh_w, _4hh_b) + self.set_mlp_gate() + + @abstractmethod + def set_mlp_gate(self): + """ + In `set_mlp_gate`, it is necessary to populate the following variables (where appropriate) + for the given model: + self.inter_up_w: inter up weight + self.inter_up_b: inter up bias + self.inter_gate_w: inter gate weight + self.inter_gate_b: inter gate bias + If the parameter does not exist in the original model, set the attribute to None. + """ + raise NotImplementedError("A set_mlp_gate() function must be defined in the model container \ + in order to set the unfused inter up and gate tensors.") + + def mlp_inter_mp(self, mp_replace, reversed_dim=False): + # Only need to alter behavior if we can't do the normal destructive copy + if self.module.mlp.inter_w is None: + params = [ + (self.module.mlp.inter_up_w, self.inter_up_w), + (self.module.mlp.inter_up_b, self.inter_up_b), + (self.module.mlp.inter_gate_w, self.inter_gate_w), + (self.module.mlp.inter_gate_b, self.inter_gate_b), + ] + for dst, src in params: + dst = mp_replace.copy(dst[:self.inter_up_w.shape[0] // mp_replace.mp_size], + src, + int8=reversed_dim, + allocate_tensor=reversed_dim) if src is not None else None + else: + self.module.mlp.inter_w = mp_replace.strided_copy(self.module.mlp.inter_w, + self._h4h_w, + num_splits=2, + int8=reversed_dim) + self.module.mlp.inter_b = mp_replace.strided_copy(self.module.mlp.inter_b, + self._h4h_b, + num_splits=2, + int8=reversed_dim) + + def release_mlp(self): + super().release_mlp() + gated_mlp_params = [ + (self.module.mlp.inter_up_w, self.inter_up_w), + (self.module.mlp.inter_up_b, self.inter_up_b), + (self.module.mlp.inter_gate_w, self.inter_gate_w), + (self.module.mlp.inter_gate_b, self.inter_gate_b), + ] + + self._release_params(gated_mlp_params) + + def reset_mlp(self): + self._h4h_w.data[:self.inter_up_w.shape[0]] = self.inter_up_w.data + self._h4h_w.data[self.inter_up_w.shape[0]:] = self.inter_gate_w.data + + if self.inter_up_b is not None: + self._h4h_b.data[:self.inter_up_b.shape[0]] = self.inter_up_b.data + self._h4h_b.data[self.inter_up_b.shape[0]:] = self.inter_gate_b.data + + inter_data = [self.inter_up_w.data, self.inter_gate_w.data] + if self.inter_up_b is not None: + inter_data.extend([self.inter_up_b.data, self.inter_gate_b.data]) + + self.inter_up_w.data = self._h4h_w.data[:self.inter_up_w.shape[0]] + self.inter_gate_w.data = self._h4h_w.data[self.inter_up_w.shape[0]:] + + if self.inter_up_b is not None: + self.inter_up_b.data = self._h4h_b.data[:self.inter_up_b.shape[0]] + self.inter_gate_b.data = self._h4h_b.data[self.inter_up_b.shape[0]:] + + for data in inter_data: + del data + + def set_mlp_params_wo_copy(self, Z3_enabled=False): + self.module.mlp.output_w = self._4hh_w + self.module.mlp.output_b = self._4hh_b + + if not Z3_enabled: + # In initialize_tensors, we create a fused inter projection with the appropriate shape + # and copy the up projection and gate projection into it + self.module.mlp.inter_w = self._h4h_w + self.module.mlp.inter_b = self._h4h_b + + self.inter_up_w.data = self._h4h_w[:self.inter_up_w.shape[0], :] + self.inter_gate_w.data = self._h4h_w[self.inter_up_w.shape[0]:, :] + + if self.inter_up_b is not None: + self.inter_up_b.data = self._h4h_b[:self.inter_up_w.shape[0]] if self._h4h_b is not None else None + self.inter_gate_b.data = self._h4h_b[self.inter_up_w.shape[0]:] if self._h4h_b is not None else None + else: + self.module.mlp.inter_up_w = self.inter_up_w + self.module.mlp.inter_up_b = self.inter_up_b + self.module.mlp.inter_gate_w = self.inter_gate_w + self.module.mlp.inter_gate_b = self.inter_gate_b + + def get_mlp_params(self): + params = super().get_mlp_params() + params.extend([self.inter_up_w, self.inter_up_b, self.inter_gate_w, self.inter_gate_b]) + return params diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..3f702abcf319a6db997b6e2607dcba26f3440841 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py @@ -0,0 +1,212 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod +from typing import List, Tuple + +import torch + + +class HybridEngineContainer(ABC): + """ + This container identifies which methods need to be overridden in addition to + the base container to enable use in the RLHF pipeline. These methods are not + necessary for inference alone. + + NOTE: If you are using this feature with a container that + also inherits from `MetaTensorContainer`, ensure that `MetaTensorContainer` + is inherited before `HybridEngineContainer` in the class definition. + """ + + def initialize_tensors(self, enable_training=False): + """ + Same purposes as the base container, but also grabs the hooks for any LoRA + parameters. If it's necessary to override specific sub-components of the model, + it's best to augment the specific `set_[component]` itself rather than modifying + the `initialize_tensors` method. See the `HybridSplitQKVContainer` for an example. + """ + super().initialize_tensors(enable_training=enable_training) + self.set_lora_params() + + def transform_for_training(self): + """ + If the views on certain parameters are largely incompatible, it may be necessary to do + more substantial transformations to the parameters. This method should be overridden to + transform the inference format to what is necessary for training. + """ + pass + + def transform_for_inference(self): + """ + If the views on certain parameters are largely incompatible, it may be necessary to do + more substantial transformations to the parameters. This method should be overridden to + transform the training format to what is necessary for inference. + """ + pass + + @abstractmethod + def set_lora_params(self): + """ + If available, set the LoRA parameters for the module. An implementation + for this would iterate over all parameters of the model and use the `maybe_get_lora` helper + method to check if the parameter does in fact have any LoRA params. + """ + raise NotImplementedError("A set_lora_params() function must be defined for the relevant parameters.") + + @abstractmethod + def get_lora_matched_pair(self): + """Get the pair of lora params and its matched model parameters.""" + raise NotImplementedError("get_lora_matched_pair() must be defined for the relevant parameters.") + + def fuse_lora(self): + """Fuse the LoRA parameters for the inference mode.""" + for maybe_lora_param, param in self.get_lora_matched_pair(): + if len(maybe_lora_param) == 3: + lora_right_weight, \ + lora_left_weight, \ + lora_scaling = maybe_lora_param + param.data += lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) + + def unfuse_lora(self): + """Unfuse the LoRA parameters for the training mode.""" + for maybe_lora_param, param in self.get_lora_matched_pair(): + if len(maybe_lora_param) == 3: + lora_right_weight, \ + lora_left_weight, \ + lora_scaling = maybe_lora_param + param.data -= lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) + + def apply_tensor_parallelism(self, mp_replace, reversed_dim=False): + """ + Add support for reversed dim in tensor parallelism. If necessary, override + the called methods to handle partitioned weights (i.e. if qkv is split, override + the `attention_qkv_mp` method). If the model component is not split, it should + be safe to use the default implementation. + """ + # Setup the new Attention module + self.attention_qkv_mp(mp_replace, reversed_dim=reversed_dim) + self.attention_o_mp(mp_replace, reversed_dim=reversed_dim) + + # Setup the new MLP module + self.mlp_inter_mp(mp_replace, reversed_dim=reversed_dim) + self.mlp_output_mp(mp_replace, reversed_dim=reversed_dim) + + # Apply weight quantization + # TODO(cmikeh2): Re-enable this once verified + #self.apply_weight_quantization() + + def _release_params(self, param_pairs: List[Tuple[torch.Tensor, torch.Tensor]]): + """ + Helper for `release_[component]` methods. Accepts a list of tuples where the first + element is the module param that needs to be deleted, and the second is the reassignment + from the container. + """ + for module_param, container_param in param_pairs: + if module_param is not None: + del module_param + module_param = container_param + + def release_memory(self): + """ + Delete module parameters if they exist and point them back to the container. The primary + purpose of this is for TP-inference with ZeRO-3. In this scenario, we need to delete the + parameters we've created for inference to free their memory. + """ + general_params = [ + (self.module.attention.attn_ow, self.dense_w), + (self.module.attention.attn_ob, self.dense_b), + (self.module.mlp.attn_nw, self.attn_nw), + (self.module.mlp.attn_nb, self.attn_nb), + (self.module.norm_w, self.input_nw), + (self.module.norm_b, self.input_nb), + ] + + self._release_params(general_params) + + self.release_qkv() + self.release_mlp() + + def release_qkv(self): + """ + Release for QKV parameters (as well as any aliases). + """ + qkv_params = [ + (self.module.attention.attn_qkvw, self.qkvw), + (self.module.attention.attn_qkvb, self.qkvb), + ] + + self._release_params(qkv_params) + + def release_mlp(self): + """ + Release for MLP parameters (as well as any aliases). + """ + mlp_params = [ + (self.module.mlp.inter_w, self._h4h_w), + (self.module.mlp.inter_b, self._h4h_b), + (self.module.mlp.output_w, self._4hh_w), + (self.module.mlp.output_b, self._4hh_b), + ] + + self._release_params(mlp_params) + + def reset_params(self): + """ + The purpose of reset params is to get the weights from the FP16 training + copy of the model and copy to them to contiguous inference view. This only needs + to be performed when the container parameters cannot be used directly for inference. + """ + self.reset_qkv() + self.reset_mlp() + + def reset_qkv(self): + """ + Perform any necessary resets of the model parameters for the QKV components. + """ + pass + + def reset_mlp(self): + """ + Perform any necessary resets of the model parameters for the MLP components. + """ + pass + + def get_lora_params(self): + """ + Return a list of all parameters that would have LoRA for the module. + """ + if not hasattr(self, "lora_params"): + self.set_lora_params() + return self.lora_params + + def set_params_wo_copy(self, Z3_enabled=False): + """ + Rather than copying into, set the parameters directly. This is necessary to provide + an inexpensive (low-memory-overhead) view onto the FP16 forward weights. + """ + self.module.mlp.attn_nw = self.attn_nw + self.module.mlp.attn_nb = self.attn_nb + self.module.norm_w = self.input_nw + self.module.norm_b = self.input_nb + self.set_attn_params_wo_copy(Z3_enabled=Z3_enabled) + self.set_mlp_params_wo_copy(Z3_enabled=Z3_enabled) + + def set_attn_params_wo_copy(self, **kwargs): + """ + Narrower sub-method for finer grained overriding. + """ + self.module.attention.attn_ow = self.dense_w + self.module.attention.attn_ob = self.dense_b + self.module.attention.attn_qkvw = self.qkvw + self.module.attention.attn_qkvb = self.qkvb + + def set_mlp_params_wo_copy(self, **kwargs): + """ + Narrower sub-method for finer grained overriding. + """ + self.module.mlp.inter_w = self._h4h_w + self.module.mlp.inter_b = self._h4h_b + self.module.mlp.output_w = self._4hh_w + self.module.mlp.output_b = self._4hh_b diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py new file mode 100644 index 0000000000000000000000000000000000000000..d40f2a6b57e86087dc13ff2b74a389c67826fa76 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from .hybrid_engine import HybridEngineContainer +from .megatron import MegatronContainer + + +class HybridMegatronContainer(MegatronContainer, HybridEngineContainer): + + def _align_qkv(self, x: torch.Tensor): + """ + Internal helper for accepting the head-contiguous weight matrix and chunking + the query, key, and value components. + """ + attention_head_size = x.shape[0] // self.num_attention_heads + new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] + x_1 = x.view(*new_x_shape) + div_dim = len(x_1.size()) - 2 if len(x.shape) == 2 else -1 + (q, k, v) = torch.split(x_1, (x_1.shape[div_dim] // 3), dim=div_dim) + if len(q.shape) > 2: + x.data.copy_( + torch.cat((q.reshape(-1, q.shape[-1]), k.reshape(-1, q.shape[-1]), v.reshape(-1, q.shape[-1])), + dim=0).reshape(x.shape)) + else: + x.data.copy_(torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)) + + def transform_for_inference(self) -> None: + """ + Overrides the HybridEngineContainer implementation. + + The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V + are sequential in memory. This is different from the default layout in which all of the Qs + are sequential, followed by all of the Ks, and then all of the Vs. Here, we take the default + layout and transform it to the inference layout. + """ + if hasattr(self.qkvw, 'ds_id'): + from deepspeed.runtime.zero import GatheredParameters + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + param_list = [self.qkvw, self.qkvb] + non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + with GatheredParameters(non_active_params): + self._align_qkv(self.qkvw) + self._align_qkv(self.qkvb) + else: + self._align_qkv(self.qkvw) + self._align_qkv(self.qkvb) + + def _partition_qkv(self, x: torch.Tensor): + """ + Internal helper for taking contiguous QKV and partitioning it for contiguous + heads. + """ + q_k_v = torch.split(x, (x.shape[0] // 3), dim=0) + attention_head_size = q_k_v[0].shape[0] // self.num_attention_heads + new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] + q, k, v = [data.view(*new_x_shape) for data in q_k_v] + if len(q.shape) > 2: + x.data.copy_(torch.cat((q, k, v), dim=-2).reshape(-1, q.shape[-1])) + else: + x.data.copy_(torch.cat((q, k, v), dim=-1).reshape(-1)) + + def transform_for_training(self): + """ + Overrides the HybridEngineContainer implementation. + + The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V + are sequential in memory. This is different from the default layout in which all of the Qs + are sequential, followed by all of the Ks, and then all of the Vs. This function takes the inference format and reverts it back to the default format. + """ + # If parameter is distributed, handle gathering it + if hasattr(self.qkvw, 'ds_id'): + from deepspeed.runtime.zero import GatheredParameters + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + param_list = [self.qkvw, self.qkvb] + non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + with GatheredParameters(non_active_params): + self._partition_qkv(self.qkvw) + self._partition_qkv(self.qkvb) + else: + self._partition_qkv(self.qkvw) + self._partition_qkv(self.qkvb) diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/megatron.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/megatron.py new file mode 100644 index 0000000000000000000000000000000000000000..4daccf7d7c8d4db06bd65243bd83a5f1fd50df2e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/megatron.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from abc import ABC + + +class MegatronContainer(ABC): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.megatron_v2 = self.policy.is_megatron_v2 + + def _align_qkv_transposed(self, x): + attention_head_size = x.shape[-1] // self.num_attention_heads + new_x_shape = x.size()[:-1] + (self.num_attention_heads, attention_head_size) + x_1 = x.view(*new_x_shape) + (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=(x_1.dim() - 1)) + if len(q.shape) > 2: + return torch.cat((q.reshape(q.shape[0], -1), k.reshape(q.shape[0], -1), v.reshape(q.shape[0], -1)), + dim=-1).reshape(x.shape) + else: + return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape) + + def transpose(self): + super().transpose() + if self.megatron_v2: + self.qkvw = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvw).contiguous()) + self.qkvb = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvb).contiguous()) diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..57b136663be370863dad81cc9d6a4a29425376b3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py @@ -0,0 +1,70 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod +from packaging import version as pkg_version +import torch + + +class MetaTensorContainer(ABC): + """ + NOTE: If you are using this feature with a container that + also inherits from `HybridEngineContainer`, ensure that `MetaTensorContainer` + is inherited before `HybridEngineContainer` in the class definition. + """ + + def __init__(self, **kwargs): + if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__): + raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") + super().__init__(**kwargs) + self.is_meta = False + self.ckpt_load_enabled = True + + def initialize_tensors(self, enable_training=False): + super().initialize_tensors(enable_training=enable_training) + self.is_meta = self.qkvw.is_meta + + def apply_tensor_parallelism(self, mp_replace, **kwargs): + if self.is_meta: + if self.qkvb is None: + self.module.attention.attn_qkvb = None + if self.dense_b is None: + self.module.attention.attn_ob = None + else: + super().apply_tensor_parallelism(mp_replace, **kwargs) + + def copy_data_to_new_module(self): + if self.is_meta: + if self.attn_nw is None: + self.module.mlp.attn_nw = self.attn_nw + self.module.mlp.attn_nb = self.attn_nb + else: + super().copy_data_to_new_module() + + def transpose(self): + if not self.is_meta: + super().transpose() + + @abstractmethod + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + """ + Load all the transformer parameter from the checkpoint file (sd). + In addition to the parameter names, we require two + more parameters to help read the data correctly + from the checkpoint and split the qkv heads in the + right order: + 1. `use_load_prefix` (Default: False): this specifies + whether we need to use the name of first abstraction + layer of the model for searching the parameter's name + in a checkpoint file. For more information of how this + is used please see + https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/load_checkpoint.py + 2. `split_qkv` (Default: True): we use this flag when splitting + the qkv parameter into heads. If it is False, it means the heads + of q, k, and v are stored together and needs to split in the + DeepSpeed-Inference API. + """ + raise NotImplementedError("A load_params() function must be defined in the model container \ + when inheriting the MetaTensorContainer feature") diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/split_qkv.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/split_qkv.py new file mode 100644 index 0000000000000000000000000000000000000000..f4c14d4e425a7e9096bae3bf5788ab7b9ca0dcc2 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/features/split_qkv.py @@ -0,0 +1,159 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import abstractmethod +import torch + +from .hybrid_engine import HybridEngineContainer + + +class HybridSplitQKVContainer(HybridEngineContainer): + + def set_attention(self, qkvw, qkvb, dense_w, dense_b): + super().set_attention(qkvw, qkvb, dense_w, dense_b) + self.set_q_k_v() + + @abstractmethod + def set_q_k_v(self): + """ + In `set_q_k_v`, it is necessary to populate the following variables (where appropriate) + for the given model: + self.qw: q weight + self.qb: q bias + self.kw: k weight + self.kb: k bias + self.vw: v weight + self.vb: v bias + """ + raise NotImplementedError("A set_q_k_v() function must be defined in the model container \ + in order to set the unfused q, k, and v tensors.") + + def attention_qkv_mp(self, mp_replace, reversed_dim=False): + # Only need to alter + if self.module.attention.attn_qkvw is None: + params = [ + (self.module.attention.attn_qw, self.qw), + (self.module.attention.attn_qb, self.qb), + (self.module.attention.attn_kw, self.kw), + (self.module.attention.attn_kb, self.kb), + (self.module.attention.attn_vw, self.vw), + (self.module.attention.attn_vb, self.vb), + ] + for dst, src in params: + dst = mp_replace.copy( + dst[:self.qw.shape[0] // mp_replace.mp_size], src, int8=reversed_dim, + allocate_tensor=reversed_dim) if src is not None else None + else: + super().attention_qkv_mp(mp_replace) + + def release_qkv(self): + super().release_qkv() + split_qkv_params = [ + (self.module.attention.attn_qw, self.qw), + (self.module.attention.attn_qb, self.qb), + (self.module.attention.attn_kw, self.kw), + (self.module.attention.attn_kb, self.kb), + (self.module.attention.attn_vw, self.vw), + (self.module.attention.attn_vb, self.vb), + ] + + self._release_params(split_qkv_params) + + def reset_qkv(self): + self.qkvw.data[:self.qw.shape[0]] = self.qw.data + self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data + self.qkvw.data[2 * self.qw.shape[0]:] = self.vw.data + + qkv_data = [self.qw.data, self.kw.data, self.vw.data] + + self.qw.data = self.qkvw.data[:self.qw.shape[0]] + self.kw.data = self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vw.data = self.qkvw.data[2 * self.qw.shape[0]:] + + if self.qkvb is not None: + self.qkvb.data[:self.qw.shape[0]] = self.qb.data + self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data + self.qkvb.data[2 * self.qw.shape[0]:] = self.vb.data + + qkv_data.extend([self.qb.data, self.kb.data, self.vb.data]) + + self.qb.data = self.qkvb.data[:self.qw.shape[0]] + self.kb.data = self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vb.data = self.qkvb.data[2 * self.qw.shape[0]:] + + for data in qkv_data: + del data + + def reset_qkv_experimental(self): + """ + WIP - experimental and likely to be changed/improved. + Unused by keeping for now. + """ + if self.module.attention.attn_qkvw is None: + self.module.attention.attn_qkvw = torch.empty(self.qw.shape[0] * 3, + self.qw.shape[0], + dtype=self.qw.dtype, + device=self.qw.device) + self.module.attention.attn_qkvb = torch.empty(self.qw.shape[0] * 3, + dtype=self.qw.dtype, + device=self.qw.device) + self.module.attention.attn_qkvw.data[:self.qw.shape[0]] = self.qw.data + self.module.attention.attn_qkvb.data[:self.qw.shape[0]] = self.qb.data + self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data + self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data + self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] = self.vw.data + self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] = self.vb.data + + qkv_data = [self.qw.data, \ + self.qb.data, \ + self.kw.data, \ + self.kb.data, \ + self.vw.data, \ + self.vb.data] + + self.qw.data = self.module.attention.attn_qkvw.data[:self.qw.shape[0]] + self.qb.data = self.module.attention.attn_qkvb.data[:self.qw.shape[0]] + self.kw.data = self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.kb.data = self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vw.data = self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] + self.vb.data = self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] + + for data in qkv_data: + del data + + def set_attn_params_wo_copy(self, Z3_enabled=False): + self.module.attention.attn_ow = self.dense_w + self.module.attention.attn_ob = self.dense_b + if not Z3_enabled: + # In initialize_tensors, we create a fused qkvw with the appropriate shape + # and copy the qw, qb, kw, kb, vw, vb into it + self.module.attention.attn_qkvw = self.qkvw + self.module.attention.attn_qkvb = self.qkvb + + # We reset the data for qw (which is the original model parameter) to point + # to the fused weight matrix we have created here + self.qw.data = self.qkvw[:self.qw.shape[0], :] + self.kw.data = self.qkvw[self.qw.shape[0]:2 * self.qw.shape[0], :] + self.vw.data = self.qkvw[self.qw.shape[0] * 2:, :] + + # Assume if one of the biases is not None, then all of them are not None + if self.qb is not None: + self.qb.data = self.qkvb[:self.qw.shape[0]] + self.kb.data = self.qkvb[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vb.data = self.qkvb[self.qw.shape[0] * 2:] + else: + # In ZeRO-3 this will be managed by ZeRO and handled separately in the + # forward of ds_attention + self.module.attention.attn_qw = self.qw + self.module.attention.attn_qb = self.qb + self.module.attention.attn_kw = self.kw + self.module.attention.attn_kb = self.kb + self.module.attention.attn_vw = self.vw + self.module.attention.attn_vb = self.vb + + def get_attn_params(self): + params = super().get_attn_params() + params.extend([self.qw, self.qb, self.kw, self.kb, self.vw, self.vb]) + return params diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/gpt2.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gpt2.py new file mode 100644 index 0000000000000000000000000000000000000000..7a19aac34b447a7d5116ccceeb7687faa61ea174 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gpt2.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +from ..policy import TransformerPolicy + + +class DS_GPT2Container(BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + +class HFGPT2LayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True): + # HuggingFace GPT2 uses convolutional layer instead of linear layer + super().__init__(inference, linear_layer=False) + self.client_module = client_module + try: + import transformers + HFGPT2LayerPolicy._orig_layer_class = transformers.models.gpt2.modeling_gpt2.GPT2Block + except: + HFGPT2LayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attn.embed_dim, \ + self.client_module.attn.num_heads, \ + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + return self.client_module.attn.c_attn.weight, \ + self.client_module.attn.c_attn.bias, \ + self.client_module.attn.c_proj.weight, \ + self.client_module.attn.c_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.c_fc.weight, \ + self.client_module.mlp.c_fc.bias, \ + self.client_module.mlp.c_proj.weight, \ + self.client_module.mlp.c_proj.bias + + def layernorm(self): + return self.client_module.ln_2.weight, \ + self.client_module.ln_2.bias, \ + self.client_module.ln_1.weight, \ + self.client_module.ln_1.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptj.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptj.py new file mode 100644 index 0000000000000000000000000000000000000000..17c0a5027a4c94dcc206b69ad6f3fa6d1d746e17 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptj.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.split_qkv import HybridSplitQKVContainer +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from ..policy import maybe_copy_qkv + +from ..policy import maybe_get_lora + + +class DS_GPTJContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.fc_in, self.policy.client_module.mlp.fc_out, + self.policy.client_module.attn.q_proj, self.policy.client_module.attn.k_proj, + self.policy.client_module.attn.v_proj, self.policy.client_module.attn.out_proj + ] + ] + + def get_lora_matched_pair(self): + fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), + (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.attn.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.attn.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.attn.v_proj.weight + self.vb = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attn.q_proj.weight', \ + 'attn.k_proj.weight', \ + 'attn.v_proj.weight', \ + 'attn.out_proj.weight', \ + 'mlp.fc_in.weight', \ + 'mlp.fc_in.bias', \ + 'mlp.fc_out.weight', \ + 'mlp.fc_out.bias', \ + 'ln_1.weight', \ + 'ln_1.bias' + ) + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + for i in range(4, 8): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(8, 10): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i + 2], + prefix + param_names[i]) + + +class HFGPTJLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True): + super().__init__(inference, scale_attention=True) + self.client_module = client_module + try: + import transformers + HFGPTJLayerPolicy._orig_layer_class = transformers.models.gptj.modeling_gptj.GPTJBlock + except: + HFGPTJLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attn.embed_dim, \ + self.client_module.attn.num_attention_heads, \ + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.attn.q_proj.weight + kw = self.client_module.attn.k_proj.weight + vw = self.client_module.attn.v_proj.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.attn.out_proj.weight, \ + None, + + def mlp(self, enable_training=False): + return self.client_module.mlp.fc_in.weight, \ + self.client_module.mlp.fc_in.bias, \ + self.client_module.mlp.fc_out.weight, \ + self.client_module.mlp.fc_out.bias + + def layernorm(self): + return None, \ + None, \ + self.client_module.ln_1.weight, \ + self.client_module.ln_1.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptneo.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptneo.py new file mode 100644 index 0000000000000000000000000000000000000000..fca673b375e18a9b74741f86a9692bfd386f6a51 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptneo.py @@ -0,0 +1,145 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.split_qkv import HybridSplitQKVContainer +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from ..policy import maybe_copy_qkv + +from ..policy import maybe_get_lora + + +class DS_GPTNEOContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.c_fc, self.policy.client_module.mlp.c_proj, + self.policy.client_module.attn.attention.q_proj, self.policy.client_module.attn.attention.k_proj, + self.policy.client_module.attn.attention.v_proj, self.policy.client_module.attn.attention.out_proj + ] + ] + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.attn.attention.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.attn.attention.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.attn.attention.v_proj.weight + self.vb = None + + def get_lora_matched_pair(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), + (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attn.attention.q_proj.weight', \ + 'attn.attention.k_proj.weight', \ + 'attn.attention.v_proj.weight', \ + 'attn.attention.out_proj.weight', \ + 'attn.attention.out_proj.bias', \ + 'mlp.c_fc.weight', \ + 'mlp.c_fc.bias', \ + 'mlp.c_proj.weight', \ + 'mlp.c_proj.bias', \ + 'ln_2.weight', \ + 'ln_2.bias', \ + 'ln_1.weight', \ + 'ln_1.bias' + ) + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 5): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + for i in range(5, 11): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + for i in range(11, 13): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + + +class HFGPTNEOLayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=True): + super().__init__(inference, scale_attention=False) + self.client_module = client_module + try: + import transformers + HFGPTNEOLayerPolicy._orig_layer_class = transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoBlock + except: + HFGPTNEOLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attn.attention.embed_dim, \ + self.client_module.attn.attention.num_heads, \ + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def get_q_k_v(self): + return self.client_module.attn.attention.q_proj.weight, \ + None, \ + self.client_module.attn.attention.k_proj.weight, \ + None, \ + self.client_module.attn.attention.v_proj.weight, \ + None + + def attention(self, enable_training=False): + qw = self.client_module.attn.attention.q_proj.weight + kw = self.client_module.attn.attention.k_proj.weight + vw = self.client_module.attn.attention.v_proj.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.attn.attention.out_proj.weight, \ + self.client_module.attn.attention.out_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.c_fc.weight, \ + self.client_module.mlp.c_fc.bias, \ + self.client_module.mlp.c_proj.weight, \ + self.client_module.mlp.c_proj.bias + + def layernorm(self): + return self.client_module.ln_2.weight, \ + self.client_module.ln_2.bias, \ + self.client_module.ln_1.weight, \ + self.client_module.ln_1.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptneox.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptneox.py new file mode 100644 index 0000000000000000000000000000000000000000..3978b166b329b341232da61811ecaf5683c5578d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/gptneox.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.meta_tensor import MetaTensorContainer +from .features.hybrid_megatron import HybridMegatronContainer +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from packaging import version as pkg_version + +from ..policy import maybe_get_lora + + +class DS_GPTNEOXContainer(MetaTensorContainer, HybridMegatronContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + + if self.megatron_v2: + self.module.config.rotate_half = True + self.module.config.rotate_every_two = False + + return self.module + + def get_lora_matched_pair(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)] + return ret + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + if GPTNEOXLayerPolicy.version == 0: + attention = self.policy.client_module.attention + else: + attention = self.policy.client_module.self_attention + + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, + attention.query_key_value, attention.dense + ] + ] + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attention.query_key_value.weight', \ + 'attention.query_key_value.bias', \ + 'attention.dense.weight', \ + 'attention.dense.bias', \ + 'mlp.dense_h_to_4h.weight', \ + 'mlp.dense_h_to_4h.bias', \ + 'mlp.dense_4h_to_h.weight', \ + 'mlp.dense_4h_to_h.bias', \ + 'post_attention_layernorm.weight', \ + 'post_attention_layernorm.bias', \ + 'input_layernorm.weight', \ + 'input_layernorm.bias' + ) + for i in range(0, 2): + maybe_copy(module.attention, + sd, + weight_quantizer, + mp_replace, + transformer_param_names[i], + prefix + param_names[i], + qkv=True, + megatron_v2=self.policy.is_megatron_v2, + split_qkv=self.policy.split_qkv, + heads=self.policy.client_module.attention.num_attention_heads) + for i in range(2, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(4, 10): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], + prefix + param_names[i]) + for i in range(10, 12): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i]) + + +class GPTNEOXLayerPolicy(TransformerPolicy): + _orig_layer_class = None + version = 0 + + def __init__(self, client_module, inference=True, megatron_v2=True, split_qkv=False): + super().__init__(inference, megatron_v2=megatron_v2, split_qkv=split_qkv) + self.client_module = client_module + if GPTNEOXLayerPolicy._orig_layer_class is None: + if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): + GPTNEOXLayerPolicy._orig_layer_class = None + else: + try: + from transformers import GPTNeoXLayer + GPTNEOXLayerPolicy._orig_layer_class = GPTNeoXLayer + except ImportError: + GPTNEOXLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + if GPTNEOXLayerPolicy.version == 0: + attention = self.client_module.attention + else: + attention = self.client_module.self_attention + + return self.client_module.attention.hidden_size, \ + self.client_module.attention.num_attention_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + if GPTNEOXLayerPolicy.version == 0: + attention = self.client_module.attention + else: + attention = self.client_module.self_attention + + return attention.query_key_value.weight, \ + attention.query_key_value.bias, \ + attention.dense.weight, \ + attention.dense.bias + + def mlp(self, enable_training=False): + return self.client_module.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.dense_4h_to_h.bias + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + self.client_module.post_attention_layernorm.bias, \ + self.client_module.input_layernorm.weight, \ + self.client_module.input_layernorm.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/internlm.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/internlm.py new file mode 100644 index 0000000000000000000000000000000000000000..31255d4b3ca5a0199ee572219b6d75fd70d478c3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/internlm.py @@ -0,0 +1,181 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import importlib + +import torch +from torch.nn.parameter import Parameter + +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +from deepspeed.utils.types import ActivationFuncType, NormType + +from ..policy import (TransformerPolicy, maybe_copy, maybe_copy_geglu, maybe_copy_qkv, maybe_get_lora, + transformer_param_names) +from .base import * +from .features import HybridGatedMLPContainer, HybridSplitQKVContainer + + +class DS_InternLMContainer(HybridGatedMLPContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + _config.rotate_half = True + _config.rotate_every_two = False + _config.rotary_dim = self.hidden_size // self.num_attention_heads + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight, + self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight, + self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight, + self.policy.client_module.self_attn.o_proj.weight + ] + ] + + def get_lora_matched_pair(self): + up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), + (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = self.policy.client_module.self_attn.q_proj.bias + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = self.policy.client_module.self_attn.k_proj.bias + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = self.policy.client_module.self_attn.v_proj.bias + + def set_mlp_gate(self): + """ + Necessary to implement for `HybridGatedMLPContainer` + """ + self.inter_up_w = self.policy.client_module.mlp.up_proj.weight + self.inter_up_b = None + self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight + self.inter_gate_b = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attn.q_proj.weight', \ + 'self_attn.k_proj.weight', \ + 'self_attn.v_proj.weight', \ + 'self_attn.o_proj.weight', \ + 'mlp.up_proj.weight', \ + 'mlp.gate_proj.weight', \ + 'mlp.down_proj.weight', \ + 'input_layernorm.weight', \ + 'post_attention_layernorm.weight' + 'self_attn.q_proj.bias', \ + 'self_attn.k_proj.bias', \ + 'self_attn.v_proj.bias', \ + 'self_attn.o_proj.bias', \ + ) + + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvb', [prefix + param_names[9], prefix + param_names[10], prefix + param_names[11]], + split_qkv=self.policy.split_qkv) + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[2], + prefix + param_names[3]) + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[3], + prefix + param_names[12]) + maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', + [prefix + param_names[4], prefix + param_names[5]]) + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) + + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) + + +class InternLMLayerPolicy(TransformerPolicy): + _orig_layer_class = [] + _orig_layer_class_inited = False + + def __init__(self, client_module, inference=True): + super().__init__( + inference, + mlp_act_func_type=ActivationFuncType.GATED_SILU, + norm_type=NormType.RMSNorm, + ) + self.client_module = client_module + + self._init_orig_layer_class_once() + + def _init_orig_layer_class_once(self): + if InternLMLayerPolicy._orig_layer_class_inited: + return + + for sub_pkg in ['', '.internlm-7b', '.internlm-chat-7b']: + try: + from transformers.utils import TRANSFORMERS_DYNAMIC_MODULE_NAME + module = importlib.import_module(f"{TRANSFORMERS_DYNAMIC_MODULE_NAME}{sub_pkg}.modeling_internlm") + if module.InternLMDecoderLayer not in InternLMLayerPolicy._orig_layer_class: + InternLMLayerPolicy._orig_layer_class.append(module.InternLMDecoderLayer) + except ImportError: + continue + + InternLMLayerPolicy._orig_layer_class_inited = True + + def get_hidden_heads(self): + return self.client_module.self_attn.q_proj.weight.shape[1], \ + self.client_module.self_attn.num_heads, \ + self.client_module.input_layernorm.variance_epsilon, \ + self.client_module.mlp.gate_proj.weight.shape[0] + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + kw = self.client_module.self_attn.k_proj.weight + vw = self.client_module.self_attn.v_proj.weight + qb = self.client_module.self_attn.q_proj.bias + kb = self.client_module.self_attn.k_proj.bias + vb = self.client_module.self_attn.v_proj.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + + return qkvw, \ + qkvb, \ + self.client_module.self_attn.o_proj.weight, \ + self.client_module.self_attn.o_proj.bias + + def mlp(self, enable_training=False): + mlp1_up = self.client_module.mlp.up_proj.weight + mlp1_gate = self.client_module.mlp.gate_proj.weight + mlp2 = self.client_module.mlp.down_proj.weight + + mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) + + return mlp1, None, mlp2, None + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + None, \ + self.client_module.input_layernorm.weight, \ + None diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/llama.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..7af333dc1ee4771f52bcc837cf01c847486bdfd7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/llama.py @@ -0,0 +1,164 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer +from deepspeed.utils.types import ActivationFuncType, NormType +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter + +from ..policy import ( + TransformerPolicy, + transformer_param_names, + maybe_copy, + maybe_copy_qkv, + maybe_copy_geglu, + maybe_get_lora, +) + + +class DS_LLAMAContainer(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, + BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + _config.rotate_half = True + _config.rotate_every_two = False + _config.rotary_dim = self.hidden_size // self.num_attention_heads + _config.rope_theta = self.policy.client_module.self_attn.rope_theta + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight, + self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight, + self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight, + self.policy.client_module.self_attn.o_proj.weight + ] + ] + + def get_lora_matched_pair(self): + up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), + (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = None + + def set_mlp_gate(self): + """ + Necessary to implement for `HybridGatedMLPContainer` + """ + self.inter_up_w = self.policy.client_module.mlp.up_proj.weight + self.inter_up_b = None + self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight + self.inter_gate_b = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attn.q_proj.weight', \ + 'self_attn.k_proj.weight', \ + 'self_attn.v_proj.weight', \ + 'self_attn.o_proj.weight', \ + 'mlp.up_proj.weight', \ + 'mlp.gate_proj.weight', \ + 'mlp.down_proj.weight', \ + 'post_attention_layernorm.weight', \ + 'input_layernorm.weight', + ) + + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', + [prefix + param_names[4], prefix + param_names[5]]) + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) + + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) + + # This line is necessary for proper output when kernels + meta tensors are used in Llama models + # TODO: Investigate root-cause and fix meta tensor loading + module.mlp.output_b = None + + +class LLAMALayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=True): + super().__init__( + inference, + mlp_act_func_type=ActivationFuncType.GATED_SILU, + norm_type=NormType.RMSNorm, + ) + self.client_module = client_module + try: + import transformers + LLAMALayerPolicy._orig_layer_class = transformers.models.llama.modeling_llama.LlamaDecoderLayer # type: ignore + except: + LLAMALayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + hidden_heads = ( + self.client_module.self_attn.q_proj.in_features, + self.client_module.self_attn.num_heads, + self.client_module.input_layernorm.variance_epsilon, + self.client_module.mlp.gate_proj.out_features, + ) + return hidden_heads + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + kw = self.client_module.self_attn.k_proj.weight + vw = self.client_module.self_attn.v_proj.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.self_attn.o_proj.weight, \ + None + + def mlp(self, enable_training=False): + mlp1_up = self.client_module.mlp.up_proj.weight + mlp1_gate = self.client_module.mlp.gate_proj.weight + mlp2 = self.client_module.mlp.down_proj.weight + + mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) + + return mlp1, None, mlp2, None + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + None, \ + self.client_module.input_layernorm.weight, \ + None diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/llama2.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/llama2.py new file mode 100644 index 0000000000000000000000000000000000000000..b531890ab859dfbff51a401819846299d8e77c08 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/llama2.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer +from deepspeed.utils.types import ActivationFuncType, NormType +from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference +import torch +from torch.nn.parameter import Parameter + +from ..policy import ( + TransformerPolicy, + transformer_param_names, + maybe_copy, + maybe_copy_qkv, + maybe_copy_geglu, + maybe_get_lora, +) + + +class DS_LLAMA2Container(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, + BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + _config.rotate_half = False + _config.rotate_every_two = True + _config.rotary_dim = self.hidden_size // self.num_attention_heads + _config.num_kv = self.policy.client_module.attention.n_kv_heads + self.module = DeepSpeedLlama2Inference(_config, mp_group=self.mp_group) + + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.feed_forward.w3.weight, self.policy.client_module.feed_forward.w1.weight, + self.policy.client_module.feed_forward.w2.weight, self.policy.client_module.attention.wq.weight, + self.policy.client_module.attention.wk.weight, self.policy.client_module.attention.wv.weight, + self.policy.client_module.attention.wo.weight + ] + ] + + def get_lora_matched_pair(self): + up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), + (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.attention.wq.weight + self.qb = None + self.kw = self.policy.client_module.attention.wk.weight + self.kb = None + self.vw = self.policy.client_module.attention.wv.weight + self.vb = None + + def set_mlp_gate(self): + """ + Necessary to implement for `HybridGatedMLPContainer` + """ + self.inter_up_w = self.policy.client_module.feed_forward.w2.weight + self.inter_up_b = None + self.inter_gate_w = self.policy.client_module.feed_forward.w1.weight + self.inter_gate_b = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'attention.wq.weight', \ + 'attention.wk.weight', \ + 'attention.wv.weight', \ + 'attention.wo.weight', \ + 'feed_forward.w3.weight', \ + 'feed_forward.w1.weight', \ + 'feed_forward.w2.weight', \ + 'ffn_norm.weight', \ + 'attention_norm.weight' + ) + + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', + [prefix + param_names[4], prefix + param_names[5]]) + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) + + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) + + +class LLAMA2LayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=True): + super().__init__( + inference, + mlp_act_func_type=ActivationFuncType.GATED_SILU, + norm_type=NormType.RMSNorm, + ) + self.client_module = client_module + try: + import llama + LLAMA2LayerPolicy._orig_layer_class = llama.model.TransformerBlock # type: ignore + except: + LLAMA2LayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.attention.wq.weight.shape[1], \ + self.client_module.n_heads, \ + self.client_module.ffn_norm.eps, \ + (self.client_module.feed_forward.w1.weight.shape[0] * \ + deepspeed.comm.get_world_size() if deepspeed.comm.is_initialized() else 1) # this is a hack to inject when model is already partitioned! + + def attention(self, enable_training=False): + qw = self.client_module.attention.wq.weight + kw = self.client_module.attention.wk.weight + vw = self.client_module.attention.wv.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.attention.wo.weight, \ + None + + def mlp(self, enable_training=False): + mlp1_up = self.client_module.feed_forward.w3.weight + mlp1_gate = self.client_module.feed_forward.w1.weight + mlp2 = self.client_module.feed_forward.w2.weight + + mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) + + return mlp1, None, mlp2, None + + def layernorm(self): + return self.client_module.ffn_norm.weight, \ + None, \ + self.client_module.attention_norm.weight, \ + None diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/megatron_gpt.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/megatron_gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..2851dd246d99c94235cc6c0abce05f0c73375138 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/megatron_gpt.py @@ -0,0 +1,117 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features.megatron import MegatronContainer +from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference +import torch +from ..policy import TransformerPolicy +from packaging import version as pkg_version + + +class DS_MegatronGPTContainer(MegatronContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + + if self.megatron_v2: + self.module.config.rotate_half = True + self.module.config.rotate_every_two = False + + return self.module + + +# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp +# TODO: Generalize MoE overall goal, expand beyond Megatron +class MegatronLayerPolicy(TransformerPolicy): + _orig_layer_class = None + version = 0 + moe_type = 'standard' + megatron_v2 = True + use_mup = False + + def __init__(self, client_module, inference=True): + super().__init__(inference, megatron_v2=MegatronLayerPolicy.megatron_v2, use_mup=MegatronLayerPolicy.use_mup) + self.client_module = client_module + # we use megatron version to differentiate between the old and new + # megatron-lm source code + if MegatronLayerPolicy._orig_layer_class is None: + if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): + MegatronLayerPolicy._orig_layer_class = None + else: + try: + from megatron.model.transformer import ParallelTransformerLayer + MegatronLayerPolicy._orig_layer_class = ParallelTransformerLayer + MegatronLayerPolicy.version = 1 + except ImportError: + MegatronLayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + if MegatronLayerPolicy.version == 0: + return self.client_module.attention.query_key_value.weight.shape[1], \ + self.client_module.attention.num_attention_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + else: + return self.client_module.self_attention.query_key_value.weight.shape[1], \ + self.client_module.self_attention.num_attention_heads, \ + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + if self.inference: + if MegatronLayerPolicy.version == 0: + attention = self.client_module.attention + else: + attention = self.client_module.self_attention + + return attention.query_key_value.weight, \ + attention.query_key_value.bias, \ + attention.dense.weight, \ + attention.dense.bias + + def mlp(self, moe_type='standard', enable_training=False): + from deepspeed.moe.utils import has_moe_layers + moe, _ = has_moe_layers(self.client_module) + + if moe: + moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \ + self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts + num_experts = len(moe_experts) + if moe_type == 'standard': + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)] + else: + + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \ + self.client_module.mlp.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.mlp.dense_4h_to_h.bias, \ + self.client_module.mlp.coefficient.weight + + else: + return self.client_module.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.dense_4h_to_h.bias + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + self.client_module.post_attention_layernorm.bias, \ + self.client_module.input_layernorm.weight, \ + self.client_module.input_layernorm.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..c4063be05b6c5f232470685a292be51667c027a8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .base_moe import * +from .features.megatron import MegatronContainer +from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference +import torch +from .megatron_gpt import MegatronLayerPolicy +from packaging import version as pkg_version + + +class DS_MegatronGPTMoEContainer(MegatronContainer, BaseTransformerMoEContainer): + + def __init__(self, policy, config, model_config, layer_id): + super().__init__(policy, config, model_config, layer_id) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + + if self.megatron_v2: + self.module.config.rotate_half = True + self.module.config.rotate_every_two = False + + return self.module + + +# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp +# TODO: Generalize MoE overall goal, expand beyond Megatron +class MegatronMoELayerPolicy(MegatronLayerPolicy): + _orig_layer_class = None + version = 0 + moe_type = 'standard' + num_experts = 1 + + def __init__(self, client_module, inference=True): + super().__init__(inference) + self.client_module = client_module + # we use megatron version to differentiate between the old and new + # megatron-lm source code + if MegatronMoELayerPolicy._orig_layer_class is None: + if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): + MegatronMoELayerPolicy._orig_layer_class = None + else: + try: + from megatron.model.transformer import ParallelTransformerLayer + MegatronMoELayerPolicy._orig_layer_class = ParallelTransformerLayer + except ImportError: + MegatronMoELayerPolicy._orig_layer_class = None + + def get_num_experts(self): + return self.num_experts + + def mlp(self, moe_type='standard', enable_training=False): + # for now, all of this is tightly coupled to megatron-deepspeed moe implementation + # todo: think and refactor this to be more general + + #from deepspeed.moe.utils import has_moe_layers + #moe, _ = has_moe_layers(self.client_module) + + moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \ + self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts + num_experts = len(moe_experts) + self.num_experts = num_experts + + if moe_type == 'standard': + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)] + else: + return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ + [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \ + self.client_module.mlp.mlp.dense_h_to_4h.weight, \ + self.client_module.mlp.mlp.dense_h_to_4h.bias, \ + self.client_module.mlp.mlp.dense_4h_to_h.weight, \ + self.client_module.mlp.mlp.dense_4h_to_h.bias, \ + self.client_module.mlp.coefficient.weight diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/opt.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/opt.py new file mode 100644 index 0000000000000000000000000000000000000000..eba619c79d6ebf5325c12f42c921b77dcb6e5ad7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/opt.py @@ -0,0 +1,160 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features import MetaTensorContainer, HybridSplitQKVContainer +from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference +import torch +from torch.nn.parameter import Parameter +from ..policy import TransformerPolicy +from ..policy import transformer_param_names +from ..policy import maybe_copy +from ..policy import maybe_copy_qkv +from ..policy import maybe_get_lora +from deepspeed.utils.types import ActivationFuncType + + +class DS_OPTContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + self.module = DeepSpeedOPTInference(_config, mp_group=self.mp_group) + self.module.config.scale_attention = self.scale_attention + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.fc1, + self.policy.client_module.fc2, + self.policy.client_module.self_attn.q_proj, + self.policy.client_module.self_attn.k_proj, + self.policy.client_module.self_attn.v_proj, + self.policy.client_module.self_attn.out_proj, + ] + ] + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = self.policy.client_module.self_attn.q_proj.bias + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = self.policy.client_module.self_attn.k_proj.bias + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = self.policy.client_module.self_attn.v_proj.bias + + def get_lora_matched_pair(self): + fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() + ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), + (k_lora, self.kw), (v_lora, self.vw)] + return ret + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attn.q_proj.weight', \ + 'self_attn.k_proj.weight', \ + 'self_attn.v_proj.weight', \ + 'self_attn.q_proj.bias', \ + 'self_attn.k_proj.bias', \ + 'self_attn.v_proj.bias', \ + 'self_attn.out_proj.weight', \ + 'self_attn.out_proj.bias', \ + 'fc1.weight', \ + 'fc1.bias', \ + 'fc2.weight', \ + 'fc2.bias', \ + 'final_layer_norm.weight', \ + 'final_layer_norm.bias', \ + 'self_attn_layer_norm.weight', \ + 'self_attn_layer_norm.bias' + ) + + for i in range(0, 6, 3): + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + transformer_param_names[i // 3], + [prefix + param_names[i], prefix + param_names[i + 1], prefix + param_names[i + 2]], + split_qkv=self.policy.split_qkv) + for i in range(6, 8): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], + prefix + param_names[i]) + for i in range(8, 14): + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], + prefix + param_names[i]) + for i in range(14, 16): + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], + prefix + param_names[i]) + + +class HFOPTLayerPolicy(TransformerPolicy): + _orig_layer_class = None + + def __init__(self, client_module, inference=True, use_load_prefix=True): + super().__init__(inference, linear_layer=True, pre_attn_norm=True, use_load_prefix=use_load_prefix) + self.client_module = client_module + try: + import transformers + HFOPTLayerPolicy._orig_layer_class = transformers.models.opt.modeling_opt.OPTDecoderLayer + except: + HFOPTLayerPolicy._orig_layer_class = None + + if hasattr(TransformerPolicy, "hf_model_config") and hasattr(TransformerPolicy.hf_model_config, + "activation_function"): + if TransformerPolicy.hf_model_config.activation_function == "relu": + self.mlp_act_func_type = ActivationFuncType.ReLU + elif TransformerPolicy.hf_model_config.activation_function in ["gelu", "gelu_new"]: + self.mlp_act_func_type = ActivationFuncType.GELU + else: + raise ValueError("Unsupported activation function: {}".format( + TransformerPolicy.hf_model_config.activation_function)) + else: + self.mlp_act_func_type = ActivationFuncType.ReLU # default + + def get_hidden_heads(self): + return self.client_module.self_attn.embed_dim, \ + self.client_module.self_attn.num_heads, \ + self.client_module.self_attn_layer_norm.eps, \ + DEFAULT_INTERMEDIATE_SIZE + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + qb = self.client_module.self_attn.q_proj.bias + + kw = self.client_module.self_attn.k_proj.weight + kb = self.client_module.self_attn.k_proj.bias + + vw = self.client_module.self_attn.v_proj.weight + vb = self.client_module.self_attn.v_proj.bias + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) + return qkvw, \ + qkvb, \ + self.client_module.self_attn.out_proj.weight, \ + self.client_module.self_attn.out_proj.bias + + def mlp(self, enable_training=False): + return self.client_module.fc1.weight, \ + self.client_module.fc1.bias, \ + self.client_module.fc2.weight, \ + self.client_module.fc2.bias + + def layernorm(self): + return self.client_module.final_layer_norm.weight, \ + self.client_module.final_layer_norm.bias, \ + self.client_module.self_attn_layer_norm.weight, \ + self.client_module.self_attn_layer_norm.bias diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/unet.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..48179265553150824855140750da3fc6e41503ef --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/unet.py @@ -0,0 +1,56 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from torch.nn.parameter import Parameter + +from ..policy import DSPolicy +from ...model_implementations.diffusers.unet import DSUNet + + +class UNetPolicy(DSPolicy): + + def __init__(self): + super().__init__() + try: + import diffusers + self._orig_layer_class = diffusers.models.unet_2d_condition.UNet2DConditionModel + except AttributeError: + self._orig_layer_class = diffusers.models.unets.unet_2d_condition.UNet2DConditionModel + except ImportError: + self._orig_layer_class = None + + def match(self, module): + return isinstance(module, self._orig_layer_class) + + def match_replaced(self, module): + return isinstance(module, DSUNet) + + def apply(self, module, enable_cuda_graph=True): + # TODO(cmikeh2): Enable cuda graph should be an inference configuration + return DSUNet(module, enable_cuda_graph=enable_cuda_graph) + + def attention(self, client_module): + qw = client_module.to_q.weight + kw = client_module.to_k.weight + vw = client_module.to_v.weight + + if qw.shape[1] == kw.shape[1]: + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=False) + + return qkvw, \ + client_module.to_out[0].weight, \ + client_module.to_out[0].bias, \ + qw.shape[-1], \ + client_module.heads + else: + #return None + #kvw = Parameter(torch.cat((kw, vw), dim=0), requires_grad=False) + return qw, \ + kw, vw, \ + client_module.to_out[0].weight, \ + client_module.to_out[0].bias, \ + qw.shape[-1], \ + client_module.heads diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/containers/vae.py b/lib/python3.12/site-packages/deepspeed/module_inject/containers/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..d26d0ef77ca9d10e27bad47859dba589190af7b8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/containers/vae.py @@ -0,0 +1,42 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..policy import DSPolicy +from ...model_implementations.diffusers.vae import DSVAE + + +class VAEPolicy(DSPolicy): + + def __init__(self): + super().__init__() + try: + import diffusers + if hasattr(diffusers.models, "autoencoders"): + # Diffusers >= 0.25.0 + # Changes location to 'autoencoders' directory + self._orig_layer_class = diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL + elif hasattr(diffusers.models.vae, "AutoencoderKL"): + # Diffusers < 0.12.0 + self._orig_layer_class = diffusers.models.vae.AutoencoderKL + else: + # Diffusers >= 0.12.0 & < 0.25.0 + # Changes location of AutoencoderKL + self._orig_layer_class = diffusers.models.autoencoder_kl.AutoencoderKL + except ImportError: + self._orig_layer_class = None + + def match(self, module): + return isinstance(module, self._orig_layer_class) + + def match_replaced(self, module): + return isinstance(module, DSVAE) + + def apply(self, module, enable_cuda_graph=True): + # TODO(cmikeh2): Enable cuda graph should be an inference configuration + return DSVAE(module, enable_cuda_graph=enable_cuda_graph) + + # NOTE (lekurile): Should we have a diffusers policy class? + def attention(self, client_module): + pass diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/fusedqkv_utils.py b/lib/python3.12/site-packages/deepspeed/module_inject/fusedqkv_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0609c6001dd2f716b3c1953500369dc9e67a4987 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/fusedqkv_utils.py @@ -0,0 +1,231 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +import torch +from deepspeed.utils.logging import warning_once +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, get_num_kv_heads, get_n_embd, get_num_attention_heads + + +def split_by_qkvlist_and_refuse(qkv_list, split_size, split_dim=0, cat_dim=0): + qkv_split_list = [torch.split(mat, split_size, dim=split_dim) for mat in qkv_list] + tp_fusedqkv_list = [ + torch.cat([qkv_s[i] for qkv_s in qkv_split_list], dim=cat_dim) for i in range(len(qkv_split_list[0])) + ] + return tp_fusedqkv_list + + +def require_tp_fused_qkvw(name, mp_size): + fused_qkvw_name_list = ['qkv_proj', 'query_key_value', 'attn.Wqkv', 'self_attn.W_pack', 'c_attn'] + + if mp_size == 1: + return False + for fused_name in fused_qkvw_name_list: + if fused_name in name: + return True + return False + + +def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index): + + module_str = str(module).strip() + if src is None: + return + fused_type_dict = { + 'CodeGenBlock': 'codegentype', + 'BloomBlock': 'bloomtype', + 'GLMBlock': 'glmtype', + "MPTBlock": 'glmtype', + "MptBlock": 'glmtype', + "BaichuanLayer": 'glmtype', + "QWenBlock": 'qwentype', + "FalconDecoderLayer": 'bloomtype', + "GPTBigCodeBlock": 'bigcodetype', + "DecoderLayer": 'glmtype', + "Phi3DecoderLayer": "phi3type" + } + + def _codegen_type_transpose(input, mp_size, codegen_mp_num=4): + # codegen_mp_num defined in https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py + assert get_num_kv_heads() % ( + mp_size * codegen_mp_num) == 0, "codgen autoTP requires num_kv_heads % (mp_size*codegen_mp_num) == 0" + #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias) + + shape = input.shape + dst_shape = get_shard_size(shape[0], mp_size) + num_mp_blocks = input.reshape(codegen_mp_num, shape[0] // codegen_mp_num, shape[1]) + + #num_mp_blocks : [codegen_mp_num, 3*hidden_dim/codegen_mp_num, :] + src_split = list(torch.split(num_mp_blocks, num_mp_blocks.shape[1] // 3, dim=1)) + src_split = [x.reshape(codegen_mp_num * mp_size, -1, shape[1]) for x in src_split] + + split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size(shape[0] // 3, mp_size), 0, 1) + tp_fuseqkv_weight = torch.cat(split_fusedqkv, dim=0).reshape(shape[0], -1) + + return tp_fuseqkv_weight[gpu_index * dst_shape:(gpu_index + 1) * dst_shape] + + def _glm_type_transpose(input, mp_size): + #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias) + + # For chatglm2 & chatglm3(kv_heads=2), need to special handle. + if get_num_kv_heads() == 2: + shape = input.shape + hidden_dim = get_n_embd() + kv_dim = (shape[0] - hidden_dim) // get_num_kv_heads() + q = input[:hidden_dim] + k = input[hidden_dim:hidden_dim + kv_dim] + v = input[hidden_dim + kv_dim:] + q_split = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0) + k_split = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0) + v_split = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0) + return torch.cat((q_split[gpu_index], k_split[gpu_index], v_split[gpu_index]), dim=0) + else: + shape = input.shape + src_split = torch.split(input, shape[0] // 3, dim=0) + + split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size)) + return split_fusedqkv[gpu_index] + + def _bloom_type_transpose(input, mp_size): + shape = input.shape + + split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size), dim=0) + return split_fusedqkv[gpu_index] + + def _qwen_type_transpose(input, mp_size, module): + if not hasattr(module, "_ds_fusedqkv_entered"): + # Adjust splitting absolute value variables + setattr(module, "_ds_fusedqkv_entered", True) + module.attn.split_size = get_shard_size(module.attn.split_size, mp_size) + return _glm_type_transpose(input, mp_size) + + def _bigcode_type_transpose(input, mp_size): + n_embd = get_n_embd() + q = input[:n_embd] + kv = input[n_embd:] + shape = q.shape + split_q = q.split(get_shard_size_list(shape[0], mp_size), dim=0) + return torch.cat((split_q[gpu_index], kv), dim=0) + + def _phi3_type_transpose(input, mp_size): + num_kv_heads = get_num_kv_heads() + num_heads = get_num_attention_heads() + hidden_size = input.shape[1] + head_dim = hidden_size // num_heads + q_pos = input.shape[0] - 2 * num_kv_heads * head_dim + q = input[:q_pos] + k = input[q_pos:q_pos + num_kv_heads * head_dim] + v = input[q_pos + num_kv_heads * head_dim:] + split_q = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0) + split_k = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0) + split_v = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0) + return torch.cat((split_q[gpu_index], split_k[gpu_index], split_v[gpu_index]), dim=0) + + def _transpose_fused_qkvw(src, mp_size, fused_qkv_type=None, module=None): + + # suppose num_heads=n, q(n)_w means the n-th q head linear weight, the weight format are as following + # bloomtype: [q(1)_w,k(1)_w,v(1)_w,q(2)_w,k(2)_w,v(2)_w,...,q(n)_w,k(n)_w,v(n)_w] + # glmtype: [q(1)_w, q(2)_w,...,q(n)_w,k(1)_w,k(2)_w,...,k(n)_w,v(1)_w,v(2)_w,...,v(n)_w] + # codegentype: [q(1)_w,q(2)_w,...,q(n/t)_w,k(1)_w,k(2)_w,...,k(n/t)_w,v(1)_2,v(2)_w,...v(n/t)_w,q(n/t+1)_w,...], where t is a const defined in model file. + + if fused_qkv_type == 'bloomtype': + return _bloom_type_transpose(src, mp_size) + elif fused_qkv_type == 'codegentype': + return _codegen_type_transpose(src, mp_size) + elif fused_qkv_type == 'glmtype': + return _glm_type_transpose(src, mp_size) + elif fused_qkv_type == 'qwentype': + return _qwen_type_transpose(src, mp_size, module) + elif fused_qkv_type == 'bigcodetype': + return _bigcode_type_transpose(src, mp_size) + elif fused_qkv_type == 'phi3type': + return _phi3_type_transpose(src, mp_size) + + raise ValueError("unknown fused_qkv_type") + + module_name_matches = [k for k in fused_type_dict.keys() if k in module_str] + if module_name_matches: + # There can be overlap with matches (e.g., "DecoderLayer" and "FalconDecoderLayer"). + # We take the longest matching module_name + module_name = max(module_name_matches, key=len) + fused_type = fused_type_dict[module_name] + return _transpose_fused_qkvw(src, mp_size, fused_type, module) + warning_once(f"Unrecognized fusedkqv weight type, default to using bloom type," + f"please check in prepare_tp_fused_qkvw() to avoid potential calculation errors") + return _bloom_type_transpose(src, mp_size) + + +# For share qk type: +# q = [q1,...,q_{n/4}, q_{n/2+1},...,q_{3n/4}, k1,...,k_{n/4}, k_{n/2+1},...,k_{3n/4}] +# k = [q_{n/4+1},...,q_{n/2}, q_{3n/4+1},...,qn, k_{n/4+1},...,k_{n/2}, k{3n/4+1},...,kn] +# Avoid modifying the modeling code. We adjust the value and oproj weight to fit this qk type. +def shard_value_with_share_qk( + weight, + bias, + rank, + world_size, + shard_value=True # True -> shard_value; False -> shard_oproj +): + if shard_value: + total_size = weight.shape[0] + weight_cat_dim = 0 + else: + total_size = weight.shape[1] + weight_cat_dim = 1 + num_heads = get_num_kv_heads() + head_dim = total_size // num_heads + assert (num_heads % world_size == 0) + if world_size > num_heads // 2: + RuntimeError(f"world_size {world_size} is larger than half of num_heads {num_heads}") + head_per_rank = num_heads // world_size + q_head_start = rank * head_per_rank + # mapping q_head to v_head + v_head_ids = [] + i = 0 + # mapping neighbor q_head to v_head + while i < head_per_rank: + v_head_ids.append(q_head_start // 2) + q_head_start += 2 + i = i + 2 + + # mapping neighbor k_head to v_head + v_head_ids.extend([i + num_heads // 2 for i in v_head_ids]) + sharded_weight = [] + sharded_bias = [] + for head_id in v_head_ids: + if shard_value: + sharded_weight.append(weight[head_id * head_dim:(head_id + 1) * head_dim]) + if bias is not None: + sharded_bias.append(bias.data[head_id * head_dim:(head_id + 1) * head_dim]) + else: + sharded_weight.append(weight[:, head_id * head_dim:(head_id + 1) * head_dim]) + sharded_weight = torch.cat(sharded_weight, dim=weight_cat_dim) + if bias is not None: + if shard_value: + sharded_bias = torch.cat(sharded_bias, dim=0) + else: + bias = bias / float(world_size) + return torch.nn.Parameter(sharded_weight), torch.nn.Parameter(sharded_bias) + else: + return torch.nn.Parameter(sharded_weight), None + + +# For phi3 with chunk mlp, adjust the weight order. +def shard_chunk_mlp( + weight, + bias, + rank, + world_size, +): + weight_gate, weight_states = weight.chunk(2, dim=0) + total_size = weight_gate.shape[0] + split_weight_gate = weight_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + split_weight_states = weight_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + shard_weight = torch.cat((split_weight_gate[rank], split_weight_states[rank]), dim=0) + if bias is not None: + bias_gate, bias_states = bias.chunk(2, dim=0) + split_bias_gate = bias_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + split_bias_states = bias_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) + return shard_weight, torch.cat((split_bias_gate[rank], split_bias_states[rank]), dim=0) + + return shard_weight, None diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/inject.py b/lib/python3.12/site-packages/deepspeed/module_inject/inject.py new file mode 100644 index 0000000000000000000000000000000000000000..401da1bd6ef76b268f65e48b77c30a8ed65bfc64 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/inject.py @@ -0,0 +1,112 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import copy +import torch +from deepspeed.ops.transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig + + +def module_inject(layer_obj, model, config, micro_batch_size, max_seq_length, seed, preln, fp16=True): + for name, child in model.named_children(): + if isinstance(child, layer_obj): + print('REPLACING BertLayer') + + cuda_config = DeepSpeedTransformerConfig(batch_size=micro_batch_size, + max_seq_length=max_seq_length, + hidden_size=config.hidden_size, + heads=config.num_attention_heads, + attn_dropout_ratio=config.attention_probs_dropout_prob, + hidden_dropout_ratio=config.hidden_dropout_prob, + num_hidden_layers=config.num_hidden_layers, + initializer_range=config.initializer_range, + seed=seed, + fp16=fp16, + pre_layer_norm=preln) + + new_module = DeepSpeedTransformerLayer(cuda_config) + + # copy relevant state from child -> new module + qw = child.attention.self.query.weight + qb = child.attention.self.query.bias + kw = child.attention.self.key.weight + kb = child.attention.self.key.bias + vw = child.attention.self.value.weight + vb = child.attention.self.value.bias + + qkvw = torch.cat((qw, kw, vw), 0) + qkvb = torch.cat((qb, kb, vb), 0) + + new_module.attn_qkvw.data = qkvw + new_module.attn_qkvb.data = qkvb + new_module.attn_ow.data = child.attention.output.dense.weight + new_module.attn_ob.data = child.attention.output.dense.bias + if preln: + attention_layerNorm = child.PostAttentionLayerNorm + else: + attention_layerNorm = child.attention.output.LayerNorm + new_module.attn_nw.data = attention_layerNorm.weight + new_module.attn_nb.data = attention_layerNorm.bias + if preln: + intermediate_FF = child.intermediate.dense_act + else: + intermediate_FF = child.intermediate.dense + new_module.inter_w.data = intermediate_FF.weight + new_module.inter_b.data = intermediate_FF.bias + new_module.output_w.data = child.output.dense.weight + new_module.output_b.data = child.output.dense.bias + if preln: + transformer_LayerNorm = child.PreAttentionLayerNorm + else: + transformer_LayerNorm = child.output.LayerNorm + new_module.norm_w.data = transformer_LayerNorm.weight + new_module.norm_b.data = transformer_LayerNorm.bias + + setattr(model, name, copy.deepcopy(new_module)) + + else: + module_inject(layer_obj, child, config, micro_batch_size, max_seq_length, seed, preln, fp16) + + return model + + +def test_hi(): + from turing.nvidia_modelingpreln import BertConfig as BertConfigPreLN + from turing.nvidia_modelingpreln import BertForQuestionAnswering as BertForQuestionAnsweringPreLN + from turing.nvidia_modelingpreln import BertLayer + bert_model_config = { + "vocab_size_or_config_json_file": 119547, + "hidden_size": 1024, + "num_hidden_layers": 1, + "num_attention_heads": 16, + "intermediate_size": 4096, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.1, + "attention_probs_dropout_prob": 0.1, + "hidden_dropout_prob": 0.1, + "attention_probs_dropout_prob": 0.1, + "max_position_embeddings": 512, + "type_vocab_size": 2, + "initializer_range": 0.02 + } + bert_config = BertConfigPreLN(**bert_model_config) + base_model = BertForQuestionAnsweringPreLN(bert_config, args=None) + + #base_model = LinearStack() + + test_model = copy.deepcopy(base_model) + test_model = module_inject(BertLayer, test_model, bert_config, 4, 384, 1234) + + print('BASE', base_model) + print('TEST', test_model) + + #base_model.eval() + #test_model.eval() + + #test_input = torch.rand(1, base_model.input_dim) + + #base_output = base_model(test_input) + #test_output = test_model(test_input) + # + #assert torch.allclose(base_output, test_output, atol=3e-8) diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/layers.py b/lib/python3.12/site-packages/deepspeed/module_inject/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..8aa45db2b137d46e1fab62bb61d9585b5f9b56ab --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/layers.py @@ -0,0 +1,818 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed import comm as dist +from torch import nn +from torch.nn import functional as F +from torch.nn.parameter import Parameter +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list +from deepspeed.runtime.zero.utils import is_zero_param +from abc import ABC, abstractmethod +from typing import Iterable, Any, Optional, List, Tuple +from .fusedqkv_utils import shard_value_with_share_qk, shard_chunk_mlp, prepare_tp_fused_qkvw +from deepspeed.runtime.tensor_parallel import AUTOTP_MODE +from copy import deepcopy +from typing import Union + +__all__ = [ + "TensorParallel_Layer", "LinearAllreduce", "LinearLayer", "LmHeadLinearAllreduce", "Yuan_LinearAllreduce", + "Yuan_LinearLayer", "GateUpPack_LinearLayer", "Conv_LinearALlreduce", "fused_LinearLayer", "conv_LinearLayer" +] + +DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE +DS_IS_REPLACED_MODULE = 'ds_is_replaced_module' +DS_TENSOR_MODEL_PARALLEL = 'tensor_model_parallel' + + +def get_auto_tp_mode(): + global DEEPSPEED_AUTOTP_MODE + return DEEPSPEED_AUTOTP_MODE + + +def is_autotp_training_mode(): + global DEEPSPEED_AUTOTP_MODE + return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING + + +def set_autotp_mode(training=False): + """ + Set the DEEPSPEED_AUTOTP_MODE based on the training flag + """ + global DEEPSPEED_AUTOTP_MODE + if training: + DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.TRAINING + else: + DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE + + +def add_bias(input, bias): + if bias is None: + return input + if is_autotp_training_mode(): + # Training mode - avoid inplace to ensure correct autograd + input = input + bias + return input + else: + input += bias + return input + + +class RowParallel(torch.autograd.Function): + """ + A custom autograd function for performing row-wise parallelism. + """ + + @staticmethod + def symbolic(graph, input): + """Symbolic function for tracing.""" + return input + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, is_inference_mode: bool) -> torch.Tensor: + """ + Forward pass. + """ + ctx.group = group + if group == None: + return input + if is_inference_mode: + dist.inference_all_reduce(input, group=group) + else: + dist.all_reduce(input.contiguous(), group=group) + return input + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None]: + """ + Backward pass. + """ + return None, grad_output, None + + +class AsyncColumnParallel(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, weight, bias) -> torch.Tensor: + """ + Forward pass. + """ + ctx.use_bias = bias is not None + ctx.group = group + output = torch.matmul(input, weight.transpose(-1, -2)) + if bias is not None: + output = add_bias(output, bias) + + ctx.save_for_backward(input, weight) + + return output + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]: + + input, weight = ctx.saved_tensors + grad_input = grad_output.matmul(weight) + handle = dist.all_reduce(grad_input.contiguous(), group=ctx.group, async_op=True) + grad_weight = grad_output.view(-1, grad_output.shape[-1]).t().matmul(input.view(-1, input.shape[-1])) + grad_bias = grad_output.sum(0) if ctx.use_bias else None + handle.wait() + return None, grad_input, grad_weight, grad_bias + + +class ColumnParallel(torch.autograd.Function): + """ + Custom autograd function for column-wise parallelism. + """ + + @staticmethod + def symbolic(graph, input): + """Symbolic function for tracing.""" + return dist.all_reduce(input.contiguous(), dist.get_tensor_model_parallel_group()) + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Tensor: + """ + Forward pass. + """ + ctx.group = group + return input + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]: + """ + Backward pass. + """ + if ctx.group == None: + return None, grad_output + + dist.all_reduce(grad_output.contiguous(), group=ctx.group) + return None, grad_output + + +class TensorParallel_Layer(nn.Module, ABC): + """ + A base class for model layers with tensor parallelism support. + This class is designed to be extended by specific layers that require distributed + operations and parameter gather/partitioning during inference or training. + + Attributes: + mode (str): The mode of operation[INFERENCE or TRAINING], default is "INFERENCE". + mp_group (Optional[dist.ProcessGroup]): The process group used for model parallelism. + tp_world_size (int): The world size of tensor parallelism, i.e., the number of parallel workers. + tp_index (int): The rank (ID) of the current worker in tensor parallelism. + support_training (bool): Flag indicating whether the layer supports training (default: False). + name (Optional[str]): The name of the layer, if provided. + """ + ##### Initialize Parameter List ##### + + # keep_module_on_host determines whether to keep the module on the host. + # Checkpoints are first loaded to the host (sometimes directly from disk to avoid filling host memory), + # so an additional copy is unnecessary. + keep_module_on_host: bool = False + + ##### Runtime Parameter List ##### + tp_overlap_comm: bool = False + """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """ + + def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any): + """ + Initializes the TensorParallel_Layer with optional model parallelism group and layer name. + + Args: + mp_group (Optional[dist.ProcessGroup]): The process group for model parallelism. + If None, no model parallelism is set. + """ + super().__init__() + self.support_training: bool = False + if mp_group is not None: + self.mp_group = mp_group + self.tp_world_size: int = dist.get_world_size(self.mp_group) + self.tp_index: int = dist.get_rank(mp_group) + + # backward compatibility + self.world_size = self.tp_world_size + self.rank = self.tp_index + + self.name = getattr(self, 'name', None) + if kwargs.get('name') is not None: + self.name = kwargs.get('name') # Set the layer name if provided. + + @classmethod + def set_keep_module_on_host(cls, value: bool): + """ + Set the static variable keep_module_on_host. + + Args: + value (bool): The new value for keep_module_on_host. + """ + cls.keep_module_on_host = value + + @abstractmethod + def forward(self, input): + """ + Forward pass method. Must be implemented by subclasses to define layer-specific operations. + """ + pass + + @abstractmethod + def gather_params(self, params_list): + """ + Gathers parameters across devices for distributed training. Must be implemented by subclasses in "TRAINING" mode. + """ + pass + + @abstractmethod + def _tp_partition(self, params_list: List[torch.Tensor]): + """ + Partitions the parameters for tensor parallelism. + It is necessary to ensure that this function only involves the logic of params partitioning. + """ + pass + + def config_requires_grad(self, weight): + if weight is not None: + if self.is_training_mode(): + if weight.requires_grad is None: + weight.requires_grad = True + else: + weight.requires_grad = False + + def config_tp_params(self, weight): + """ + Configures the weight tensor for training with tensor parallelism. This includes enabling gradients + and associating necessary methods for parameter gathering and partitioning. + + Args: + weight (Optional[torch.Tensor]): The weight tensor to configure for tensor parallelism. + If None, no action is taken. + """ + # # The RNG states have already been synchronized in init_inference. + if self.is_training_mode(): + assert self.support_training, "No implementation of backward." + if weight is not None: + self.config_requires_grad(weight) + weight.gather_params = self.gather_params + weight._tp_partition = self._tp_partition + setattr(weight, DS_TENSOR_MODEL_PARALLEL, True) + setattr(weight, DS_IS_REPLACED_MODULE, True) + + def is_training_mode(self): + global DEEPSPEED_AUTOTP_MODE + return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING + + def __deepcopy__(self, memo): + # This function is designed for + # 'mp_group' (a 'ProcessGroup') cannot be pickled during deepcopy in some usage. + cls = self.__class__ + new_obj = cls.__new__(cls) + + for key, value in vars(self).items(): + if key == 'mp_group': + new_obj.mp_group = self.mp_group + else: + setattr(new_obj, key, deepcopy(value, memo)) + + memo[id(self)] = new_obj + return new_obj + + def extra_repr(self): + out_features, in_features = None, None + if self.weight is not None: + out_features, in_features = self.weight.ds_shape[-2:] if is_zero_param( + self.weight) else self.weight.shape[-2:] + dtype = self.weight.dtype if self.weight is not None else None + return "in_features={}, out_features={}, bias={}, dtype={}".format(in_features, out_features, self.bias + is not None, dtype) + + def move(self, tensor): + # TODO: consider the timing of deletion + # to save host resources when DP > 1。 + + # keep_module_on_host is used to keep the module on the host. Checkpoints are loaded to the host first (in some + # cases it can be done from the disk even to prevent filling host's memory), thus no need to create a new copy. + if tensor.is_meta: + # Keep tensor in meta device if tensor is meta. + return tensor + else: + device = 'cpu' if self.__class__.keep_module_on_host else get_accelerator().current_device_name() + return_new_copy = not self.__class__.keep_module_on_host + + # Using new tensors help in freeing memory (after split for example) was done before by calling clone(). + # Using copy=True instead of clone() will help in case of cpu --> cpu. + # Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced. + cloned_tensor = tensor.to(device, copy=return_new_copy) + + if return_new_copy: + # free the memory of the original tensor to reduce memory peak + # Equivalent to directly deleting the tensor reference outside the function. + # see https://github.com/microsoft/DeepSpeed/pull/4353 + tensor.data = torch.empty(0, device=tensor.device) + return cloned_tensor + + +def configure_tensor_parallel_runtime(config): + runtime_keys = ['tp_overlap_comm'] + for key in runtime_keys: + if hasattr(config, key): + setattr(TensorParallel_Layer, key, getattr(config, key)) + + +class GatherReplacedLayerParams: + """ + A context manager for gathering parameters of a replaced layer, enabling partitioning and gathering functionality + based on the configuration of the model. + """ + + def __init__(self, + params: Union[Iterable[torch.Tensor], torch.Tensor], + module: torch.nn.Module, + enabled: bool = True): + """ + Initialize the context manager to handle parameter gathering and partitioning for a replaced layer. + + Args: + params (Iterable or torch.Tensor): A collection or single parameter to manage. + module (torch.nn.Module): The module that these parameters belong to. + enabled (bool): Flag indicating whether the parameter management is enabled (default: True). + """ + self.enabled = enabled + self.module = module + if not enabled: + return + + # Ensure params is a list, whether it's a single param or iterable (e.g., model.parameters()) + if isinstance(params, Iterable) and not isinstance(params, torch.Tensor): + self.params: List[torch.Tensor] = list(params) # Convert generators to a list for multiple iterations + else: + self.params: List[torch.Tensor] = [params] # Wrap single parameter in a list for uniform processing + + # Check if the parameters belong to a replaced layer (indicated by a specific attribute) + if not any(self._is_replaced_module_weight(p) for p in params): + self.enabled = False + return + + def _is_replaced_module_weight(self, param: torch.Tensor) -> bool: + """ + Helper function to determine if a parameter belongs to a replaced module. + + Args: + param (torch.Tensor): The parameter to check. + + Returns: + bool: True if the parameter belongs to a replaced module, False otherwise. + """ + return getattr(param, DS_IS_REPLACED_MODULE, False) + + def __enter__(self) -> None: + """ + Enter the context manager. If enabled, gather parameters for the replaced module. + """ + if self.enabled: + self.params[0].gather_params(self.params) + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """ + Exit the context manager. If enabled, partition the parameters for the replaced module. + """ + #TODO : Check whether there are any missing attributes. + if self.enabled: + self.params[0]._tp_partition(self.params) + + +class LinearAllreduce(TensorParallel_Layer): + + def __init__(self, module, mp_group, **kwargs): + super(LinearAllreduce, self).__init__(mp_group, **kwargs) + self.weight = module.weight + self.bias = module.bias + + self._tp_partition([self.weight, self.bias]) + self.support_training = True + self.config_tp_params(self.weight) + if self.bias is not None: + # bias here is not tp params + self.config_requires_grad(self.bias) + + def forward(self, input): + output = torch.matmul(input, self.weight.transpose(-1, -2)) + output = RowParallel.apply(self.mp_group, output, not self.is_training_mode()) + if self.bias is not None: + output = add_bias(output, self.bias) + return output + + @torch.no_grad() + def gather_params(self, params_list): + + for idx, param in enumerate(params_list): + if param is None or idx > 0: + # don't gather bias + return + params_list[idx].data_partition = param.data + param = param.transpose(0, 1).contiguous() + + output_param = torch.empty(self.tp_world_size * param.shape[0], + param.shape[1], + dtype=param.dtype, + device=param.device) + dist.all_gather_into_tensor(output_param, param, group=self.mp_group) + params_list[idx].data = output_param.transpose(0, 1).contiguous() + return + + @torch.no_grad() + def _tp_partition(self, params_list): + + if not self.is_training_mode(): + self.uneven_partition(params_list) + return + + else: + for idx, param in enumerate(params_list): + if param is None: + # don't slipt bias + return + if idx > 0: # move bias to device at initialization + _partition = self.move(param).detach() + params_list[idx].data = _partition + return + + _partition = torch.chunk(param, self.tp_world_size, dim=-1)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + def uneven_partition(self, params_list): + for idx, param in enumerate(params_list): + if param is None or idx > 0: + # don't slipt bias + return + assert self.name is not None, "The module name must be provided in the initialization." + _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[1], self.tp_world_size, + self.name), + dim=1)[self.tp_index] + + _partition = self.move(_partition).detach() + params_list[idx].data = _partition + + +#remove kwargs from partition. +class LinearLayer(TensorParallel_Layer): + + def __init__(self, module, mp_group=None, skip_partition=False, **kwargs): + super(LinearLayer, self).__init__(mp_group, **kwargs) + self.weight = module.weight + self.bias = module.bias + if not skip_partition: + self._tp_partition([self.weight, self.bias]) + self.support_training = True + self.config_tp_params(self.weight) + if self.bias is not None: + self.config_tp_params(self.bias) + + def forward(self, input): + if not self.__class__.tp_overlap_comm: + if getattr(self, 'mp_group', None) is not None: + input = ColumnParallel.apply(self.mp_group, input) + output = torch.matmul(input, self.weight.transpose(-1, -2)) + if self.bias is not None: + output = add_bias(output, self.bias) + else: + output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias) + + return output + + @torch.no_grad() + def gather_params(self, params_list): + # Does not support uneven shard. + for idx, param in enumerate(params_list): + + params_list[idx].data_partition = param.data + output_param = torch.empty((self.tp_world_size * param.shape[0], *param.shape[1:]), + dtype=param.dtype, + device=param.device) + dist.all_gather_into_tensor(output_param, param, group=self.mp_group) + params_list[idx].data = output_param.contiguous() + + @torch.no_grad() + def _tp_partition(self, params_list): + + if not self.is_training_mode(): + self.uneven_partition(params_list) + return + for idx, param in enumerate(params_list): + if param is None: + return + #split bias if provide + _partition = torch.chunk(param, self.tp_world_size, dim=0)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + def uneven_partition(self, params_list): + + for idx, param in enumerate(params_list): + if param is None: + #split bias if provide + return + assert self.name is not None, "The module name must be provided in the initialization." + _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[0], self.tp_world_size, + self.name), + dim=0)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + # for bwc + @classmethod + def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=None): + if weight is not None: + in_features = weight.shape[1] + out_features = weight.shape[0] + linear = nn.Linear(in_features, out_features, bias=(bias is not None)) + linear.weight.data = weight + if bias is not None: + linear.bias.data = bias + else: + in_features = weight_shape[1] + out_features = weight_shape[0] + linear = nn.Linear(in_features, out_features, bias=(bias is not None)) + return cls(linear, skip_partition=True) + + +class FusedModuleWrapper: + + def __init__(self, fused_module: nn.Module): + self.fused_module = fused_module + + def __getattr__(self, module): + return self.fused_module + + +class fused_LinearLayer(LinearLayer): + + def __init__(self, module, mp_group, skip_partition=False, **kwargs): + assert kwargs.get('fused_module') is not None, "'fused_module' is required but not provided" + # Use the warp class to avoid module circular references. + self.fused_module = FusedModuleWrapper(kwargs.get('fused_module')) + super().__init__(module, mp_group, skip_partition, **kwargs) + + @torch.no_grad() + def _tp_partition(self, params_list): + for idx, param in enumerate(params_list): + if param is None: + return + + _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index) + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + +class conv_LinearLayer(LinearLayer): + + @torch.no_grad() + def _tp_partition(self, params_list): + weight = None + bias = None + if len(params_list) == 1: + weight = params_list[0] + elif len(params_list) == 2: + weight, bias = params_list[0], params_list[1] + _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.name), + dim=1)[self.tp_index] + _partition = self.move(_partition).detach() + weight.data = _partition + + if bias is not None: + _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.name), + dim=0)[self.tp_index] + _partition = self.move(_partition).detach() + + bias.data = _partition + + +#override the subclasses related to weight splitting. +class Yuan_LinearAllreduce(LinearAllreduce): + + #Yuan2 + @torch.no_grad() + def _tp_partition(self, params_list): + weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index, + self.tp_world_size, False) + params_list[0].data = weight + if bias is not None: + params_list[1].data = bias + + +class Yuan_LinearLayer(LinearLayer): + #Yuan2 + @torch.no_grad() + def _tp_partition(self, params_list): + weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index, + self.tp_world_size, True) + params_list[0].data = self.move(weight).detach() + if bias is not None: + params_list[1].data = self.move(bias).detach() + + +class GateUpPack_LinearLayer(LinearLayer): + # chatGLM2, chatGLM2 + @torch.no_grad() + def _tp_partition(self, params_list): + weight, bias = shard_chunk_mlp(params_list[0].data, params_list[1], self.tp_index, self.tp_world_size) + params_list[0].data = self.move(weight).detach() + if bias is not None: + params_list[1].data = self.move(bias).detach() + + +class Conv_LinearALlreduce(LinearAllreduce): + + @torch.no_grad() + def _tp_partition(self, params_list): + for idx, param in enumerate(params_list): + if param is None: + return + param.data = param.data.transpose(-1, -2).contiguous() + + _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.name), + dim=1)[self.tp_index] + + _partition = self.move(_partition).detach() + + params_list[idx].data = _partition + + +#override the subclasses related to fwd/bwd. +class LmHeadLinearAllreduce(LinearAllreduce): + + def __init__(self, module, mp_group, **kwargs): + # set the fixed name before partition + self.name = "lm_head" + + # In some tied_embedding cases, only the lm head is sharded, while the word embedding is not. + # Reinitialization is used to decouple them and prevent the word embedding from being sharded. + # This should also be effective for cases where both are sharded in tied_embedding scenarios. + + # TODO: Training scenario-related tests, is it necessary to re-implement the vocab parallel module? + module.weight = nn.Parameter(module.weight.clone().detach()) + if hasattr(module, 'bias') and module.bias is not None: + module.bias = nn.Parameter(module.bias.clone().detach()) + super().__init__(module, mp_group, **kwargs) + + def forward(self, input): + input_shard_size = get_shard_size(input.shape[-1], self.tp_world_size, "lm_head") + input_shard_offset = sum(get_shard_size_list(input.shape[-1], self.tp_world_size, "lm_head")[0:self.tp_index]) + output = torch.matmul(input[:, :, input_shard_offset:input_shard_offset + input_shard_size], + self.weight.transpose(-1, -2)) + if self.mp_group is not None: + dist.inference_all_reduce(output, group=self.mp_group) + if self.bias is not None: + output = add_bias(output, self.bias) + return output + + +class TensorParallelConv2d(nn.Module): + + def __init__(self, conv, rank, world_size, shard_by_oc): + super().__init__() + self.rank = rank + self.world_size = world_size + self.shard_by_oc = shard_by_oc + self.shard_weights(conv) + + # Split along the input/output channel depending on whether it is the last conv layer. + def shard_weights(self, conv): + if self.shard_by_oc: + total_size = conv.weight.shape[0] + else: + total_size = conv.weight.shape[1] + bias_data = None + cols_per_rank = [0] + for i in range(self.world_size - 1, -1, -1): + cols = total_size // self.world_size + if i < total_size % self.world_size: + cols += 1 + cols_per_rank.append(cols_per_rank[-1] + cols) + weight_data = conv.weight.data + if self.shard_by_oc: + # not last conv layer, split output channel + weight_data = weight_data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] + if conv.bias is not None: + bias_data = conv.bias.data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] + else: + # last conv layer, split input channel + weight_data = weight_data[:, cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] + if conv.bias is not None: + bias_data = conv.bias.data / float(self.world_size) + self.conv = nn.Conv2d(weight_data.shape[1], weight_data.shape[0], conv.kernel_size, conv.stride, conv.padding, + conv.dilation, conv.groups, conv.bias is not None, conv.padding_mode) + self.conv.weight = torch.nn.Parameter(weight_data) + if conv.bias is not None: + self.conv.bias = torch.nn.Parameter(bias_data) + del conv + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return self.conv(input) + + +class TensorParallelOcShardConv2d(TensorParallelConv2d): + + def __init__(self, conv, rank, world_size): + super().__init__(conv, rank, world_size, True) + + +class TensorParallelIcShardConv2d(TensorParallelConv2d): + + def __init__(self, conv, rank, world_size): + super().__init__(conv, rank, world_size, False) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + out = self.conv(input) + if self.world_size > 1: + dist.inference_all_reduce(out) + return out + + +class Normalize(nn.Module): + + def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None, bias=None): + super(Normalize, self).__init__() + if weight is not None: + self.weight = weight + self.bias = bias + else: + self.norm = nn.LayerNorm(dim, eps=eps).to(dtype).to(get_accelerator().current_device_name()) + self.weight = self.norm.weight + self.bias = self.norm.bias + + self.eps = eps + + def forward(self, input): + return nn.functional.layer_norm(input, input.shape[-1:], self.weight, self.bias, eps=self.eps) + + +class EmbeddingLayer(nn.Module): + + def __init__(self, weight_shape=None, dtype=torch.half, weight=None, bias=None): + super(EmbeddingLayer, self).__init__() + if weight is None: + self.weight = Parameter( + torch.empty(weight_shape[0], + weight_shape[1], + dtype=dtype, + device=get_accelerator().current_device_name())) + else: + self.weight = weight + + def forward(self, input): + return F.embedding(input, self.weight) + + +class OPTEmbedding(EmbeddingLayer): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, weight_shape=None, weight=None, bias=None): + # OPT is set up so that if padding_idx is specified then offset the embedding ids by 2 + # and adjust num_embeddings appropriately. Other models don't have this hack + self.offset = 2 + super().__init__(weight_shape, weight=weight) + + def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0, position_ids: int = 0): + """`input_ids_shape` is expected to be [bsz x seqlen].""" + attention_mask = attention_mask.long() + + # create positions depending on attention_mask + positions = (torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask).long() - 1 + + # cut positions if `past_key_values_length` is > 0 + positions = positions[:, past_key_values_length:] + + return super().forward(positions + self.offset) + + +class RMSNormalize(nn.Module): + + def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None): + super(RMSNormalize, self).__init__() + if weight is not None: + self.weight = weight + else: + self.weight = nn.Parameter(torch.ones(dim, dtype=dtype, device=get_accelerator().current_device_name())) + + self.eps = eps + + def forward(self, hidden_states): + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return hidden_states * self.weight diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/load_checkpoint.py b/lib/python3.12/site-packages/deepspeed/module_inject/load_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..862628fa7b4b46ad2cd8bf2c3f5706233bda0627 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/load_checkpoint.py @@ -0,0 +1,285 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from torch import nn +from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference +from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference +from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference +from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference + +import deepspeed.ops.transformer as transformer_inference +from .layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding, RMSNormalize +import torch +import gc +from deepspeed.accelerator import get_accelerator +import re + + +def load_model_with_checkpoint(r_module, + sd, + mp_replace, + ckpt_type, + ckpt_mp_size, + weight_quantizer=None, + rank=0, + container=None): + error_msgs = [] + + def prefix_check(): + # if keys start with 'model.' or 'transformer.', don't skip level 0 prefix + for key in sd[0].keys(): + # OPT models + if re.match("^model[.]", key): + return False + # BLOOM models + if re.match("^transformer[.]", key): + return False + return True + + skip_level_0_prefix = prefix_check() and container.policy.use_load_prefix + + def transpose(data): + with torch.no_grad(): + data = data.contiguous() + data1 = data.transpose(-1, -2).reshape(-1) + data.reshape(-1).copy_(data1) + data1 = None + return data.reshape(data.shape[-1], data.shape[-2]) + + def load(module, prefix): + args = (sd[0], prefix, {}, True, [], [], error_msgs) + + if hasattr(module, 'weight'): + module.weight = mp_replace.copy(module.weight.data, sd[0][prefix + 'weight']) + if prefix + 'bias' in sd[0].keys(): + if module.bias.data.is_meta: + # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here + module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"), + requires_grad=module.bias.data.requires_grad) + module.bias = mp_replace.copy(module.bias.data, sd[0][prefix + 'bias']) + args = None + gc.collect() + + def load_transformer_layer(module, prefix): + if ckpt_type == "tp": + + def load_parameters(module, prefix): + for n, p in module.named_parameters(): + if prefix + n in sd[0] and len(n.split('.')) == 1: + if type(sd[0][prefix + n]) is list: + tmp_data, scale = sd[0][prefix + n] + tmp_data = tmp_data + scale = scale.to(get_accelerator().current_device_name()) + # set the quantizer number of groups using the checkpoint scale shape + weight_quantizer.num_groups = scale.shape[0] + else: + tmp_data = sd[0][prefix + n].to(get_accelerator().current_device_name()) + scale = None + src_shape = tmp_data.shape + dst_shape = p.shape + inner_dim = 1 if tmp_data.dtype == torch.int8 else 0 + outer_dim = 0 if tmp_data.dtype == torch.int8 else 1 + if (len(src_shape) == 2 and len(dst_shape) == 2): + if (src_shape[inner_dim] == dst_shape[0] and src_shape[outer_dim] == dst_shape[1]): + if tmp_data.dtype != torch.int8: + p = weight_quantizer.quantize( + transpose(tmp_data) if weight_quantizer.q_int8 else tmp_data) + else: + p = torch.nn.parameter.Parameter(tmp_data, requires_grad=False) + p.scale = scale + setattr(module, n, p) + else: + dim = inner_dim if src_shape[inner_dim] != dst_shape[0] else outer_dim + dim1 = 0 if src_shape[inner_dim] != dst_shape[0] else 1 + if src_shape[dim] > dst_shape[dim1]: + weight_partition = torch.split(tmp_data, dst_shape[dim1], dim=dim)[rank].to( + get_accelerator().current_device_name()) + assert tmp_data.dtype != torch.int8 or scale.numel() > weight_quantizer.num_groups * (rank+1), \ + '''ERROR: We require the quantization scales for larger TP-size when loading INT8 checkpoint!\ + Please use the FP16 checkpoint to generate INT8 checkpoint with the sharding parameters!''' + scale = scale.view(-1)[weight_quantizer.num_groups * (rank + 1):].reshape( + weight_quantizer.num_groups, -1).contiguous() + else: + assert tmp_data.dtype != torch.int8, \ + '''Merging of the checkpoints are not supported when using INT8 checkpoint! \ + Please use a as many GPUs as TP-size for the checkpoint''' + all_data = [ + sd[j][prefix + n] if type(sd[j][prefix + n]) is list else sd[j][prefix + n].to( + get_accelerator().current_device_name()) for j in range(len(sd)) + ] + # Check if the weight tensor is for the QKV parameter + if src_shape[1] == (3 * src_shape[0]) // ckpt_mp_size: + qkv_size = src_shape[outer_dim] // 3 + src_split = [ + torch.split(src[0].data, qkv_size, dim=outer_dim) for src in all_data + ] + + weight_partition = torch.cat([ + torch.cat([qkv_s[i] for qkv_s in src_split], axis=outer_dim) + for i in range(len(src_split[0])) + ], + dim=dim) + else: + weight_partition = torch.cat([ + ad[0].to(get_accelerator().current_device_name()) + if type(ad) is list else ad for ad in all_data + ], + dim=dim) + if tmp_data.dtype == torch.int8: + scale = torch.cat( + [ad[1].to(get_accelerator().current_device_name()) for ad in all_data], + dim=dim) + + if tmp_data.dtype != torch.int8: + weight_partition = weight_quantizer.quantize( + transpose(weight_partition), \ + parallel_dim=(0 if dim == 1 else 1)) if weight_quantizer.q_int8 else \ + weight_quantizer.quantize(weight_partition) + else: + weight_partition = torch.nn.parameter.Parameter(weight_partition, + requires_grad=False) + weight_partition.scale = scale + setattr(module, n, weight_partition) + else: + if src_shape[0] == dst_shape[0]: + p.data.copy_(tmp_data) + else: + if src_shape[0] > dst_shape[0]: + bias_split = torch.split(tmp_data, dst_shape[-1])[rank].to( + get_accelerator().current_device_name()).contiguous() + p.data.copy_(bias_split) + else: + # Check if the weight tensor is for the QKV parameter + if src_shape[0] == (3 * r_module.config.hidden_size) // ckpt_mp_size: + qkv_size = src_shape[0] // 3 + src_split = [ + torch.split(sd[j][prefix + n], qkv_size, dim=0) for j in range(len(sd)) + ] + + p.data.copy_( + torch.cat([ + torch.cat([qkv_s[i] for qkv_s in src_split], axis=0) + for i in range(len(src_split[0])) + ], + dim=0).to(get_accelerator().current_device_name()).contiguous()) + else: + p.data.copy_( + torch.cat([sd[j][prefix + n] for j in range(len(sd))], + dim=0).to(get_accelerator().current_device_name()).contiguous()) + + load_parameters(module, prefix) + for n, child in module.named_children(): + load_parameters(child, prefix + n + '.') + else: + container.load_params(module, sd[0], weight_quantizer, mp_replace, prefix) + + try: + import transformers + OPTLearnedPositionalEmbedding = transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding + if hasattr(transformers.models, "llama"): + LlamaRMSNorm = transformers.models.llama.modeling_llama.LlamaRMSNorm + else: + LlamaRMSNorm = None + except: + OPTLearnedPositionalEmbedding = None + try: + from fairscale.nn.model_parallel.layers import ( + ColumnParallelLinear, + ParallelEmbedding, + RowParallelLinear, + ) + except: + ColumnParallelLinear = None + ParallelEmbedding = None + RowParallelLinear = None + try: + from llama.model import RMSNorm + except: + RMSNorm = None + layer_policies = { + nn.Linear: load, + nn.Embedding: load, + nn.LayerNorm: load, + EmbeddingLayer: load, + LinearLayer: load, + Normalize: load, + transformer_inference.DeepSpeedTransformerInference: load_transformer_layer, + DeepSpeedBloomInference: load_transformer_layer, + DeepSpeedGPTInference: load_transformer_layer, + DeepSpeedBERTInference: load_transformer_layer, + DeepSpeedMegatronGPTInference: load_transformer_layer, + DeepSpeedOPTInference: load_transformer_layer, + DeepSpeedLlama2Inference: load_transformer_layer, + OPTLearnedPositionalEmbedding: load, + OPTEmbedding: load, + LlamaRMSNorm: load, + RMSNormalize: load, + ColumnParallelLinear: load, + ParallelEmbedding: load, + RowParallelLinear: load, + RMSNorm: load + } + + all_ds_ids = {} + + def load_module_recursive(module, prefix='', level=0): + for name, child in module.named_children(): + if child.__class__ in layer_policies: + checking_key = prefix + name + '.' + if not any(checking_key in item for item in sd[0].keys()): + if hasattr(child, 'weight') and \ + (hasattr(child.weight, 'ds_id') and \ + child.weight.ds_id in all_ds_ids): + prefix1 = all_ds_ids[child.weight.ds_id] + if child.__class__ is nn.Linear: + child = LinearLayer.from_weights(weight=all_ds_ids[child.weight.ds_id]) + setattr(module, name, child) + continue + child_params = list(child.parameters()) + if len(child_params) > 0 and (child_params[0].numel() == 0 or child_params[0].is_meta): + if child.weight.is_meta: + ds_shape = child.weight.shape + else: + ds_shape = child.weight.ds_shape + if child.__class__ is nn.LayerNorm: + child = Normalize(dim=ds_shape[-1], dtype=child.weight.dtype, eps=child.eps) + setattr(module, name, child) + elif child.__class__ in [nn.Linear, ColumnParallelLinear, RowParallelLinear]: + child = LinearLayer.from_weights(weight_shape=child.weight.shape, + dtype=child.weight.dtype, + bias=child.bias) + setattr(module, name, child) + elif child.__class__ is OPTLearnedPositionalEmbedding: + child = OPTEmbedding(weight_shape=ds_shape) + setattr(module, name, child) + elif child.__class__ in [LlamaRMSNorm, RMSNorm]: + child = RMSNormalize(dim=ds_shape[-1], + dtype=child.weight.dtype, + eps=child.eps if hasattr(child, 'eps') else child.variance_epsilon) + setattr(module, name, child) + else: + ds_id = None + if hasattr(child.weight, 'ds_id'): + ds_id = child.weight.ds_id + child = EmbeddingLayer(weight_shape=ds_shape, dtype=child.weight.dtype) + if ds_id is not None: + all_ds_ids[ds_id] = child.weight + setattr(module, name, child) + layer_policies[child.__class__](child, prefix + name + '.') + else: + load_module_recursive( + child, + prefix if (level == 0 and ckpt_type == 'pp') and skip_level_0_prefix else \ + prefix + name + '.', + level + 1) + + load_module_recursive(r_module) + + for sd_ in sd: + del sd_ + sd = None + gc.collect() diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/module_quantize.py b/lib/python3.12/site-packages/deepspeed/module_inject/module_quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..1f5b2f8a1d2861a50878e262a61b783a9324ad4e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/module_quantize.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + + +def quantize_transformer_layer(orig_layer_impl, model, megatron=False, preln=False): + """ Quantize bert-style transformer layers with DeepSpeed's transformer layer + Arguments: + orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for, + e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer + model (torch.nn.Module): user's nn.module representing their model + + megatron (bool): megatron model-parallel implementation (this is supported for inference only) + preln (bool): does the original layer implementation do pre or post layer norm? + + Note: For Bert kind of models, we inject based on the DeepSpeed-Example models, if not setting huggingface flag. + + Returns: + Updated nn.module with quantized transformer layers + """ + + def quantize_weight(weight): + return weight.to(torch.int8) + + def megatron_layer_quantize(layer): + layer.attention.query_key_value.weight.data = quantize_weight(layer.attention.query_key_value.weight.data) + layer.attention.dense.weight.data = quantize_weight(layer.attention.dense.weight.data) + layer.mlp.dense_h_to_4h.weight.data = quantize_weight(layer.mlp.dense_h_to_4h.weight.data) + layer.mlp.dense_4h_to_h.weight.data = quantize_weight(layer.mlp.dense_4h_to_h.weight.data) + + def bert_layer_quantize(layer): + layer.attention.self.query.weight.data = quantize_weight(layer.attention.self.query.weight.data) + layer.attention.self.key.weight.data = quantize_weight(layer.attention.self.key.weight.data) + layer.attention.self.value.weight.data = quantize_weight(layer.attention.self.value.weight.data) + layer.attention.output.dense.weight.data = quantize_weight(layer.attention.output.dense.weight.data) + if preln: + layer.intermediate.dense_act.weight.data = quantize_weight(layer.intermediate.dense_act.weight.data) + else: + layer.intermediate.dense.weight.data = quantize_weight(layer.intermediate.dense.weight.data) + layer.output.dense.weight.data = quantize_weight(layer.output.dense.weight.data) + + def quantize_fn(child): + if megatron: + # Quantize megatron GPT2 / GPT3 trained model + megatron_layer_quantize(child) + else: + # Quantize either DeepSpeed or HuggingFace trained model + bert_layer_quantize(child) + + return child + + return quantize_module(model=model, orig_class=orig_layer_impl, quantize_fn=quantize_fn) + + +def quantize_module(model, orig_class, quantize_fn): + policy = {orig_class: quantize_fn} + return _quantize_module(model, policy) + + +def _quantize_module(model, policies): + for name, child in model.named_children(): + if child.__class__ in policies: + orig = repr(child) + setattr(model, name, policies[child.__class__](child)) + new = getattr(model, name) + else: + _quantize_module(child, policies) + + return model diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/policy.py b/lib/python3.12/site-packages/deepspeed/module_inject/policy.py new file mode 100644 index 0000000000000000000000000000000000000000..41df2b85dc0c16fb1b1c21942f5c963fcdc02ad1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/policy.py @@ -0,0 +1,224 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod +from deepspeed.utils.types import ActivationFuncType, NormType +import torch +from deepspeed.accelerator import get_accelerator + +transformer_param_names = ( + 'attn_qkvw', \ + 'attn_qkvb', \ + 'attn_ow' , \ + 'attn_ob', \ + 'inter_w', \ + 'inter_b', \ + 'output_w', \ + 'output_b', \ + 'attn_nw', \ + 'attn_nb', \ + 'norm_w', \ + 'norm_b') + + +class DSPolicy(ABC): + _orig_layer_class = None + + def __init__(self): + self.cuda_graph_supported = False + + @abstractmethod + def attention(self): + """ + Returns attention qkv and dense parameters + weight: (3*hidden, hidden) and (hidden, hidden) + bias: (3*hidden) and (hidden) + """ + raise NotImplementedError + + +class TransformerPolicy(DSPolicy): + # a static class variable containing the HuggingFace model configuration. + # see e.g., transformers.models.opt.configuration_opt.OPTConfig + hf_model_config = None + + def __init__( + self, + inference=True, + linear_layer=True, + scale_attention=True, + megatron_v2=False, + use_mup=False, + # the type of activation function used in MLP + mlp_act_func_type=ActivationFuncType.GELU, + # applies layer norm before attention if `pre_attn_norm` is set to True + pre_attn_norm=True, + # this flag shows whether or not using prefix in loading the checkpoint + use_load_prefix=False, + # whether or not the qkv is stored in the split-format + split_qkv=True, + # Type of normalization to perform + norm_type=NormType.LayerNorm): + super().__init__() + self.cuda_graph_supported = False + self.inference = inference + self.linear_layer = linear_layer + self.scale_attention = scale_attention + self.is_megatron_v2 = megatron_v2 + self.use_mup = use_mup + self.mlp_act_func_type = mlp_act_func_type + self.pre_attn_norm = pre_attn_norm + self.use_load_prefix = use_load_prefix + self.split_qkv = split_qkv + self.norm_type = norm_type + + @abstractmethod + def attention(self): + """ + Returns attention qkv and dense parameters + weight: (3*hidden, hidden) and (hidden, hidden) + bias: (3*hidden) and (hidden) + """ + raise NotImplementedError + + @abstractmethod + def get_hidden_heads(self): + """ + return hidden_size and number of heads + """ + raise NotImplementedError + + @abstractmethod + def mlp(self): + """ + Returns mlp intermediate and output + weight: (intermediate, hidden) and (hidden, intermediate) + bias: (intermediate) and (hidden) + """ + raise NotImplementedError + + @abstractmethod + def layernorm(self): + """ + Returns LayerNorms used in transformer layer + Post-Attention and pre/post layer norm + gamma and beta with shape: (hidden) + """ + raise NotImplementedError + + +# TODO (lekurile): This function exists in base container as well, consolidate as some point +def transpose(data): + with torch.no_grad(): + data = data.contiguous() + data1 = data.transpose(-1, -2).reshape(-1) + data.reshape(-1).copy_(data1) + data1 = None + return data.reshape(data.shape[-1], data.shape[-2]) + + +# TODO (lekurile): This function exists in megatron feature container as well, consolidate as some point +def _transpose(x, heads=1, mp_replace=None): + heads = heads // mp_replace.mp_size # type: ignore + outer_dim = -1 + attention_head_size = x.shape[outer_dim] // heads + new_x_shape = x.size()[:outer_dim] + (heads, attention_head_size) + x_1 = x.view(*new_x_shape) + (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=-1) + if len(q.shape) > 2: + new_shape = (q.shape[0], ) + (-1, ) + return torch.cat((q.reshape(new_shape), k.reshape(new_shape), v.reshape(new_shape)), + dim=outer_dim).reshape(x.shape) + else: + return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape) + + +# This checks if the parameter exits in the checkpoint file and maybe copies it into the corresponding destination tensor. +# Note that not all parameters are saved in one checkpoint, that's why we always need to check if they exist! +def maybe_copy(module, + sd, + weight_quantizer, + mp_replace, + dst_name, + src_name, + qkv=False, + megatron_v2=False, + split_qkv=False, + heads=1): + if src_name in sd: + dst = getattr(module, dst_name) + tmp = sd[src_name] + if len(dst.shape) == 1: + if split_qkv: + dst = mp_replace.strided_copy(dst, tmp, num_splits=3) + else: + dst = mp_replace.copy(dst, tmp) + if qkv and megatron_v2: + dst = torch.nn.parameter.Parameter(_transpose(dst, heads=heads, mp_replace=mp_replace).contiguous()) + else: + if split_qkv: + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ + (transpose(tmp).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) + else: + if qkv and megatron_v2: + tmp = _transpose(transpose(tmp), heads=heads, mp_replace=mp_replace).contiguous() + if weight_quantizer.q_int8: + tmp = transpose(tmp) + dst = mp_replace.copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ + transpose(tmp)), int8=weight_quantizer.q_int8) + setattr(module, dst_name, dst) + + +# Extending the maybe_copy function for when the q, k, and v are in separate parameters! +def maybe_copy_qkv(module, sd, weight_quantizer, mp_replace, dst_name, src_names, split_qkv=False): + if src_names[0] in sd: + q = sd[src_names[0]] + k = sd[src_names[1]] + v = sd[src_names[2]] + qkv_data = torch.cat((q, k, v), dim=0) + dst = getattr(module, dst_name) + if len(dst.shape) == 1: + if split_qkv: + dst = mp_replace.strided_copy(dst, qkv_data.contiguous(), num_splits=3) + else: + dst = mp_replace.copy(dst, qkv_data) + else: + if split_qkv: + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + ((transpose(qkv_data)).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) + else: + dst = mp_replace.copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + transpose(qkv_data)), int8=weight_quantizer.q_int8) + setattr(module, dst_name, dst) + + +# Extending the `maybe_copy` function for when mlp1 is in separate parameters for GeGLU +def maybe_copy_geglu(module, sd, weight_quantizer, mp_replace, dst_name, src_names): + if src_names[0] in sd: + reg_proj = sd[src_names[0]] + gate_proj = sd[src_names[1]] + + mlp1_data = torch.cat((reg_proj, gate_proj), dim=0) + dst = getattr(module, dst_name) + + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(mlp1_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + transpose(mlp1_data)), num_splits=2, int8=weight_quantizer.q_int8) + setattr(module, dst_name, dst) + + +def pack_lora_weights(p): + return [ + p.lora_right_weight, \ + p.lora_left_weight, \ + p.lora_scaling + ] + + +def maybe_get_lora(p): + if hasattr(p, 'lora_right_weight'): + lora_param = pack_lora_weights(p) + else: + lora_param = [] + return lora_param diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/replace_module.py b/lib/python3.12/site-packages/deepspeed/module_inject/replace_module.py new file mode 100644 index 0000000000000000000000000000000000000000..ed94a5021feecaac6be0383b104892e80f489b31 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/replace_module.py @@ -0,0 +1,723 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch +import tqdm +import deepspeed +import deepspeed.ops.transformer as transformer_inference +from deepspeed.ops.transformer.inference.diffusers_attention import DeepSpeedDiffusersAttention +from deepspeed.ops.transformer.inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock +from deepspeed.ops.transformer.inference.diffusers_2d_transformer import Diffusers2DTransformerConfig +from deepspeed.accelerator import get_accelerator +from .replace_policy import replace_policies, generic_policies +from .auto_tp import AutoTP, ReplaceWithTensorSlicing, Loading +from .layers import TensorParallelOcShardConv2d, TensorParallelIcShardConv2d +from deepspeed.module_inject.layers import is_autotp_training_mode +from deepspeed import comm as dist +from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads, set_tp_grain_size + +from .load_checkpoint import load_model_with_checkpoint +import time + +from .utils import policy_to_ds_container +import gc + + +def get_transformer_name(replaced_module): + from .containers import supported_models + from torch.nn import ModuleList + transformer_name = '' + for n, c in replaced_module.named_children(): + if c.__class__ in supported_models: + transformer_name += n + '.' + for name, child in c.named_children(): + if child.__class__ is ModuleList: + transformer_name += name + break + break + return transformer_name + + +class GroupQuantizer: + + def __init__(self, q_int8=True, group_size=1, num_bits=8, num_groups=0): + self.group_size = group_size + self.num_bits = num_bits + self.q_int8 = q_int8 + + self.num_groups = num_groups + + def quantize(self, inputs, qkv=True, count=1, parallel_dim=0): + if not self.q_int8 or not qkv: + inputs = torch.nn.Parameter(inputs, requires_grad=False) + inputs.scale = torch.empty(1) + return inputs + q_range = 2**self.num_bits + num_groups = self.num_groups if self.num_groups > 0 else inputs.shape[0] // self.group_size + inputs = inputs.to(get_accelerator().current_device_name()) + input_flat = inputs.reshape(num_groups, -1).contiguous() + input_min = torch.min(input_flat, dim=1, keepdim=True)[0].float() + input_max = torch.max(input_flat, dim=1, keepdim=True)[0].float() + scale = torch.max(input_min.abs(), input_max.abs()) * 2.0 / (q_range) + input_flat = (input_flat / scale).round().clamp(-q_range // 2, q_range // 2 - 1) + inputs_q = input_flat.reshape(inputs.shape).to(torch.int8).contiguous() + out = torch.nn.Parameter(inputs_q, requires_grad=False) + inputs_split = inputs.split(inputs.shape[parallel_dim] // 2, dim=parallel_dim) + input_flat = [inputs_split[i].reshape(num_groups, -1).contiguous() for i in range(2)] + input_min = [torch.min(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)] + input_max = [torch.max(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)] + scale1 = [(torch.max(input_min[i].abs(), input_max[i].abs()) * 2.0 / (q_range)).squeeze().unsqueeze(0) + for i in range(2)] + + out.scale = torch.cat([scale.squeeze().unsqueeze(0), scale1[0], scale1[1]], dim=0).reshape(num_groups, + -1).contiguous() + return out + + +def _module_match(module): + for policy in generic_policies: + policy = policy() + if policy.match(module): + return policy + return None + + +def generic_injection(module, dtype=None, enable_cuda_graph=True): + + def replace_attn(child, policy): + policy_attn = policy.attention(child) + if policy_attn is None: + return child + if len(policy_attn) == 5: + qkvw, attn_ow, attn_ob, hidden_size, heads = policy_attn + else: + qw, kw, vw, attn_ow, attn_ob, hidden_size, heads = policy_attn + + config = transformer_inference.DeepSpeedInferenceConfig( + hidden_size=hidden_size, + heads=heads, + dtype=dtype, + triangular_masking=False, + max_out_tokens=4096, + ) + attn_module = DeepSpeedDiffusersAttention(config) + + def transpose(data): + data = data.contiguous() + data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1)) + data = data.reshape(data.shape[-1], data.shape[-2]) + data.to(get_accelerator().current_device_name()) + return data + + if len(policy_attn) == 5: + attn_module.attn_qkvw.data = transpose(qkvw.data) + else: + attn_module.attn_qkvw = None + attn_module.attn_qw.data = transpose(qw.data) + attn_module.attn_kw.data = transpose(kw.data) + attn_module.attn_vw.data = transpose(vw.data) + + attn_module.attn_qkvb = None + attn_module.attn_ow.data = transpose(attn_ow.data) + attn_module.attn_ob.data.copy_(attn_ob.data.to(get_accelerator().current_device_name())) + return attn_module + + def replace_attn_block(child, policy): + config = Diffusers2DTransformerConfig() + return DeepSpeedDiffusersTransformerBlock(child, config) + + if isinstance(module, torch.nn.Module): + pass + else: + if dtype not in [torch.float16, torch.half]: + raise ValueError("Generic injection only supported with FP16") + + try: + import diffusers + if hasattr(diffusers.models.attention, 'CrossAttention'): + cross_attention = diffusers.models.attention.CrossAttention + else: + cross_attention = diffusers.models.attention_processor.Attention + attention_block = diffusers.models.attention.BasicTransformerBlock + new_policies = { + cross_attention: replace_attn, + attention_block: replace_attn_block, + } + except ImportError: + new_policies = {} + + #replace_transformer_layer(None, + # module.text_encoder, + # training=False, + # replace_with_kernel_inject=True, + # triangular_masking=True, + # max_out_tokens=8192) + from ..model_implementations.transformers.clip_encoder import DSClipEncoder + cg_encoder = DSClipEncoder(module.text_encoder, enable_cuda_graph=enable_cuda_graph) + setattr(module, 'text_encoder', cg_encoder) + for name in module.__dict__.keys(): + sub_module = getattr(module, name) + policy = _module_match(sub_module) + + if policy is not None: + + def _replace_module(module, policy): + for name, child in module.named_children(): + _replace_module(child, policy) + if child.__class__ in new_policies: + replaced_module = new_policies[child.__class__](child, policy) + setattr(module, name, replaced_module) + + _replace_module(sub_module, policy) + new_module = policy.apply(sub_module, enable_cuda_graph=enable_cuda_graph) + print(f"**** found and replaced {name} w. {type(new_module)}") + setattr(module, name, new_module) + + +container_g = None + + +def replace_transformer_layer(orig_layer_impl, model, checkpoint_dict, config, model_config): + """ Replace bert-style transformer layers with DeepSpeed's transformer layer + Arguments: + orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for, + e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer + model (torch.nn.Module): user's nn.module representing their model + checkpoint_dict: Dictionary for checkpoint passed from the Inference Engine + config: top-level DS Inference config defined in inference/config.py + model_config: HuggingFace model config passed from the inference/engine.py + Returns: + Updated nn.module with replaced transformer layers + """ + # defining globals as internally defined functions inherit these everywhere + quantize = (config.dtype == torch.int8) + # todo: Refactor later. In future, let's minimize the style used above and use config.** instead + + linear_layer_setting = None + ''' + linear_layer_setting (tuple of modules) [Optional]: shows which two classes are used for linear layers and embedding layers + ''' + micro_batch_size = -1 + seed = -1 + local_rank = -1 + + mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group, + mp_size=config.tensor_parallel.tp_size) #, out_dim=0, in_dim=1) + + def replace_with_policy(child, policy_cls, triangular_masking, inference=False, layer_id=0): + policy = policy_cls(child, inference=inference) + if not policy.cuda_graph_supported: + # policy says cuda graph is not supported raise an error if set + assert not config.enable_cuda_graph, "cuda graph is not supported with this model, please disable" + + from deepspeed.moe.layer import MoE + moe = False + if hasattr(child, 'mlp') and isinstance(child.mlp, MoE): + num_experts = child.mlp.num_experts + moe = True + + # 1. Create a model-specific container object using the policy object. + _container = policy_to_ds_container(policy=policy, + config=config, + model_config=model_config, + layer_id=layer_id, + child=child) + _container.set_moe(moe) + + # 2. Set the tensor parallelism config + _container.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group) + + # 3. Initialize tensors + _container.initialize_tensors() + + # 4. deal with data types -- needs refactor to use dtype instead of fp16 + if config.dtype in [torch.float16, torch.bfloat16, torch.int8]: + _container.convert_to_required_dtype() + + # 5. Set the quantization config + quantizer = GroupQuantizer(q_int8=quantize) + _container.set_quantization_config(quantizer) + + # 6. create a DS Inference config object + _container.create_ds_model_config() + + # 7. use the config and create the module + _container.create_module() + + # 8. transpose the weights and bias if needed + _container.transpose() + + # 9. deal with tensor parallelism. + _container.apply_tensor_parallelism(mp_replace) + + # 10. copy the tensors from the model-specific container to the new module + _container.copy_data_to_new_module() + + # 11. set global for generic checkpoint loading + global container_g + + if container_g is None: + container_g = _container + + return _container.module + + def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None): + #mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group) + + # 1. Create AutoTP object + _autotp = AutoTP(module, all_reduce_linears, prefix, state_dict, linear_layer_setting, orig_layer_impl, + config.keep_module_on_host) + + # 2. Set the tensor parallelism config + _autotp.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group) + + # 3. Try to get num_key_heads from model_config.num_key_value_heads + if hasattr(model_config, "vision_config"): + if "MllamaVisionEncoderLayer" in str(module): + num_kv_heads = _autotp.get_model_num_kv_heads(model_config.vision_config) + elif hasattr(model_config, "text_config"): + num_kv_heads = _autotp.get_model_num_kv_heads(model_config.text_config) + else: + num_kv_heads = _autotp.get_model_num_kv_heads(model_config) + else: + num_kv_heads = _autotp.get_model_num_kv_heads(model_config) + + # 4. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division + set_num_kv_heads(num_kv_heads) + + # 4.1 Get n_embd + n_embd = None + multi_query_n_embd_names = ['n_embd', 'hidden_size'] + for name in multi_query_n_embd_names: + if hasattr(model_config, name): + n_embd = getattr(model_config, name) + if n_embd != None: + break + + # 4.2 set n_embd + set_n_embd(n_embd) + + # 4.3 set attention_heads + if hasattr(model_config, 'num_attention_heads'): + set_num_attention_heads(getattr(model_config, 'num_attention_heads')) + + # 4.4 set tp_grain_size + set_tp_grain_size(config.tensor_parallel.tp_grain_size) + + # 5. Set linear policies + _autotp.update_linear_policies() + + # 6. Replace modules + if "lm_head" in all_reduce_linears or "embed_out" in all_reduce_linears: + return _autotp._replace_last_linear_module(module) + return _autotp._replace_module(module) + + def replace_fn(child, _policy, layer_id=0, prefix="", state_dict=None): + training = False # todo: refactor this part to go in the config + if training: + # copy relevant state from child -> new module + new_module = replace_with_policy(child, _policy, config.triangular_masking) + + else: + # copy relevant state from child -> new module + if not is_autotp_training_mode() and config.replace_with_kernel_inject: + new_module = replace_with_policy(child, + _policy, + config.triangular_masking, + inference=True, + layer_id=layer_id) + else: + new_module = replace_wo_policy(child, _policy, prefix=prefix, state_dict=state_dict) + + return new_module + + def set_lm_head(module): + if is_autotp_training_mode(): + # we need to handle autoTP training mode separately. + return + + embedding_weight = None + for n, p in module.named_parameters(): + if "word_embeddings." in n or "embed_tokens." in n or "wte." in n: + embedding_weight = p + if embedding_weight is not None and hasattr(module, "lm_head") and hasattr( + module.lm_head, "weight") and module.lm_head.weight.is_meta: + module.lm_head.weight = embedding_weight + # enable tensor parallel for the last linear + if hasattr(module, "lm_head") and hasattr(module.lm_head, "weight") and isinstance( + module.lm_head, torch.nn.Linear): + module = replace_wo_policy(module, ("lm_head", ), 0, "lm_head") + elif hasattr(module, "embed_out") and hasattr(module.embed_out, "weight") and isinstance( + module.embed_out, torch.nn.Linear): + module = replace_wo_policy(module, ("embed_out", ), 0, "embed_out") + elif hasattr(module, "language_model") and hasattr(module.language_model, "lm_head"): + module = replace_wo_policy(module.language_model, ("lm_head", ), 0, "lm_head") + return module + + def conv2d_parallel_shard_weights(model, rank, world_size): + # add conv policy + shard_oc_name = ["conv1"] + shard_ic_name = ["conv2"] + for name, sub_m in model.named_children(): + for l_name, l_sub_m in sub_m.named_children(): + if l_name in shard_oc_name: + TPConv2d = TensorParallelOcShardConv2d( + l_sub_m, + rank, + world_size, + ) + setattr(sub_m, l_name, TPConv2d) + if l_name in shard_ic_name: + TPConv2d = TensorParallelIcShardConv2d( + l_sub_m, + rank, + world_size, + ) + setattr(sub_m, l_name, TPConv2d) + conv2d_parallel_shard_weights(sub_m, rank, world_size) + + if checkpoint_dict is not None and not config.replace_with_kernel_inject: + # AutoTP shard loading + checkpoint = checkpoint_dict["checkpoints"] + pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards") + for i in range(len(checkpoint)): + checkpoint_file = os.path.join(config.base_dir, checkpoint[i]) + replaced_module = replace_module(model=model, + orig_class=orig_layer_impl, + replace_fn=replace_fn, + _replace_policy=config.injection_policy_tuple, + checkpoint=checkpoint_file) + pbar.update(1) + gc.collect() + # conv2d tp module replace + # Now is for yuan model. Add model list and conv policy to decide whether to replace conv. + if 'Yuan' in str(replaced_module): + conv2d_parallel_shard_weights(replaced_module, dist.get_rank(), dist.get_world_size()) + else: + replaced_module = replace_module(model=model, + orig_class=orig_layer_impl, + replace_fn=replace_fn, + _replace_policy=config.injection_policy_tuple) + # AutoTP default set lm_head tp + if not config.replace_with_kernel_inject: + replaced_module = set_lm_head(replaced_module) + + quantizer = GroupQuantizer(q_int8=quantize) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + if checkpoint_dict is not None and config.replace_with_kernel_inject: + assert container_g.ckpt_load_enabled, \ + f"Meta Tensor checkpoint loading not supported in {container_g.__class__.__name__} container" + start_time = time.time() + checkpoint = checkpoint_dict['checkpoints'] + ckpt_list = checkpoint["tp"] if type(checkpoint) is dict else checkpoint + ckpt_type = checkpoint_dict.get('parallelization', 'pp') + ckpt_mp_size = checkpoint_dict.get('tp_size', len(ckpt_list)) + ckpt_mp_size = checkpoint_dict.get('mp_size', ckpt_mp_size) + base_dir1 = checkpoint_dict.get('base_dir', config.base_dir) + + if ckpt_type == 'pp' and type(checkpoint) is list: + pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards") + + for i in range(len(checkpoint)): + sd = [torch.load(os.path.join(base_dir1, checkpoint[i]), map_location='cpu', weights_only=False)] + load_model_with_checkpoint(replaced_module, + sd, + mp_replace, + ckpt_type, + ckpt_mp_size, + quantizer, + container=container_g) + pbar.update(1) + else: + num_checkpoints = len(ckpt_list) // ckpt_mp_size + tp_split_size = (world_size / ckpt_mp_size) + sd_offset = int(rank / tp_split_size) + sd_count = int((rank + max(1, tp_split_size)) / tp_split_size) - sd_offset + pbar = tqdm.tqdm(total=num_checkpoints, desc=f"Loading {num_checkpoints} checkpoint shards") + for i in range(num_checkpoints): + pbar.update(1) + ckpt_index = i * ckpt_mp_size + sd_offset + ckpt_files = [ + os.path.join(base_dir1, ckpt_list[ckpt_index + j]) if base_dir1 else ckpt_list[ckpt_index + j] + for j in range(sd_count) + ] + sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False) for ckpt_file in ckpt_files] + load_model_with_checkpoint(replaced_module, + sds, + mp_replace, + ckpt_type, + ckpt_mp_size, + quantizer, + int(rank % tp_split_size), + container=container_g) + sds = [None for _ in sds] + gc.collect() + + if "non_tp" in checkpoint: + pbar = tqdm.tqdm(total=len(checkpoint["non_tp"]), + desc=f"Loading {len(checkpoint['non_tp'])} checkpoint shards") + + for i in range(len(checkpoint["non_tp"])): + pbar.update(1) + ckpt_file = os.path.join(base_dir1, + checkpoint["non_tp"][i]) if base_dir1 else checkpoint["non_tp"][i] + sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False)] + load_model_with_checkpoint(replaced_module, + sds, + mp_replace, + ckpt_type, + ckpt_mp_size, + quantizer, + int(rank % tp_split_size), + container=container_g) + sds = [None for _ in sds] + gc.collect() + set_lm_head(replaced_module) + print(f"checkpoint loading time at rank {rank}: {time.time()-start_time} sec") + + if not is_autotp_training_mode() and config.save_mp_checkpoint_path is not None: + from collections import OrderedDict + import json + num_partitions = 8 + + if checkpoint_dict is None: + ckpt_name = "ds_model" + try: + from transformers.models.bloom.modeling_bloom import BloomForCausalLM + if isinstance(model, BloomForCausalLM): + ckpt_name = "bloom" + except ImportError: + ckpt_name = "ds_model" + else: + ckpt_name = checkpoint_dict['type'] + if dist.is_initialized(): + dist.barrier() + transformer_name = get_transformer_name(replaced_module) + non_tp_ckpt_name = f'non-tp.pt' + ckpt_files = [non_tp_ckpt_name] + os.makedirs(config.save_mp_checkpoint_path, exist_ok=True) + + if not dist.is_initialized() or dist.get_rank() == 0: + print("Saving tp-sharded checkpoints") + torch.save( + OrderedDict({ + k: v + for k, v in dict(replaced_module.state_dict()).items() if transformer_name not in k + }), f'{config.save_mp_checkpoint_path}/{non_tp_ckpt_name}') + + dtype_reprs = { + torch.float32: 'float32', + torch.float16: 'float16', + torch.int8: 'int8', + torch.bfloat16: 'bfloat16' + } + + ckpt_config = json.dumps({ + 'type': ckpt_name, + 'base_dir': f'{config.save_mp_checkpoint_path}', + 'checkpoints': { + "non_tp": ckpt_files, + "tp": [f'tp_{r:0>2d}_{m:0>2d}.pt' for m in range(num_partitions) for r in range(world_size)] + }, + 'version': 1.0, + 'parallelization': 'tp', + 'tp_size': world_size, + 'dtype': dtype_reprs[config.dtype] + }) + with open(f"{config.save_mp_checkpoint_path}/ds_inference_config.json", "w") as cfg: + cfg.write(ckpt_config) + + rep_sd = replaced_module.state_dict() + for n, p in replaced_module.named_parameters(): + if hasattr(p, 'scale'): + rep_sd[n] = [p, p.scale] + keys = list(rep_sd.keys()) + partition_size = (len(keys) // num_partitions + 1) + for m in range(num_partitions): + torch.save( + OrderedDict({ + k: [rep_sd[k], rep_sd[k].scale] if hasattr(rep_sd[k], 'scale') else rep_sd[k] + for k in keys[m * partition_size:(m + 1) * partition_size] if transformer_name in k + }), f'{config.save_mp_checkpoint_path}/tp_{rank:0>2d}_{m:0>2d}.pt') + + return replaced_module + + +def revert_transformer_layer(orig_layer_impl, model, config, preln=False): + """ Revert DeepSpeed's transformer layer back to original bert-style transformer layer + Arguments: + orig_layer_impl (torch.nn.Module): the original transformer layer implementation that was replaced, + e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer + model (torch.nn.Module): user's nn.module representing their model + config (dict): model config containing hidden size, attention heads, etc. + Returns: + Updated nn.module with original bert-style transformer layers + """ + + def replace_fn(child, _replace_policy, layer_id): + #from turing.nvidia_modelingpreln import BertLayer + orig_module = orig_layer_impl(config) + + # copy relevant state from child -> original module + qkvw = child.attn_qkvw.data + qkvb = child.attn_qkvb.data + + qw, kw, vw = torch.chunk(qkvw, 3, axis=0) + qb, kb, vb = torch.chunk(qkvb, 3, axis=0) + + orig_module.attention.self.query.weight.data = qw + orig_module.attention.self.query.bias.data = qb + orig_module.attention.self.key.weight.data = kw + orig_module.attention.self.key.bias.data = kb + orig_module.attention.self.value.weight.data = vw + orig_module.attention.self.value.bias.data = vb + + orig_module.attention.output.dense.weight.data = child.attn_ow.data + orig_module.attention.output.dense.bias.data = child.attn_ob.data + + attn_ln_w = child.attn_nw.data + attn_ln_b = child.attn_nb.data + if preln: + orig_module.PostAttentionLayerNorm.weight.data = attn_ln_w + orig_module.PostAttentionLayerNorm.bias.data = attn_ln_b + else: + orig_module.attention.output.LayerNorm.weight.data = attn_ln_w + orig_module.attention.output.LayerNorm.bias.data = attn_ln_b + + inter_ff_w = child.inter_w.data + inter_ff_b = child.inter_b.data + if preln: + orig_module.intermediate.dense_act.weight.data = inter_ff_w + orig_module.intermediate.dense_act.bias.data = inter_ff_b + else: + orig_module.intermediate.dense.weight.data = inter_ff_w + orig_module.intermediate.dense.bias.data = inter_ff_b + + orig_module.output.dense.weight.data = child.output_w.data + orig_module.output.dense.bias.data = child.output_b.data + + transformer_ln_w = child.norm_w.data + transformer_ln_b = child.norm_b.data + if preln: + orig_module.PreAttentionLayerNorm.weight.data = transformer_ln_w + orig_module.PreAttentionLayerNorm.bias.data = transformer_ln_b + else: + orig_module.output.LayerNorm.weight.data = transformer_ln_w + orig_module.output.LayerNorm.bias.data = transformer_ln_b + return orig_module + + return replace_module(model=model, + orig_class=deepspeed.DeepSpeedTransformerLayer, + replace_fn=replace_fn, + _replace_policy=None) + + +def replace_module(model, orig_class, replace_fn, _replace_policy, checkpoint=None): + """ Scan the model for instances of ``orig_clas:`` to replace using ``replace_fn``. + Arguments: + model (torch.nn.Module): the model to augment + orig_class (torch.nn.Module): the module to search for + replace_fn (method): a method to convert instances of ``orig_class`` to the + desired type and return a new instance. + Returns: + A modified ``model``. + """ + sd = None + if checkpoint is not None: + if checkpoint.endswith(".safetensors"): + from safetensors.torch import load_file + sd = load_file(checkpoint) + else: + sd = torch.load(checkpoint, map_location='cpu', weights_only=False) + + policy = {} + if orig_class is not None: + policy.update({orig_class: (replace_fn, _replace_policy)}) + else: + for plcy in replace_policies: + # instantiate a throw-away policy in order to populate the _orig_layer_class + _ = plcy(None) + if isinstance(plcy._orig_layer_class, list): + for orig_layer_class in plcy._orig_layer_class: + policy.update({orig_layer_class: (replace_fn, plcy)}) + elif plcy._orig_layer_class is not None: + policy.update({plcy._orig_layer_class: (replace_fn, plcy)}) + assert len(policy.items()) > 0,\ + "No default policy found! Please specify your policy injection_policy (like {BertLayer:HFBEertLayerPolicy})." +\ + "You can find some samples here: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/replace_policy.py" + + replaced_module, _ = _replace_module(model, policy, state_dict=sd) + return replaced_module + + +from ..pipe import PipelineModule + +import re + + +def skip_level_0_prefix(model, state_dict): + model = str(model) + key = re.search(r": (.*?)Model", model) + if key is None: + key = re.search(r": (.*?)Stack", model) + if key is None: + key = re.match(r"(.*?)Model", model) + # if keys start with 'model.', don't skip level 0 prefix + if state_dict is not None: + for item in state_dict.keys(): + if re.match("^model[.]", item): + return False + if key is not None and key.group(1).lower() in ["bloom", "opt"]: + return True + return False + + +def _replace_module(model, policies, prefix='', layer_id=0, level_id=0, state_dict=None): + """ Traverse model's children recursively and apply any transformations in ``policies``. + Arguments: + model (torch.nn.Module): model to augment + policies (dict): Mapping of source class to replacement function. + Returns: + Modified ``model``. + """ + for name, child in model.named_children(): + if child.__class__ in policies: + replaced_module = policies[child.__class__][0](child, + policies[child.__class__][-1], + layer_id, + prefix=prefix + name, + state_dict=state_dict) + setattr(model, name, replaced_module) + if isinstance(model, PipelineModule): + assert hasattr(model, 'forward_funcs'),\ + "we require pipe-module to have the list of fwd_functions" + model.forward_funcs[model.fwd_map[name]] = replaced_module + layer_id += 1 + else: + checking_key = prefix + name + '.' + if Loading.is_load_module(child) and state_dict is not None: + if any(checking_key in item for item in state_dict): + Loading.load( + child, + state_dict, + checking_key, + ) + else: + continue + if len(child._buffers) != 0 and state_dict is not None: + Loading.load_buffer(child, state_dict, checking_key) + _, layer_id = _replace_module(child, + policies, + prefix if level_id == 0 and skip_level_0_prefix(model, state_dict) else \ + prefix + name + '.', + layer_id=layer_id, + level_id=level_id + 1, + state_dict=state_dict) + + # Add the reset_cache func to the model, so that it can be called in the beginning of text-generation. + model.reset_cache = transformer_inference.DeepSpeedTransformerInference.reset_cache + return model, layer_id diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/replace_policy.py b/lib/python3.12/site-packages/deepspeed/module_inject/replace_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..2c06e31aaa41ae61b5091a7766b709fac2c534f5 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/replace_policy.py @@ -0,0 +1,30 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .containers import HFGPT2LayerPolicy +from .containers import HFBertLayerPolicy +from .containers import BLOOMLayerPolicy +from .containers import HFGPTJLayerPolicy +from .containers import HFGPTNEOLayerPolicy +from .containers import GPTNEOXLayerPolicy +from .containers import HFOPTLayerPolicy +from .containers import MegatronLayerPolicy +from .containers import HFDistilBertLayerPolicy +from .containers import HFCLIPLayerPolicy +from .containers import LLAMALayerPolicy +from .containers import UNetPolicy +from .containers import VAEPolicy +from .containers import LLAMA2LayerPolicy +from .containers import InternLMLayerPolicy + +# transformer-based policies +replace_policies = [ + HFBertLayerPolicy, HFGPTNEOLayerPolicy, GPTNEOXLayerPolicy, HFGPTJLayerPolicy, MegatronLayerPolicy, + HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy, + LLAMALayerPolicy, LLAMA2LayerPolicy, InternLMLayerPolicy +] + +# non-transformer-based policies +generic_policies = [UNetPolicy, VAEPolicy] diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/tp_shard.py b/lib/python3.12/site-packages/deepspeed/module_inject/tp_shard.py new file mode 100644 index 0000000000000000000000000000000000000000..ded262edcf61bbb39d5d3e0c1575430aa8d9b080 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/tp_shard.py @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed import comm as dist +global num_kv_heads + + +def set_num_kv_heads(num): + global num_kv_heads + num_kv_heads = num + + +def set_num_attention_heads(num): + global num_attention_heads + num_attention_heads = num + + +def set_n_embd(num): + global n_embd + n_embd = num + + +def set_tp_grain_size(num): + global tp_grain_size + tp_grain_size = num + + +def get_num_kv_heads(): + global num_kv_heads + if 'num_kv_heads' in globals(): + return num_kv_heads + return None + + +def get_num_attention_heads(): + global num_attention_heads + return num_attention_heads + + +def get_shard_size(total_size, mp_size, name=None, rank=None): + global num_kv_heads + last_linear = ["lm_head", "embed_out"] + # MoE MLP layer use near even division will get better perf. + moe_mlp_layer = ["gate_proj", "up_proj", "down_proj", "w1", "w2", "w3"] + not_moe_mlp_layer = True + if name != None and any(s in str(name) for s in moe_mlp_layer): + not_moe_mlp_layer = False + # When we have num_kv_heads defined, uneven division is possible, otherwise enforce near even division + if rank == None: + rank = dist.get_rank() + if num_kv_heads != None and total_size % num_kv_heads == 0 and "mlp" not in str(name) and str( + name) not in last_linear and not_moe_mlp_layer: + my_slices = (num_kv_heads // mp_size) + (1 if rank < (num_kv_heads % mp_size) else 0) + return total_size * my_slices // num_kv_heads + else: + if total_size >= tp_grain_size: + grain_size = total_size // tp_grain_size + return (grain_size // mp_size + (1 if rank < (grain_size % mp_size) else 0)) * tp_grain_size + else: + return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0) + + +def get_n_embd(): + global n_embd + return n_embd + + +def get_shard_size_list(total_size, mp_size, name=None): + shard_sizes = [] + for i in range(mp_size): + shard_sizes.append(get_shard_size(total_size, mp_size, name, i)) + return shard_sizes diff --git a/lib/python3.12/site-packages/deepspeed/module_inject/utils.py b/lib/python3.12/site-packages/deepspeed/module_inject/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..42822128f9e11c84660dcd57a2473753b3bc6642 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/module_inject/utils.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.utils import log_dist + + +# helper function to map between DS policies and DS containers +def policy_to_ds_container(**kwargs): + from .containers import HFGPT2LayerPolicy, DS_GPT2Container + from .containers import HFBertLayerPolicy, DS_BERTContainer + from .containers import BLOOMLayerPolicy, DS_BloomContainer + from .containers import HFGPTJLayerPolicy, DS_GPTJContainer + from .containers import HFGPTNEOLayerPolicy, DS_GPTNEOContainer + from .containers import GPTNEOXLayerPolicy, DS_GPTNEOXContainer + from .containers import HFOPTLayerPolicy, DS_OPTContainer + from .containers import MegatronLayerPolicy, DS_MegatronGPTContainer + from .containers import HFDistilBertLayerPolicy, DS_DistilBERTContainer + from .containers import LLAMALayerPolicy, DS_LLAMAContainer + from .containers import LLAMA2LayerPolicy, DS_LLAMA2Container + from .containers import InternLMLayerPolicy, DS_InternLMContainer + + policy_to_container = { + HFGPT2LayerPolicy: DS_GPT2Container, + HFBertLayerPolicy: DS_BERTContainer, + BLOOMLayerPolicy: DS_BloomContainer, + HFGPTJLayerPolicy: DS_GPTJContainer, + HFGPTNEOLayerPolicy: DS_GPTNEOContainer, + GPTNEOXLayerPolicy: DS_GPTNEOXContainer, + HFOPTLayerPolicy: DS_OPTContainer, + MegatronLayerPolicy: DS_MegatronGPTContainer, + HFDistilBertLayerPolicy: DS_DistilBERTContainer, + LLAMALayerPolicy: DS_LLAMAContainer, + LLAMA2LayerPolicy: DS_LLAMA2Container, + InternLMLayerPolicy: DS_InternLMContainer + } + + container = None + policy = kwargs['policy'] + assert policy is not None, "Policy cannot be None" + policy_type = type(policy) + + if policy_type not in policy_to_container: + log_dist(f"Policy type {policy_type} not supported", [0]) + else: + container = policy_to_container[policy_type](**kwargs) + + return container diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__init__.py b/lib/python3.12/site-packages/deepspeed/monitor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e8199df973191ef009acb6c67cbb85fc3f327be Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/comet.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/comet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6dc29950790004301ae39a09fc4944d4cafeb09e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/comet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..261d1a81c19fe6be151910154c5ff5645621e0f1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1eb6fd7072204d89b395190c897315b801c9d59a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d44ceaaadcaf972a349965e1fae2c3dfef855b23 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1632b014b75a59870f2d21aaba915fb3fe1feae5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04bc5bfb5a77e379ab147cd608560ab09770993e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2de7229c8524c31b80524d6f120059f7e876626b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/monitor/comet.py b/lib/python3.12/site-packages/deepspeed/monitor/comet.py new file mode 100644 index 0000000000000000000000000000000000000000..d8bc4017800f3914fd1a1881b0510c3cdee75324 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/comet.py @@ -0,0 +1,92 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import TYPE_CHECKING, Any, Tuple, List, Dict, Optional + +from .utils import check_comet_availability +from .monitor import Monitor + +import deepspeed.comm as dist + +if TYPE_CHECKING: + import comet_ml + from .config import CometConfig + +Name = str +Value = Any +GlobalSamples = int +Event = Tuple[Name, Value, GlobalSamples] + + +class CometMonitor(Monitor): + + def __init__(self, comet_config: "CometConfig"): + super().__init__(comet_config) + check_comet_availability() + import comet_ml + + self.enabled = comet_config.enabled + self._samples_log_interval = comet_config.samples_log_interval + self._experiment: Optional["comet_ml.ExperimentBase"] = None + + if self.enabled and dist.get_rank() == 0: + self._experiment = comet_ml.start( + api_key=comet_config.api_key, + project=comet_config.project, + workspace=comet_config.workspace, + experiment_key=comet_config.experiment_key, + mode=comet_config.mode, + online=comet_config.online, + ) + + if comet_config.experiment_name is not None: + self._experiment.set_name(comet_config.experiment_name) + + self._events_log_scheduler = EventsLogScheduler(comet_config.samples_log_interval) + + @property + def experiment(self) -> Optional["comet_ml.ExperimentBase"]: + return self._experiment + + @property + def samples_log_interval(self) -> int: + return self._samples_log_interval + + def write_events(self, event_list: List[Event]) -> None: + if not self.enabled or dist.get_rank() != 0: + return None + + for event in event_list: + name = event[0] + value = event[1] + engine_global_samples = event[2] + + if self._events_log_scheduler.needs_logging(name, engine_global_samples): + self._experiment.__internal_api__log_metric__( + name=name, + value=value, + step=engine_global_samples, + ) + + +class EventsLogScheduler: + + def __init__(self, samples_log_interval: int): + self._samples_log_interval = samples_log_interval + self._last_logged_events_samples: Dict[str, int] = {} + + def needs_logging(self, name: str, current_sample: int) -> bool: + if name not in self._last_logged_events_samples: + self._last_logged_events_samples[name] = current_sample + return True + + last_logged_sample = self._last_logged_events_samples[name] + samples_delta = current_sample - last_logged_sample + + if samples_delta >= self._samples_log_interval: + self._last_logged_events_samples[name] = current_sample + return True + + return False diff --git a/lib/python3.12/site-packages/deepspeed/monitor/config.py b/lib/python3.12/site-packages/deepspeed/monitor/config.py new file mode 100644 index 0000000000000000000000000000000000000000..960ce1ba997a231edec53e84855942f3d6ece1d8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/config.py @@ -0,0 +1,144 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Optional + +from pydantic import model_validator +from deepspeed.runtime.config_utils import DeepSpeedConfigModel + + +def get_monitor_config(param_dict): + monitor_dict = {key: param_dict.get(key, {}) for key in ("tensorboard", "wandb", "csv_monitor", "comet")} + return DeepSpeedMonitorConfig(**monitor_dict) + + +class TensorBoardConfig(DeepSpeedConfigModel): + """Sets parameters for TensorBoard monitor.""" + + enabled: bool = False + """ Whether logging to Tensorboard is enabled. Requires `tensorboard` package is installed. """ + + output_path: str = "" + """ + Path to where the Tensorboard logs will be written. If not provided, the + output path is set under the training script’s launching path. + """ + + job_name: str = "DeepSpeedJobName" + """ Name for the current job. This will become a new directory inside `output_path`. """ + + +class WandbConfig(DeepSpeedConfigModel): + """Sets parameters for WandB monitor.""" + + enabled: bool = False + """ Whether logging to WandB is enabled. Requires `wandb` package is installed. """ + + group: Optional[str] = None + """ Name for the WandB group. This can be used to group together runs. """ + + team: Optional[str] = None + """ Name for the WandB team. """ + + project: str = "deepspeed" + """ Name for the WandB project. """ + + +class CSVConfig(DeepSpeedConfigModel): + """Sets parameters for CSV monitor.""" + + enabled: bool = False + """ Whether logging to local CSV files is enabled. """ + + output_path: str = "" + """ + Path to where the csv files will be written. If not provided, the output + path is set under the training script’s launching path. + """ + + job_name: str = "DeepSpeedJobName" + """ Name for the current job. This will become a new directory inside `output_path`. """ + + +class CometConfig(DeepSpeedConfigModel): + """ + Sets parameters for Comet monitor. For logging data Comet uses + experiment object. + https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/ + """ + + enabled: bool = False + """ Whether logging to Comet is enabled. Requires `comet_ml` package is installed. """ + + samples_log_interval: int = 100 + """ Metrics will be submitted to Comet after processing every `samples_log_intervas` samples""" + + project: Optional[str] = None + """ + Comet project name. Can be set through .comet.config file or environment variable COMET_PROJECT_NAME + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + workspace: Optional[str] = None + """ + Comet workspace name. Can be set through .comet.config file or environment variable COMET_WORKSPACE + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + api_key: Optional[str] = None + """ + Comet API key. Can be set through .comet.config file or environment variable COMET_API_KEY + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + experiment_name: Optional[str] = None + """ + The name for comet experiment to be used for logging. + Can be set through .comet.config file or environment variable COMET_EXPERIMENT_NAME + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + experiment_key: Optional[str] = None + """ + The key for comet experiment to be used for logging. Must be an alphanumeric string whose length is between 32 and 50 characters. + Can be set through .comet.config or environment variable COMET_EXPERIMENT_KEY + https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options + """ + + online: Optional[bool] = None + """ + If True, the data will be logged to Comet server, otherwise it will be stored locally in offline experiment + Defaults to True. + """ + + mode: Optional[str] = None + """ + Control how the Comet experiment is started, 3 options are possible.: + - "get": Continue logging to an existing experiment identified by the `experiment_key` value. + - "create": Always creates of a new experiment, useful for HPO sweeps. + - "get_or_create" (default): Starts a fresh experiment if required, or persists logging to an existing one. + """ + + +class DeepSpeedMonitorConfig(DeepSpeedConfigModel): + """Sets parameters for various monitoring methods.""" + + tensorboard: TensorBoardConfig = {} + """ TensorBoard monitor, requires `tensorboard` package is installed. """ + + comet: CometConfig = {} + """ Comet monitor, requires `comet_ml` package is installed """ + + wandb: WandbConfig = {} + """ WandB monitor, requires `wandb` package is installed. """ + + csv_monitor: CSVConfig = {} + """ Local CSV output of monitoring data. """ + + @model_validator(mode="after") + def check_enabled(self): + enabled = self.tensorboard.enabled or self.wandb.enabled or self.csv_monitor.enabled or self.comet.enabled + self.__dict__["enabled"] = enabled + return self diff --git a/lib/python3.12/site-packages/deepspeed/monitor/csv_monitor.py b/lib/python3.12/site-packages/deepspeed/monitor/csv_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..c7a19b14ad8227fc8187bbe161667285b7d0c717 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/csv_monitor.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .monitor import Monitor +import os + +import deepspeed.comm as dist + + +class csvMonitor(Monitor): + + def __init__(self, csv_config): + super().__init__(csv_config) + self.filenames = [] + self.enabled = csv_config.enabled + self.output_path = csv_config.output_path + self.job_name = csv_config.job_name + self.log_dir = self.setup_log_dir() + + def setup_log_dir(self, base=os.path.join(os.path.expanduser("~"), "csv_monitor")): + if self.enabled and dist.get_rank() == 0: + if self.output_path is not None: + log_dir = os.path.join(self.output_path, self.job_name) + # NOTE: This code path currently is never used since the default tensorboard_output_path is an empty string and not None. Saving it in case we want this functionality in the future. + else: + if "DLWS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLWS_JOB_ID"] + elif "DLTS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLTS_JOB_ID"] + else: + infra_job_id = "unknown-job-id" + + csv_monitor_dir_name = os.path.join(infra_job_id, "logs") + log_dir = os.path.join(base, csv_monitor_dir_name, self.job_name) + os.makedirs(log_dir, exist_ok=True) + return log_dir + + def write_events(self, event_list): + if self.enabled and dist.get_rank() == 0: + import csv + # We assume each event_list element is a tensorboard-style tuple in the format: (log_name: String, value, step: Int) + for event in event_list: + log_name = event[0] + value = event[1] + step = event[2] + + # Set the header to the log_name + # Need this check because the deepspeed engine currently formats log strings to separate with '/' + if '/' in log_name: + record_splits = log_name.split('/') + header = record_splits[len(record_splits) - 1] + else: + header = log_name + + # sanitize common naming conventions into filename + filename = log_name.replace('/', '_').replace(' ', '_') + fname = self.log_dir + '/' + filename + '.csv' + + # Open file and record event. Insert header if this is the first time writing + with open(fname, 'a+') as csv_monitor_file: + csv_monitor_writer = csv.writer(csv_monitor_file) + if filename not in self.filenames: + self.filenames.append(filename) + csv_monitor_writer.writerow(['step', header]) + csv_monitor_writer.writerow([step, value]) diff --git a/lib/python3.12/site-packages/deepspeed/monitor/monitor.py b/lib/python3.12/site-packages/deepspeed/monitor/monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..e7e26dc483d9707bfc04febda348bff30d25e53b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/monitor.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Support different forms of monitoring such as wandb and tensorboard +""" + +from abc import ABC, abstractmethod +import deepspeed.comm as dist + + +class Monitor(ABC): + + @abstractmethod + def __init__(self, monitor_config): + self.monitor_config = monitor_config + + @abstractmethod + def write_events(self, event_list): + pass + + +from .wandb import WandbMonitor +from .tensorboard import TensorBoardMonitor +from .csv_monitor import csvMonitor +from .comet import CometMonitor + + +class MonitorMaster(Monitor): + + def __init__(self, monitor_config): + super().__init__(monitor_config) + self.tb_monitor = None + self.wandb_monitor = None + self.csv_monitor = None + self.comet_monitor = None + self.enabled = monitor_config.enabled + + if dist.get_rank() == 0: + if monitor_config.tensorboard.enabled: + self.tb_monitor = TensorBoardMonitor(monitor_config.tensorboard) + if monitor_config.wandb.enabled: + self.wandb_monitor = WandbMonitor(monitor_config.wandb) + if monitor_config.csv_monitor.enabled: + self.csv_monitor = csvMonitor(monitor_config.csv_monitor) + if monitor_config.comet.enabled: + self.comet_monitor = CometMonitor(monitor_config.comet) + + def write_events(self, event_list): + if dist.get_rank() == 0: + if self.tb_monitor is not None: + self.tb_monitor.write_events(event_list) + if self.wandb_monitor is not None: + self.wandb_monitor.write_events(event_list) + if self.csv_monitor is not None: + self.csv_monitor.write_events(event_list) + if self.comet_monitor is not None: + self.comet_monitor.write_events(event_list) diff --git a/lib/python3.12/site-packages/deepspeed/monitor/tensorboard.py b/lib/python3.12/site-packages/deepspeed/monitor/tensorboard.py new file mode 100644 index 0000000000000000000000000000000000000000..985c9ed44b6f5be28785699bca4a3638ed6063bd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/tensorboard.py @@ -0,0 +1,56 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .utils import check_tb_availability +from .monitor import Monitor +import os + +import deepspeed.comm as dist + + +class TensorBoardMonitor(Monitor): + + def __init__(self, tensorboard_config): + super().__init__(tensorboard_config) + check_tb_availability() + + self.summary_writer = None + self.enabled = tensorboard_config.enabled + self.output_path = tensorboard_config.output_path + self.job_name = tensorboard_config.job_name + + if self.enabled and dist.get_rank() == 0: + self.get_summary_writer() + + def get_summary_writer(self, base=os.path.join(os.path.expanduser("~"), "tensorboard")): + if self.enabled and dist.get_rank() == 0: + from torch.utils.tensorboard import SummaryWriter + if self.output_path is not None: + log_dir = os.path.join(self.output_path, self.job_name) + # NOTE: This code path currently is never used since the default output_path is an empty string and not None. Saving it in case we want this functionality in the future. + else: + if "DLWS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLWS_JOB_ID"] + elif "DLTS_JOB_ID" in os.environ: + infra_job_id = os.environ["DLTS_JOB_ID"] + else: + infra_job_id = "unknown-job-id" + + summary_writer_dir_name = os.path.join(infra_job_id, "logs") + log_dir = os.path.join(base, summary_writer_dir_name, self.output_path) + os.makedirs(log_dir, exist_ok=True) + self.summary_writer = SummaryWriter(log_dir=log_dir) + return self.summary_writer + + def write_events(self, event_list, flush=True): + if self.enabled and self.summary_writer is not None and dist.get_rank() == 0: + for event in event_list: + self.summary_writer.add_scalar(*event) + if flush: + self.summary_writer.flush() + + def flush(self): + if self.enabled and self.summary_writer is not None and dist.get_rank() == 0: + self.summary_writer.flush() diff --git a/lib/python3.12/site-packages/deepspeed/monitor/utils.py b/lib/python3.12/site-packages/deepspeed/monitor/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f5530e8532e1a6982ce8acdf20b98aa8fdbb5e81 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/utils.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from packaging import version as pkg_version + + +def check_tb_availability(): + try: + # torch.utils.tensorboard will fail if `tensorboard` is not available, + # see their docs for more details: https://pytorch.org/docs/1.8.0/tensorboard.html + import tensorboard # noqa: F401 # type: ignore + except ImportError: + print('If you want to use tensorboard logging, please `pip install tensorboard`') + raise + + +def check_wandb_availability(): + try: + import wandb # noqa: F401 # type: ignore + except ImportError: + print( + 'If you want to use wandb logging, please `pip install wandb` and follow the instructions at https://docs.wandb.ai/quickstart' + ) + raise + + +def check_comet_availability(): + try: + import comet_ml + comet_version = pkg_version.parse(comet_ml.__version__) + if comet_version < pkg_version.Version("3.41.0"): + raise ImportError("`comet_ml` must have at least version 3.41.0") + except ImportError: + print('If you want to use comet logging, please `pip install "comet_ml>=3.41.0"`') + raise diff --git a/lib/python3.12/site-packages/deepspeed/monitor/wandb.py b/lib/python3.12/site-packages/deepspeed/monitor/wandb.py new file mode 100644 index 0000000000000000000000000000000000000000..30209191171afc4c0ad3ff7088639f7d678cd505 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/monitor/wandb.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .utils import check_wandb_availability +from .monitor import Monitor + +import deepspeed.comm as dist + + +class WandbMonitor(Monitor): + + def __init__(self, wandb_config): + super().__init__(wandb_config) + check_wandb_availability() + import wandb + + self.enabled = wandb_config.enabled + self.group = wandb_config.group + self.team = wandb_config.team + self.project = wandb_config.project + + if self.enabled and dist.get_rank() == 0: + wandb.init(project=self.project, group=self.group, entity=self.team) + + def log(self, data, step=None, commit=None, sync=None): + if self.enabled and dist.get_rank() == 0: + import wandb + return wandb.log(data, step=step, commit=commit, sync=sync) + + def write_events(self, event_list): + if self.enabled and dist.get_rank() == 0: + for event in event_list: + label = event[0] + value = event[1] + step = event[2] + self.log({label: value}, step=step) diff --git a/lib/python3.12/site-packages/deepspeed/nebula/__init__.py b/lib/python3.12/site-packages/deepspeed/nebula/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/nebula/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/nebula/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/nebula/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50e09ae38c80300275f4683b6c21e19d668fa6a4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/nebula/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/nebula/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/nebula/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d40df9c7d83c2829aaf1db0b1225a23ac81c57e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/nebula/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/nebula/config.py b/lib/python3.12/site-packages/deepspeed/nebula/config.py new file mode 100644 index 0000000000000000000000000000000000000000..dc49185738c92a3173f2fd5c68f1d6ab5a32dd92 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/nebula/config.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject +from deepspeed.nebula.constants import * + + +class DeepSpeedNebulaConfig(DeepSpeedConfigObject): + + def __init__(self, param_dict): + super(DeepSpeedNebulaConfig, self).__init__() + + self.enabled = None + self.persistent_storage_path = None + self.persistent_time_interval = None + self.num_of_version_in_retention = None + self.enable_nebula_load = None + + if NEBULA in param_dict.keys(): + nebula_dict = param_dict[NEBULA] + else: + nebula_dict = {} + + self._initialize(nebula_dict) + + def _initialize(self, nebula_dict): + self.enabled = get_scalar_param(nebula_dict, NEBULA_ENABLED, NEBULA_ENABLED_DEFAULT) + + self.load_path = get_scalar_param(nebula_dict, NEBULA_LOAD_PATH, NEBULA_LOAD_PATH_DEFAULT) + + self.enable_nebula_load = get_scalar_param(nebula_dict, NEBULA_ENABLE_NEBULA_LOAD, + NEBULA_ENABLE_NEBULA_LOAD_DEFAULT) + + self.persistent_storage_path = get_scalar_param(nebula_dict, NEBULA_PERSISTENT_STORAGE_PATH, + NEBULA_PERSISTENT_STORAGE_PATH_DEFAULT) + + self.persistent_time_interval = get_scalar_param(nebula_dict, NEBULA_PERSISTENT_TIME_INTERVAL, + NEBULA_PERSISTENT_TIME_INTERVAL_DEFAULT) + + self.num_of_version_in_retention = get_scalar_param(nebula_dict, NEBULA_NUM_OF_VERSION_IN_RETENTION, + NEBULA_NUM_OF_VERSION_IN_RETENTION_DEFAULT) diff --git a/lib/python3.12/site-packages/deepspeed/nebula/constants.py b/lib/python3.12/site-packages/deepspeed/nebula/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..9fa5769b55979e4dcd5c80ead06d2117dcc2ec40 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/nebula/constants.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######################################### +# nebula +######################################### +# Nebula. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +NEBULA_FORMAT = ''' +nebula should be enabled as: +"session_params": { + "nebula": { + "enabled": true, + "persistent_storage_path": "/foo/bar", + "persistent_time_interval": 100, + "num_of_version_in_retention": 2, + "enable_nebula_load": true + } +} +''' + +NEBULA = "nebula" + +NEBULA_ENABLED = "enabled" +NEBULA_ENABLED_DEFAULT = False + +# There is a case where customer want to load the checkpoint saved +# by raw torch. Because nebula cannot load torch checkpoint directly +# as they have different folder structures to bring the gap for +# loading(the data are totally same in bytes for torch and nebula +# saving). +# In this case, we must disable nebula load to use raw torch load. +# Customer can just set NEBULA_ENABLE_NEBULA_LOAD to False. Then use +# original way of deepspeed to load, i.e. set the value of "--load". +NEBULA_ENABLE_NEBULA_LOAD = "enable_nebula_load" +NEBULA_ENABLE_NEBULA_LOAD_DEFAULT = True + +# When you want to resume the previous checkpoint saved by nebula, +# you can set NEBULA_LOAD_PATH as the parent folder of checkpoint. +# If NEBULA_LOAD_PATH is None, the NEBULA_PERSISTENT_STORAGE_PATH +# will be the default path to load. +NEBULA_LOAD_PATH = "nebula_load_path" +NEBULA_LOAD_PATH_DEFAULT = None + +# Nebula will save the checkpoint under NEBULA_LOAD_PATH in the +# asynchronous way. +NEBULA_PERSISTENT_STORAGE_PATH = "persistent_storage_path" +NEBULA_PERSISTENT_STORAGE_PATH_DEFAULT = None + +# Time interval to trigger the nebula persistence. +NEBULA_PERSISTENT_TIME_INTERVAL = "persistent_time_interval" +NEBULA_PERSISTENT_TIME_INTERVAL_DEFAULT = 100 + +# Checkpoint number which will be kept in memory. Let us say, +# if the value is 2. Then we have checkpoints 1 and 2 are ready +# now. When it comes to checkpoint 3, the 1 will be removed if +# 1 has been persisted to disk. +NEBULA_NUM_OF_VERSION_IN_RETENTION = "num_of_version_in_retention" +NEBULA_NUM_OF_VERSION_IN_RETENTION_DEFAULT = 2 + +# Nebula envs +NEBULA_EXPORT_ENVS = [ + 'DLTS_JOB_ID', 'DLTS_NUM_WORKER', 'NEBULA_PERSISTENT_STORAGE_PATH', 'NEBULA_PERSISTENT_TIME_INTERVAL', + 'AML_RUN_ID', 'AZUREML_RUN_TOKEN', 'AZUREML_WORKSPACE_SCOPE', 'AZUREML_EXPERIMENT_SCOPE', + 'AZUREML_RUN_HISTORY_SERVICE_ENDPOINT', 'AZUREML_RUN_ID', 'NEBULA_MEMORY_BUFFER_SIZE', + 'AZUREML_PARAMETER_ITPJOB_NAME', 'FC_TASKROLE_NAME', 'FC_TASK_INDEX', 'MASTER_HOST', 'LOCAL_HOST', + 'AZUREML_BLOB_ACCOUNT_NAME', 'AZUREML_BLOB_ACCOUNT_KEY' +] + +# ITP env files +DLTS_POD_ENV_PATH = '/dlts-runtime/env/pod.env' diff --git a/lib/python3.12/site-packages/deepspeed/pipe/__init__.py b/lib/python3.12/site-packages/deepspeed/pipe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f25471713b4568b43216fe16786a46904575ae5e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/pipe/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..runtime.pipe import PipelineModule, LayerSpec, TiedLayerSpec diff --git a/lib/python3.12/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b60cdef13e705eb58407aaec45483f6b5ee65299 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/profiling/__init__.py b/lib/python3.12/site-packages/deepspeed/profiling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/profiling/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c906ab83ce47e07aaf2adc824dfa1842e0f1e730 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24bdcefb7ed785e1e548fcfb7da7a9ca78a1bdd5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63458030541226dea8e25b57fbd973f35f6f3ee6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/profiling/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/profiling/config.py b/lib/python3.12/site-packages/deepspeed/profiling/config.py new file mode 100644 index 0000000000000000000000000000000000000000..e4f06630ea6f6801883046efb5f61df3e7132f63 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/profiling/config.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject +from deepspeed.profiling.constants import * + + +class DeepSpeedFlopsProfilerConfig(DeepSpeedConfigObject): + + def __init__(self, param_dict): + super(DeepSpeedFlopsProfilerConfig, self).__init__() + + self.enabled = None + self.recompute_fwd_factor = None + self.profile_step = None + self.module_depth = None + self.top_modules = None + + if FLOPS_PROFILER in param_dict.keys(): + flops_profiler_dict = param_dict[FLOPS_PROFILER] + else: + flops_profiler_dict = {} + + self._initialize(flops_profiler_dict) + + def _initialize(self, flops_profiler_dict): + self.enabled = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_ENABLED, FLOPS_PROFILER_ENABLED_DEFAULT) + + self.recompute_fwd_factor = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR, + FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR_DEFAULT) + + self.profile_step = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_PROFILE_STEP, + FLOPS_PROFILER_PROFILE_STEP_DEFAULT) + + self.module_depth = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_MODULE_DEPTH, + FLOPS_PROFILER_MODULE_DEPTH_DEFAULT) + + self.top_modules = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_TOP_MODULES, + FLOPS_PROFILER_TOP_MODULES_DEFAULT) + + self.detailed = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_DETAILED, FLOPS_PROFILER_DETAILED_DEFAULT) + + self.output_file = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_OUTPUT_FILE, + FLOPS_PROFILER_OUTPUT_FILE_DEFAULT) diff --git a/lib/python3.12/site-packages/deepspeed/profiling/constants.py b/lib/python3.12/site-packages/deepspeed/profiling/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..0374303d7d969c76cd6771b2b18bf46f94df8d04 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/profiling/constants.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######################################### +# flops profiler +######################################### +# Flops profiler. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +FLOPS_PROFILER_FORMAT = ''' +flops profiler should be enabled as: +"session_params": { + "flops_profiler": { + "enabled": true, + "recompute_fwd_factor": 0.0, + "profile_step": 1, + "module_depth": -1, + "top_modules": 3, + "detailed": true, + "output_file": null + } +} +''' + +FLOPS_PROFILER = "flops_profiler" + +FLOPS_PROFILER_ENABLED = "enabled" +FLOPS_PROFILER_ENABLED_DEFAULT = False + +FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR = "recompute_fwd_factor" +FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR_DEFAULT = 0.0 + +FLOPS_PROFILER_PROFILE_STEP = "profile_step" +FLOPS_PROFILER_PROFILE_STEP_DEFAULT = 1 + +FLOPS_PROFILER_MODULE_DEPTH = "module_depth" +FLOPS_PROFILER_MODULE_DEPTH_DEFAULT = -1 + +FLOPS_PROFILER_TOP_MODULES = "top_modules" +FLOPS_PROFILER_TOP_MODULES_DEFAULT = 1 + +FLOPS_PROFILER_DETAILED = "detailed" +FLOPS_PROFILER_DETAILED_DEFAULT = True + +FLOPS_PROFILER_OUTPUT_FILE = "output_file" +FLOPS_PROFILER_OUTPUT_FILE_DEFAULT = None diff --git a/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__init__.py b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..980d8a5e92fc4c3491f822479d05c5e2776233df --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .profiler import * diff --git a/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c957d0652ceae7199617a09238b16e9a1f8f3bca Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fce1cd28df9950fc4c153f7cb620b4f8d0d0078b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/profiler.py b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..c6b607aad813d557258ed40a2730a1d9d1c1db55 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/profiling/flops_profiler/profiler.py @@ -0,0 +1,1293 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import time +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial +from typing import List, Optional +from collections import OrderedDict +import numpy as np +from deepspeed.accelerator import get_accelerator +from deepspeed.utils import logger +from deepspeed.moe.layer import MoE +from deepspeed.utils.timer import FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, STEP_GLOBAL_TIMER +from deepspeed.utils.torch import required_torch_version +import einops + +Tensor = torch.Tensor + +module_flop_count = [] +module_mac_count = [] +old_functions = {} + +DEFAULT_PRECISION = 2 + + +class FlopsProfiler(object): + """Measures the latency, number of estimated floating-point operations and parameters of each module in a PyTorch model. + + The flops-profiler profiles the forward pass of a PyTorch model and prints the model graph with the measured profile attached to each module. It shows how latency, flops and parameters are spent in the model and which modules or layers could be the bottleneck. It also outputs the names of the top k modules in terms of aggregated latency, flops, and parameters at depth l with k and l specified by the user. The output profile is computed for each batch of input. + The DeepSpeed flops profiler can be used with the DeepSpeed runtime or as a standalone package. + When using DeepSpeed for model training, the flops profiler can be configured in the deepspeed_config file and no user code change is required. + + If using the profiler as a standalone package, one imports the flops_profiler package and use the APIs. + + Here is an example for usage in a typical training workflow: + + .. code-block:: python + + model = Model() + prof = FlopsProfiler(model) + + for step, batch in enumerate(data_loader): + if step == profile_step: + prof.start_profile() + + loss = model(batch) + + if step == profile_step: + flops = prof.get_total_flops(as_string=True) + params = prof.get_total_params(as_string=True) + prof.print_model_profile(profile_step=profile_step) + prof.end_profile() + + loss.backward() + optimizer.step() + + To profile a trained model in inference, use the `get_model_profile` API. + + Args: + object (torch.nn.Module): The PyTorch model to profile. + """ + + def __init__(self, model, ds_engine=None, recompute_fwd_factor=0.0): + self.model = model + self.ds_engine = ds_engine + self.recompute_fwd_factor = recompute_fwd_factor + self.started = False + self.func_patched = False + + def start_profile(self, ignore_list=None): + """Starts profiling. + + Extra attributes are added recursively to all the modules and the profiled torch.nn.functionals are monkey patched. + + Args: + ignore_list (list, optional): the list of modules to ignore while profiling. Defaults to None. + """ + logger.info("Flops profiler started") + self.reset_profile() + _patch_functionals() + _patch_tensor_methods() + _patch_miscellaneous_operations() + + def register_module_hooks(module, ignore_list): + if ignore_list and type(module) in ignore_list: + return + + # if computing the flops of a module directly + if type(module) in MODULE_HOOK_MAPPING: + if not hasattr(module, "__flops_handle__"): + module.__flops_handle__ = module.register_forward_hook(MODULE_HOOK_MAPPING[type(module)]) + return + + # if computing the flops of the functionals in a module + def pre_hook(module, input): + module_flop_count.append([]) + module_mac_count.append([]) + + if not hasattr(module, "__pre_hook_handle__"): + module.__pre_hook_handle__ = module.register_forward_pre_hook(pre_hook) + + def post_hook(module, input, output): + if module_flop_count: + module.__flops__ += sum([elem[1] for elem in module_flop_count[-1]]) + module_flop_count.pop() + module.__macs__ += sum([elem[1] for elem in module_mac_count[-1]]) + module_mac_count.pop() + + if not hasattr(module, "__post_hook_handle__"): + module.__post_hook_handle__ = module.register_forward_hook(post_hook) + + def start_time_hook(module, input): + get_accelerator().synchronize() + module.__start_time__ = time.time() + + if not hasattr(module, "__start_time_hook_handle__"): + module.__start_time_hook_handle__ = module.register_forward_pre_hook(start_time_hook) + + def end_time_hook(module, input, output): + get_accelerator().synchronize() + module.__duration__ += time.time() - module.__start_time__ + + if not hasattr(module, "__end_time_hook_handle__"): + module.__end_time_hook_handle__ = module.register_forward_hook(end_time_hook) + + self.model.apply(partial(register_module_hooks, ignore_list=ignore_list)) + self.started = True + self.func_patched = True + + def stop_profile(self): + """Stop profiling. + + All torch.nn.functionals are restored to their originals. + """ + if self.started and self.func_patched: + _reload_functionals() + _reload_tensor_methods() + _reload_miscellaneous_operations() + self.func_patched = False + + def remove_profile_attrs(module): + if hasattr(module, "__pre_hook_handle__"): + module.__pre_hook_handle__.remove() + del module.__pre_hook_handle__ + if hasattr(module, "__post_hook_handle__"): + module.__post_hook_handle__.remove() + del module.__post_hook_handle__ + if hasattr(module, "__flops_handle__"): + module.__flops_handle__.remove() + del module.__flops_handle__ + if hasattr(module, "__start_time_hook_handle__"): + module.__start_time_hook_handle__.remove() + del module.__start_time_hook_handle__ + if hasattr(module, "__end_time_hook_handle__"): + module.__end_time_hook_handle__.remove() + del module.__end_time_hook_handle__ + + self.model.apply(remove_profile_attrs) + + def reset_profile(self): + """Resets the profiling. + + Adds or resets the extra attributes. + """ + + def get_param_count_and_ep(param): + """ + Return the number of parameters in the layer, whether the layer is an MoE layer, + and its expert parallelism size if so + """ + prefix = 'ep_size_' + offset = len(prefix) + expert_parallelism = 0 + if getattr(param, "group_name", "").startswith(prefix): + try: + expert_parallelism = int(param.group_name[offset:]) + except ValueError: + pass + return param.numel(), expert_parallelism, param.element_size() + + def add_or_reset_attrs(module): + module.__flops__ = 0 + module.__macs__ = 0 + module.__params__ = module.__expert_params__ = module.__model_expert_params__ = 0 + parameters = (get_param_count_and_ep(p) for p in module.parameters()) + for num_params, expert_parallelism, per_param_size in parameters: + params = num_params if not expert_parallelism else 0 + expert_params = num_params if expert_parallelism else 0 + # number of expert parameters taking into account other expert parallel groups + model_expert_params = num_params * expert_parallelism + module.__params__ += params + module.__expert_params__ += expert_params + module.__model_expert_params__ += model_expert_params + module.__start_time__ = 0 + module.__duration__ = 0 + + self.model.apply(add_or_reset_attrs) + + def end_profile(self): + """Ends profiling. + + The added attributes and handles are removed recursively on all the modules. + """ + if not self.started: + return + self.stop_profile() + self.started = False + + def remove_profile_attrs(module): + if hasattr(module, "__flops__"): + del module.__flops__ + if hasattr(module, "__macs__"): + del module.__macs__ + if hasattr(module, "__params__"): + del module.__params__ + if hasattr(module, "__expert_params__"): + del module.__expert_params__ + if hasattr(module, "__model_expert_params__"): + del module.__model_expert_params__ + if hasattr(module, "__start_time__"): + del module.__start_time__ + if hasattr(module, "__duration__"): + del module.__duration__ + + self.model.apply(remove_profile_attrs) + logger.info("Flops profiler finished") + + def get_total_flops(self, as_string=False): + """Returns the total flops of the model. + + Args: + as_string (bool, optional): whether to output the flops as string. Defaults to False. + + Returns: + The number of multiply-accumulate operations of the model forward pass. + """ + total_flops = get_module_flops(self.model) + return number_to_string(total_flops) if as_string else total_flops + + def get_total_macs(self, as_string=False): + """Returns the total MACs of the model. + + Args: + as_string (bool, optional): whether to output the flops as string. Defaults to False. + + Returns: + The number of multiply-accumulate operations of the model forward pass. + """ + total_macs = get_module_macs(self.model) + return macs_to_string(total_macs) if as_string else total_macs + + def get_total_duration(self, as_string=False): + """Returns the total duration of the model forward pass. + + Args: + as_string (bool, optional): whether to output the duration as string. Defaults to False. + + Returns: + The latency of the model forward pass. + """ + total_duration = get_module_duration(self.model) + return duration_to_string(total_duration) if as_string else total_duration + + def get_total_params(self, as_string=False): + """Returns the total number of parameters stored per rank. + + Args: + as_string (bool, optional): whether to output the parameters as string. Defaults to False. + + Returns: + The total number of parameters stored per rank. + """ + total_params = self.model.__expert_params__ + self.model.__params__ + return params_to_string(total_params) if as_string else total_params + + def is_expert_tensor_parallelism_enabled(self): + for _, module in self.model.named_modules(): + if isinstance(module, MoE) and hasattr(module, 'enable_expert_tensor_parallelism'): + return module.enable_expert_tensor_parallelism + return False + + def print_model_profile(self, profile_step=1, module_depth=-1, top_modules=1, detailed=True, output_file=None): + """Prints the model graph with the measured profile attached to each module. + + Args: + profile_step (int, optional): The global training step at which to profile. Note that warm up steps are needed for accurate time measurement. + module_depth (int, optional): The depth of the model to which to print the aggregated module information. When set to -1, it prints information from the top to the innermost modules (the maximum depth). + top_modules (int, optional): Limits the aggregated profile output to the number of top modules specified. + detailed (bool, optional): Whether to print the detailed model profile. + output_file (str, optional): Path to the output file. If None, the profiler prints to stdout. + """ + if not self.started: + return + import sys + import os.path + original_stdout = None + f = None + if output_file and output_file != "": + dir_path = os.path.dirname(os.path.abspath(output_file)) + if not os.path.exists(dir_path): + os.makedirs(dir_path) + original_stdout = sys.stdout + f = open(output_file, "w") + sys.stdout = f + + total_flops = self.get_total_flops() + total_macs = self.get_total_macs() + total_duration = self.get_total_duration() + total_params = self.get_total_params() + expert_tensor_parallelism = None # silence the linters + total_model_expert_params = total_model_nonexpert_params = 0 + if self.ds_engine: + total_model_nonexpert_params = self.model.__params__ * self.ds_engine.mp_world_size + if self.ds_engine.has_moe_layers: + expert_tensor_parallelism = self.ds_engine.mp_world_size if self.is_expert_tensor_parallelism_enabled( + ) else 1 + total_model_expert_params = self.model.__model_expert_params__ * expert_tensor_parallelism + + self.flops = total_flops + self.macs = total_macs + self.params = total_params + + print("\n-------------------------- DeepSpeed Flops Profiler --------------------------") + print(f'Profile Summary at step {profile_step}:') + print("Notations:\n" + "data parallel size (dp_size), model parallel size(mp_size),\n" + "number of parameters (params), number of multiply-accumulate operations(MACs),\n" + "number of floating-point operations (flops), floating-point operations per second (FLOPS),\n" + "fwd latency (forward propagation latency), bwd latency (backward propagation latency),\n" + "step (weights update latency), iter latency (sum of fwd, bwd and step latency)\n") + line_fmt = '{:<70} {:<8}' + if self.ds_engine: + print(line_fmt.format('world size: ', self.ds_engine.world_size)) + print(line_fmt.format('data parallel size: ', self.ds_engine.dp_world_size)) + print(line_fmt.format('model parallel size: ', self.ds_engine.mp_world_size)) + print(line_fmt.format('batch size per GPU: ', self.ds_engine.train_micro_batch_size_per_gpu())) + if self.ds_engine.has_moe_layers: + print(line_fmt.format('expert tensor parallelism enabled: ', expert_tensor_parallelism > 1)) + + print(line_fmt.format('params per GPU: ', params_to_string(total_params))) + if total_model_expert_params > 0: + print( + line_fmt.format('params of model: ', + params_to_string(total_model_nonexpert_params + total_model_expert_params))) + print(line_fmt.format(' non-expert params of model: ', params_to_string(total_model_nonexpert_params))) + print(line_fmt.format(' expert params of model: ', params_to_string(total_model_expert_params))) + else: + print( + line_fmt.format('params of model = params per GPU * mp_size: ', + params_to_string(total_model_nonexpert_params))) + + print(line_fmt.format('fwd MACs per GPU: ', macs_to_string(total_macs))) + + print(line_fmt.format('fwd flops per GPU: ', number_to_string(total_flops))) + + print( + line_fmt.format('fwd flops of model = fwd flops per GPU * mp_size: ', + number_to_string(total_flops * (self.ds_engine.mp_world_size if self.ds_engine else 1)))) + + fwd_latency = self.get_total_duration() + if self.ds_engine and self.ds_engine.wall_clock_breakdown(): + fwd_latency = self.ds_engine.timers(FORWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 + print(line_fmt.format('fwd latency: ', duration_to_string(fwd_latency))) + print( + line_fmt.format('fwd FLOPS per GPU = fwd flops per GPU / fwd latency: ', + flops_to_string(total_flops / fwd_latency))) + + if self.ds_engine and self.ds_engine.wall_clock_breakdown(): + bwd_factor = 2 + self.recompute_fwd_factor + bwd_latency = self.ds_engine.timers(BACKWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 + step_latency = self.ds_engine.timers(STEP_GLOBAL_TIMER).elapsed(False) / 1000.0 + print(line_fmt.format('bwd latency: ', duration_to_string(bwd_latency))) + print( + line_fmt.format(f'bwd FLOPS per GPU = {bwd_factor:g} * fwd flops per GPU / bwd latency: ', + flops_to_string(bwd_factor * total_flops / bwd_latency))) + print( + line_fmt.format( + f'fwd+bwd FLOPS per GPU = {bwd_factor + 1:g} * fwd flops per GPU / (fwd+bwd latency): ', + flops_to_string((bwd_factor + 1) * total_flops / (fwd_latency + bwd_latency)))) + + print(line_fmt.format('step latency: ', duration_to_string(step_latency))) + + iter_latency = fwd_latency + bwd_latency + step_latency + print(line_fmt.format('iter latency: ', duration_to_string(iter_latency))) + print( + line_fmt.format(f'FLOPS per GPU = {bwd_factor + 1:g} * fwd flops per GPU / iter latency: ', + flops_to_string((bwd_factor + 1) * total_flops / iter_latency))) + + samples_per_iter = self.ds_engine.train_micro_batch_size_per_gpu() * self.ds_engine.world_size + print(line_fmt.format('samples/second: ', round(samples_per_iter / iter_latency, DEFAULT_PRECISION))) + + def flops_repr(module): + params = module.__params__ + module.__expert_params__ + flops = get_module_flops(module) + macs = get_module_macs(module) + duration = get_module_duration(module) + items = [ + "{} = {:g}% Params".format( + params_to_string(params), + round(100 * params / total_params, DEFAULT_PRECISION) if total_params else 0), + "{} = {:g}% MACs".format(macs_to_string(macs), + round(100 * macs / total_macs, DEFAULT_PRECISION) if total_macs else 0), + "{} = {:g}% latency".format( + duration_to_string(duration), + round(100 * duration / total_duration, DEFAULT_PRECISION) if total_duration else 0), + flops_to_string(round(flops / duration, DEFAULT_PRECISION) if duration else 0), + ] + original_extra_repr = module.original_extra_repr() + if original_extra_repr: + items.append(original_extra_repr) + return ", ".join(items) + + def add_extra_repr(module): + flops_extra_repr = flops_repr.__get__(module) + if module.extra_repr != flops_extra_repr: + module.original_extra_repr = module.extra_repr + module.extra_repr = flops_extra_repr + assert module.extra_repr != module.original_extra_repr + + def del_extra_repr(module): + if hasattr(module, "original_extra_repr"): + module.extra_repr = module.original_extra_repr + del module.original_extra_repr + + self.model.apply(add_extra_repr) + + print("\n----------------------------- Aggregated Profile per GPU -----------------------------") + self.print_model_aggregated_profile(module_depth=module_depth, top_modules=top_modules) + + if detailed: + print("\n------------------------------ Detailed Profile per GPU ------------------------------") + print( + "Each module profile is listed after its name in the following order: \nparams, percentage of total params, MACs, percentage of total MACs, fwd latency, percentage of total fwd latency, fwd FLOPS" + ) + print( + "\nNote: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). They are not counted as submodules, thus not to be printed out. However they make up the difference between a parent's MACs (or latency) and the sum of its submodules'.\n2. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.\n3. The fwd latency listed in the top module's profile is directly captured at the module forward function in PyTorch, thus it's less than the fwd latency shown above which is captured in DeepSpeed.\n" + ) + print(self.model) + + self.model.apply(del_extra_repr) + + print("------------------------------------------------------------------------------") + + if output_file: + sys.stdout = original_stdout + f.close() + + def print_model_aggregated_profile(self, module_depth=-1, top_modules=1): + """Prints the names of the top top_modules modules in terms of aggregated time, flops, and parameters at depth module_depth. + + Args: + module_depth (int, optional): the depth of the modules to show. Defaults to -1 (the innermost modules). + top_modules (int, optional): the number of top modules to show. Defaults to 1. + """ + info = {} + if not hasattr(self.model, "__flops__"): + print("no __flops__ attribute in the model, call this function after start_profile and before end_profile") + return + + def walk_module(module, curr_depth, info): + if curr_depth not in info: + info[curr_depth] = {} + if module.__class__.__name__ not in info[curr_depth]: + info[curr_depth][module.__class__.__name__] = [ + 0, + 0, + 0, + ] # macs, params, time + info[curr_depth][module.__class__.__name__][0] += get_module_macs(module) + info[curr_depth][module.__class__.__name__][1] += module.__params__ + module.__expert_params__ + info[curr_depth][module.__class__.__name__][2] += get_module_duration(module) + has_children = len(module._modules.items()) != 0 + if has_children: + for child in module.children(): + walk_module(child, curr_depth + 1, info) + + walk_module(self.model, 0, info) + + depth = module_depth + if module_depth == -1: + depth = len(info) - 1 + + print(f'Top {top_modules} modules in terms of params, MACs or fwd latency at different model depths:') + + for d in range(depth): + num_items = min(top_modules, len(info[d])) + + sort_macs = { + k: macs_to_string(v[0]) + for k, v in sorted(info[d].items(), key=lambda item: item[1][0], reverse=True)[:num_items] + } + sort_params = { + k: params_to_string(v[1]) + for k, v in sorted(info[d].items(), key=lambda item: item[1][1], reverse=True)[:num_items] + } + sort_time = { + k: duration_to_string(v[2]) + for k, v in sorted(info[d].items(), key=lambda item: item[1][2], reverse=True)[:num_items] + } + + print(f"depth {d}:") + print(f" params - {sort_params}") + print(f" MACs - {sort_macs}") + print(f" fwd latency - {sort_time}") + + +def _prod(dims): + p = 1 + for v in dims: + p *= v + return p + + +def _linear_flops_compute(input, weight, bias=None): + out_features = weight.shape[0] + macs = input.numel() * out_features + return 2 * macs, macs + + +def _relu_flops_compute(input, inplace=False): + return input.numel(), 0 + + +def _prelu_flops_compute(input: Tensor, weight: Tensor): + return input.numel(), 0 + + +def _elu_flops_compute(input: Tensor, alpha: float = 1.0, inplace: bool = False): + return input.numel(), 0 + + +def _leaky_relu_flops_compute(input: Tensor, negative_slope: float = 0.01, inplace: bool = False): + return input.numel(), 0 + + +def _relu6_flops_compute(input: Tensor, inplace: bool = False): + return input.numel(), 0 + + +def _silu_flops_compute(input: Tensor, inplace: bool = False): + return input.numel(), 0 + + +def _gelu_flops_compute(input, **kwargs): + return input.numel(), 0 + + +def _pool_flops_compute(input, + kernel_size, + stride=None, + padding=0, + dilation=None, + ceil_mode=False, + count_include_pad=True, + divisor_override=None, + return_indices=None): + return input.numel(), 0 + + +def _conv_flops_compute(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): + assert weight.shape[1] * groups == input.shape[1] + + batch_size = input.shape[0] + in_channels = input.shape[1] + out_channels = weight.shape[0] + kernel_dims = list(weight.shape[2:]) + input_dims = list(input.shape[2:]) + + length = len(input_dims) + + strides = stride if type(stride) is tuple else (stride, ) * length + dilations = dilation if type(dilation) is tuple else (dilation, ) * length + if isinstance(padding, str): + if padding == 'valid': + paddings = (0, ) * length + elif padding == 'same': + paddings = () + for d, k in zip(dilations, kernel_dims): + total_padding = d * (k - 1) + paddings += (total_padding // 2, ) + elif isinstance(padding, tuple): + paddings = padding + else: + paddings = (padding, ) * length + + output_dims = [] + for idx, input_dim in enumerate(input_dims): + output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] * + (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1 + output_dims.append(output_dim) + + filters_per_channel = out_channels // groups + conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel + active_elements_count = batch_size * int(_prod(output_dims)) + overall_conv_macs = conv_per_position_macs * active_elements_count + overall_conv_flops = 2 * overall_conv_macs + + bias_flops = 0 + if bias is not None: + bias_flops = out_channels * active_elements_count + + return int(overall_conv_flops + bias_flops), int(overall_conv_macs) + + +def _conv_trans_flops_compute( + input, + weight, + bias=None, + stride=1, + padding=0, + output_padding=0, + groups=1, + dilation=1, +): + batch_size = input.shape[0] + in_channels = input.shape[1] + out_channels = weight.shape[1] + kernel_dims = list(weight.shape[2:]) + input_dims = list(input.shape[2:]) + + length = len(input_dims) + + paddings = padding if type(padding) is tuple else (padding, ) * length + strides = stride if type(stride) is tuple else (stride, ) * length + dilations = dilation if type(dilation) is tuple else (dilation, ) * length + + output_dims = [] + for idx, input_dim in enumerate(input_dims): + + output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] * + (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1 + output_dims.append(output_dim) + + paddings = padding if type(padding) is tuple else (padding, padding) + strides = stride if type(stride) is tuple else (stride, stride) + dilations = dilation if type(dilation) is tuple else (dilation, dilation) + + filters_per_channel = out_channels // groups + conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel + active_elements_count = batch_size * int(_prod(input_dims)) + overall_conv_macs = conv_per_position_macs * active_elements_count + overall_conv_flops = 2 * overall_conv_macs + + bias_flops = 0 + if bias is not None: + bias_flops = out_channels * batch_size * int(_prod(output_dims)) + + return int(overall_conv_flops + bias_flops), int(overall_conv_macs) + + +def _batch_norm_flops_compute( + input, + running_mean, + running_var, + weight=None, + bias=None, + training=False, + momentum=0.1, + eps=1e-05, +): + has_affine = weight is not None + if training: + # estimation + return input.numel() * (5 if has_affine else 4), 0 + flops = input.numel() * (2 if has_affine else 1) + return flops, 0 + + +def _layer_norm_flops_compute( + input: Tensor, + normalized_shape: List[int], + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + eps: float = 1e-5, +): + has_affine = weight is not None + # estimation + return input.numel() * (5 if has_affine else 4), 0 + + +def _group_norm_flops_compute(input: Tensor, + num_groups: int, + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + eps: float = 1e-5): + has_affine = weight is not None + # estimation + return input.numel() * (5 if has_affine else 4), 0 + + +def _instance_norm_flops_compute( + input: Tensor, + running_mean: Optional[Tensor] = None, + running_var: Optional[Tensor] = None, + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + use_input_stats: bool = True, + momentum: float = 0.1, + eps: float = 1e-5, +): + has_affine = weight is not None + # estimation + return input.numel() * (5 if has_affine else 4), 0 + + +def _upsample_flops_compute(*args, **kwargs): + input = args[0] + size = kwargs.get('size', None) + if size is None and len(args) > 1: + size = args[1] + + if size is not None: + if isinstance(size, tuple) or isinstance(size, list): + return int(_prod(size)), 0 + else: + return int(size), 0 + + scale_factor = kwargs.get('scale_factor', None) + if scale_factor is None and len(args) > 2: + scale_factor = args[2] + assert scale_factor is not None, "either size or scale_factor should be defined" + + flops = input.numel() + if isinstance(scale_factor, tuple) and len(scale_factor) == len(input): + flops *= int(_prod(scale_factor)) + else: + flops *= scale_factor**len(input) + return flops, 0 + + +def _softmax_flops_compute(input, dim=None, _stacklevel=3, dtype=None): + return input.numel(), 0 + + +def _embedding_flops_compute( + input, + weight, + padding_idx=None, + max_norm=None, + norm_type=2.0, + scale_grad_by_freq=False, + sparse=False, +): + return 0, 0 + + +def _dropout_flops_compute(input, p=0.5, training=True, inplace=False): + return 0, 0 + + +def _matmul_flops_compute(input, other, *, out=None): + """ + Count flops for the matmul operation. + """ + macs = _prod(input.shape) * other.shape[-1] + return 2 * macs, macs + + +def _addmm_flops_compute(input, mat1, mat2, *, beta=1, alpha=1, out=None): + """ + Count flops for the addmm operation. + """ + macs = _prod(mat1.shape) * mat2.shape[-1] + return 2 * macs + _prod(input.shape), macs + + +def _einsum_flops_compute(equation, *operands): + """ + Count flops for the einsum operation. + """ + equation = equation.replace(" ", "") + input_shapes = [o.shape for o in operands] + + # Re-map equation so that same equation with different alphabet + # representations will look the same. + letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys() + mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)} + equation = equation.translate(mapping) + + np_arrs = [np.zeros(s) for s in input_shapes] + optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1] + for line in optim.split("\n"): + if "optimized flop" in line.lower(): + flop = int(float(line.split(":")[-1])) + return flop, 0 + raise NotImplementedError("Unsupported einsum operation.") + + +def _einops_einsum_flops_compute(*args): + """ + Count flops for the einops.einsum operation. + """ + *operands, equation = args + input_shapes = [o.shape for o in operands] + + # Re-map equation so that same equation with different alphabet + # representations will look the same. + letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys() + mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)} + equation = equation.translate(mapping) + + np_arrs = [np.zeros(s) for s in input_shapes] + optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1] + for line in optim.split("\n"): + if "optimized flop" in line.lower(): + flop = int(float(line.split(":")[-1])) + return flop, 0 + + raise NotImplementedError("Unsupported einops.einsum operation.") + + +def _tensor_addmm_flops_compute(self, mat1, mat2, *, beta=1, alpha=1, out=None): + """ + Count flops for the tensor addmm operation. + """ + macs = _prod(mat1.shape) * mat2.shape[-1] + return 2 * macs + _prod(self.shape), macs + + +def _mul_flops_compute(input, other, *, out=None): + return _elementwise_flops_compute(input, other) + + +def _add_flops_compute(input, other, *, alpha=1, out=None): + return _elementwise_flops_compute(input, other) + + +def _elementwise_flops_compute(input, other): + if not torch.is_tensor(input): + if torch.is_tensor(other): + return _prod(other.shape), 0 + else: + return 1, 0 + elif not torch.is_tensor(other): + return _prod(input.shape), 0 + else: + dim_input = len(input.shape) + dim_other = len(other.shape) + max_dim = max(dim_input, dim_other) + + final_shape = [] + for i in range(max_dim): + in_i = input.shape[i] if i < dim_input else 1 + ot_i = other.shape[i] if i < dim_other else 1 + if in_i > ot_i: + final_shape.append(in_i) + else: + final_shape.append(ot_i) + flops = _prod(final_shape) + return flops, 0 + + +def _attn_flops_compute(q, k, v, *args, **kwargs): + """ + Count flops for the scaled_dot_product_attention operation. + """ + macs = _prod(q.shape) * k.shape[-2] + macs += _prod(q.shape[:-1]) * k.shape[-2] * v.shape[-1] + return 2 * macs, macs + + +def wrapFunc(func, funcFlopCompute): + oldFunc = func + name = func.__str__ + old_functions[name] = oldFunc + + def newFunc(*args, **kwds): + flops, macs = funcFlopCompute(*args, **kwds) + if module_flop_count: + module_flop_count[-1].append((name, flops)) + if module_mac_count and macs: + module_mac_count[-1].append((name, macs)) + return oldFunc(*args, **kwds) + + newFunc.__str__ = func.__str__ + + return newFunc + + +def _patch_functionals(): + # FC + F.linear = wrapFunc(F.linear, _linear_flops_compute) + + # convolutions + F.conv1d = wrapFunc(F.conv1d, _conv_flops_compute) + F.conv2d = wrapFunc(F.conv2d, _conv_flops_compute) + F.conv3d = wrapFunc(F.conv3d, _conv_flops_compute) + + # conv transposed + F.conv_transpose1d = wrapFunc(F.conv_transpose1d, _conv_trans_flops_compute) + F.conv_transpose2d = wrapFunc(F.conv_transpose2d, _conv_trans_flops_compute) + F.conv_transpose3d = wrapFunc(F.conv_transpose3d, _conv_trans_flops_compute) + + # activations + F.relu = wrapFunc(F.relu, _relu_flops_compute) + F.prelu = wrapFunc(F.prelu, _prelu_flops_compute) + F.elu = wrapFunc(F.elu, _elu_flops_compute) + F.leaky_relu = wrapFunc(F.leaky_relu, _leaky_relu_flops_compute) + F.relu6 = wrapFunc(F.relu6, _relu6_flops_compute) + if hasattr(F, "silu"): + F.silu = wrapFunc(F.silu, _silu_flops_compute) + F.gelu = wrapFunc(F.gelu, _gelu_flops_compute) + + # Normalizations + F.batch_norm = wrapFunc(F.batch_norm, _batch_norm_flops_compute) + F.layer_norm = wrapFunc(F.layer_norm, _layer_norm_flops_compute) + F.instance_norm = wrapFunc(F.instance_norm, _instance_norm_flops_compute) + F.group_norm = wrapFunc(F.group_norm, _group_norm_flops_compute) + + # poolings + F.avg_pool1d = wrapFunc(F.avg_pool1d, _pool_flops_compute) + F.avg_pool2d = wrapFunc(F.avg_pool2d, _pool_flops_compute) + F.avg_pool3d = wrapFunc(F.avg_pool3d, _pool_flops_compute) + F.max_pool1d = wrapFunc(F.max_pool1d, _pool_flops_compute) + F.max_pool2d = wrapFunc(F.max_pool2d, _pool_flops_compute) + F.max_pool3d = wrapFunc(F.max_pool3d, _pool_flops_compute) + F.adaptive_avg_pool1d = wrapFunc(F.adaptive_avg_pool1d, _pool_flops_compute) + F.adaptive_avg_pool2d = wrapFunc(F.adaptive_avg_pool2d, _pool_flops_compute) + F.adaptive_avg_pool3d = wrapFunc(F.adaptive_avg_pool3d, _pool_flops_compute) + F.adaptive_max_pool1d = wrapFunc(F.adaptive_max_pool1d, _pool_flops_compute) + F.adaptive_max_pool2d = wrapFunc(F.adaptive_max_pool2d, _pool_flops_compute) + F.adaptive_max_pool3d = wrapFunc(F.adaptive_max_pool3d, _pool_flops_compute) + + # upsample + F.upsample = wrapFunc(F.upsample, _upsample_flops_compute) + F.interpolate = wrapFunc(F.interpolate, _upsample_flops_compute) + + # softmax + F.softmax = wrapFunc(F.softmax, _softmax_flops_compute) + + # embedding + F.embedding = wrapFunc(F.embedding, _embedding_flops_compute) + + # attn - scaled_dot_product_attention added in torch 2.0+ + if required_torch_version(min_version=2.0): + F.scaled_dot_product_attention = wrapFunc(F.scaled_dot_product_attention, _attn_flops_compute) + + +def _patch_tensor_methods(): + torch.matmul = wrapFunc(torch.matmul, _matmul_flops_compute) + torch.Tensor.matmul = wrapFunc(torch.Tensor.matmul, _matmul_flops_compute) + torch.Tensor.__matmul__ = wrapFunc(torch.Tensor.__matmul__, _matmul_flops_compute) + torch.mm = wrapFunc(torch.mm, _matmul_flops_compute) + torch.Tensor.mm = wrapFunc(torch.Tensor.mm, _matmul_flops_compute) + torch.bmm = wrapFunc(torch.bmm, _matmul_flops_compute) + torch.Tensor.bmm = wrapFunc(torch.Tensor.bmm, _matmul_flops_compute) + + torch.addmm = wrapFunc(torch.addmm, _addmm_flops_compute) + torch.Tensor.addmm = wrapFunc(torch.Tensor.addmm, _tensor_addmm_flops_compute) + + torch.mul = wrapFunc(torch.mul, _mul_flops_compute) + torch.Tensor.mul = wrapFunc(torch.Tensor.mul, _mul_flops_compute) + + torch.add = wrapFunc(torch.add, _add_flops_compute) + torch.Tensor.add = wrapFunc(torch.Tensor.add, _add_flops_compute) + + torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) + + torch.baddbmm = wrapFunc(torch.baddbmm, _tensor_addmm_flops_compute) + + +def _patch_miscellaneous_operations(): + einops.einsum = wrapFunc(einops.einsum, _einops_einsum_flops_compute) + + +def _reload_functionals(): + # torch.nn.functional does not support importlib.reload() + F.linear = old_functions[F.linear.__str__] + F.conv1d = old_functions[F.conv1d.__str__] + F.conv2d = old_functions[F.conv2d.__str__] + F.conv3d = old_functions[F.conv3d.__str__] + F.conv_transpose1d = old_functions[F.conv_transpose1d.__str__] + F.conv_transpose2d = old_functions[F.conv_transpose2d.__str__] + F.conv_transpose3d = old_functions[F.conv_transpose3d.__str__] + F.relu = old_functions[F.relu.__str__] + F.prelu = old_functions[F.prelu.__str__] + F.elu = old_functions[F.elu.__str__] + F.leaky_relu = old_functions[F.leaky_relu.__str__] + F.relu6 = old_functions[F.relu6.__str__] + if hasattr(F, "silu"): + F.silu = old_functions[F.silu.__str__] + F.gelu = old_functions[F.gelu.__str__] + F.batch_norm = old_functions[F.batch_norm.__str__] + F.layer_norm = old_functions[F.layer_norm.__str__] + F.instance_norm = old_functions[F.instance_norm.__str__] + F.group_norm = old_functions[F.group_norm.__str__] + F.avg_pool1d = old_functions[F.avg_pool1d.__str__] + F.avg_pool2d = old_functions[F.avg_pool2d.__str__] + F.avg_pool3d = old_functions[F.avg_pool3d.__str__] + F.max_pool1d = old_functions[F.max_pool1d.__str__] + F.max_pool2d = old_functions[F.max_pool2d.__str__] + F.max_pool3d = old_functions[F.max_pool3d.__str__] + F.adaptive_avg_pool1d = old_functions[F.adaptive_avg_pool1d.__str__] + F.adaptive_avg_pool2d = old_functions[F.adaptive_avg_pool2d.__str__] + F.adaptive_avg_pool3d = old_functions[F.adaptive_avg_pool3d.__str__] + F.adaptive_max_pool1d = old_functions[F.adaptive_max_pool1d.__str__] + F.adaptive_max_pool2d = old_functions[F.adaptive_max_pool2d.__str__] + F.adaptive_max_pool3d = old_functions[F.adaptive_max_pool3d.__str__] + F.upsample = old_functions[F.upsample.__str__] + F.interpolate = old_functions[F.interpolate.__str__] + F.softmax = old_functions[F.softmax.__str__] + F.embedding = old_functions[F.embedding.__str__] + + +def _reload_tensor_methods(): + torch.matmul = old_functions[torch.matmul.__str__] + torch.Tensor.matmul = old_functions[torch.Tensor.matmul.__str__] + torch.mm = old_functions[torch.mm.__str__] + torch.Tensor.mm = old_functions[torch.Tensor.mm.__str__] + torch.bmm = old_functions[torch.matmul.__str__] + torch.Tensor.bmm = old_functions[torch.Tensor.bmm.__str__] + torch.addmm = old_functions[torch.addmm.__str__] + torch.Tensor.addmm = old_functions[torch.Tensor.addmm.__str__] + torch.mul = old_functions[torch.mul.__str__] + torch.Tensor.mul = old_functions[torch.Tensor.mul.__str__] + torch.add = old_functions[torch.add.__str__] + torch.Tensor.add = old_functions[torch.Tensor.add.__str__] + + torch.einsum = old_functions[torch.einsum.__str__] + + torch.baddbmm = old_functions[torch.baddbmm.__str__] + + +def _reload_miscellaneous_operations(): + einops.einsum = old_functions[einops.einsum.__str__] + + +def _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size): + gates_size = w_ih.shape[0] + # matrix matrix mult ih state and internal state + flops += 2 * w_ih.shape[0] * w_ih.shape[1] - gates_size + # matrix matrix mult hh state and internal state + flops += 2 * w_hh.shape[0] * w_hh.shape[1] - gates_size + if isinstance(rnn_module, (nn.RNN, nn.RNNCell)): + # add both operations + flops += rnn_module.hidden_size + elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)): + # hadamard of r + flops += rnn_module.hidden_size + # adding operations from both states + flops += rnn_module.hidden_size * 3 + # last two hadamard _product and add + flops += rnn_module.hidden_size * 3 + elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)): + # adding operations from both states + flops += rnn_module.hidden_size * 4 + # two hadamard _product and add for C state + flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size + # final hadamard + flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size + return flops + + +def _rnn_forward_hook(rnn_module, input, output): + flops = 0 + # input is a tuple containing a sequence to process and (optionally) hidden state + inp = input[0] + batch_size = inp.shape[0] + seq_length = inp.shape[1] + num_layers = rnn_module.num_layers + + for i in range(num_layers): + w_ih = rnn_module.__getattr__("weight_ih_l" + str(i)) + w_hh = rnn_module.__getattr__("weight_hh_l" + str(i)) + if i == 0: + input_size = rnn_module.input_size + else: + input_size = rnn_module.hidden_size + flops = _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size) + if rnn_module.bias: + b_ih = rnn_module.__getattr__("bias_ih_l" + str(i)) + b_hh = rnn_module.__getattr__("bias_hh_l" + str(i)) + flops += b_ih.shape[0] + b_hh.shape[0] + + flops *= batch_size + flops *= seq_length + if rnn_module.bidirectional: + flops *= 2 + rnn_module.__flops__ += int(flops) + + +def _rnn_cell_forward_hook(rnn_cell_module, input, output): + flops = 0 + inp = input[0] + batch_size = inp.shape[0] + w_ih = rnn_cell_module.__getattr__("weight_ih") + w_hh = rnn_cell_module.__getattr__("weight_hh") + input_size = inp.shape[1] + flops = _rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size) + if rnn_cell_module.bias: + b_ih = rnn_cell_module.__getattr__("bias_ih") + b_hh = rnn_cell_module.__getattr__("bias_hh") + flops += b_ih.shape[0] + b_hh.shape[0] + + flops *= batch_size + rnn_cell_module.__flops__ += int(flops) + + +MODULE_HOOK_MAPPING = { + # RNN + nn.RNN: _rnn_forward_hook, + nn.GRU: _rnn_forward_hook, + nn.LSTM: _rnn_forward_hook, + nn.RNNCell: _rnn_cell_forward_hook, + nn.LSTMCell: _rnn_cell_forward_hook, + nn.GRUCell: _rnn_cell_forward_hook, +} + + +def macs_to_string(macs, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(macs, units=units, precision=precision)}MACs" + + +def number_to_string(num, units=None, precision=DEFAULT_PRECISION): + if units is None: + if num >= 1e12: + magnitude, units = 1e12, "T" + elif num >= 1e9: + magnitude, units = 1e9, "G" + elif num >= 1e6: + magnitude, units = 1e6, "M" + elif num >= 1e3: + magnitude, units = 1e3, "K" + elif num >= 1 or num == 0: + magnitude, units = 1, "" + elif num >= 1e-3: + magnitude, units = 1e-3, "m" + else: + magnitude, units = 1e-6, "u" + else: + if units == "T": + magnitude = 1e12 + elif units == "G": + magnitude = 1e9 + elif units == "M": + magnitude = 1e6 + elif units == "K": + magnitude = 1e3 + elif units == "m": + magnitude = 1e-3 + elif units == "u": + magnitude = 1e-6 + else: + magnitude = 1 + return f"{round(num / magnitude, precision):g} {units}" + + +def flops_to_string(flops, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(flops, units=units, precision=precision)}FLOPS" + + +def bytes_to_string(b, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(b, units=units, precision=precision)}B" + + +def params_to_string(params_num, units=None, precision=DEFAULT_PRECISION): + units = units.replace("B", "G") if units else units + return number_to_string(params_num, units=units, precision=precision).replace("G", "B").strip() + + +def duration_to_string(duration, units=None, precision=DEFAULT_PRECISION): + return f"{number_to_string(duration, units=units, precision=precision)}s" + + + # can not iterate over all submodules using self.model.modules() + # since modules() returns duplicate modules only once +def get_module_flops(module): + sum = module.__flops__ + # iterate over immediate children modules + for child in module.children(): + sum += get_module_flops(child) + return sum + + +def get_module_macs(module): + sum = module.__macs__ + # iterate over immediate children modules + for child in module.children(): + sum += get_module_macs(child) + return sum + + +def get_module_duration(module): + duration = module.__duration__ + if duration == 0: # e.g. ModuleList + for m in module.children(): + duration += get_module_duration(m) + return duration + + +def get_model_profile(model, + input_shape=None, + args=[], + kwargs={}, + print_profile=True, + detailed=True, + module_depth=-1, + top_modules=1, + warm_up=1, + as_string=True, + output_file=None, + ignore_modules=None, + mode='forward'): + """Returns the total floating-point operations, MACs, and parameters of a model. + + Example: + + .. code-block:: python + + model = torchvision.models.alexnet() + batch_size = 256 + flops, macs, params = get_model_profile(model=model, input_shape=(batch_size, 3, 224, 224))) + + Args: + model ([torch.nn.Module]): the PyTorch model to be profiled. + input_shape (tuple): input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument. + args (list): list of positional arguments to the model. + kwargs (dict): dictionary of keyword arguments to the model. + print_profile (bool, optional): whether to print the model profile. Defaults to True. + detailed (bool, optional): whether to print the detailed model profile. Defaults to True. + module_depth (int, optional): the depth into the nested modules. Defaults to -1 (the inner most modules). + top_modules (int, optional): the number of top modules to print in the aggregated profile. Defaults to 3. + warm_up (int, optional): the number of warm-up steps before measuring the latency of each module. Defaults to 1. + as_string (bool, optional): whether to print the output as string. Defaults to True. + output_file (str, optional): path to the output file. If None, the profiler prints to stdout. + ignore_modules ([type], optional): the list of modules to ignore during profiling. Defaults to None. + + Returns: + The number of floating-point operations, multiply-accumulate operations (MACs), and parameters in the model. + """ + assert isinstance(model, nn.Module), "model must be a PyTorch module" + prof = FlopsProfiler(model) + model.eval() + + if input_shape is not None: + assert type(input_shape) is tuple, "input_shape must be a tuple" + assert len(input_shape) >= 1, "input_shape must have at least one element" + try: + input = torch.ones(()).new_empty( + (*input_shape, ), + dtype=next(model.parameters()).dtype, + device=next(model.parameters()).device, + ) + except StopIteration: + input = torch.ones(()).new_empty((*input_shape, )) + + args = [input] + assert (len(args) > 0) or (len(kwargs) > 0), "args and/or kwargs must be specified if input_shape is None" + + logger.info("Flops profiler warming-up...") + for _ in range(warm_up): + if kwargs: + if mode == 'forward': + _ = model(*args, **kwargs) + if mode == 'generate': + _ = model.generate(*args, **kwargs) + else: + if mode == 'forward': + _ = model(*args) + if mode == 'generate': + _ = model.generate(*args) + prof.start_profile(ignore_list=ignore_modules) + + if kwargs: + if mode == 'forward': + _ = model(*args, **kwargs) + if mode == 'generate': + _ = model.generate(*args, **kwargs) + else: + if mode == 'forward': + _ = model(*args) + if mode == 'generate': + _ = model.generate(*args) + + flops = prof.get_total_flops() + macs = prof.get_total_macs() + params = prof.get_total_params() + if print_profile: + prof.print_model_profile(profile_step=warm_up, + module_depth=module_depth, + top_modules=top_modules, + detailed=detailed, + output_file=output_file) + + prof.end_profile() + if as_string: + return number_to_string(flops), macs_to_string(macs), params_to_string(params) + + return flops, macs, params diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..28c35a798ae7e98225145a682e6db93852eb0651 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9028ea0a939cd0badec1b73962232f26db3d9521 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4924204ebf8725f2a636dff296870a156b23338 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8436ae506430f8dae081bedf1ac482a7f2bcaa0 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c685f652cf686cbdbbd25b714afc2083e2c562b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eab855bf2c6e555435fdb81454227a570549fda9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a01c7bc42b2a907ab4868dde3cb5444d59ce5ead Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f85219de302992fecdb840f4d8208ff54a5eb7cf Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25fb2dbcc495f768856b7c9e223dc42d478ef24a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93de61abd7af1dab2e9c309c5c1e11b0ee952dc5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..224c3577e256a3aa762389f7184a484b7f078491 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4af56b93efeaa6ed7e8cdd8393534cb87927dcb Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a54d7b9d0fdc1acb3e32ef9ff415e4ab0fc3440d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3754051f51081980ac5f4ebd7a34b4cea897cf8d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aba8beb5033680c76fec95a2fd97f845b26cb16d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45878a5db828995f374757b8ff84d45b4e51cc5b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f875733808f703b400d00642de21e5156a89fda Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e6b462526e8a7943fc15435968f6f8469dc77ef Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b5f1288680195844cd776ad5ba3774645ae5458 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac96e74939596c0f08cb5c2985eadbcf98443ab1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e3447b4b57fa2c794753f5e4577409b348dc6fd Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..a341dffdf692cab4b49a0b7a5e6c7bb8c2658672 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py @@ -0,0 +1,30 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + + +class CheckpointEngine(object): + + # init checkpoint engine for save/load + def __init__(self, config_params=None): + pass + + def create(self, tag): + # create checkpoint on give tag for save/load. + pass + + def makedirs(self, path, exist_ok=False): + os.makedirs(path, exist_ok=exist_ok) + + def save(self, state_dict, path: str): + pass + + def load(self, path: str, map_location=None): + pass + + def commit(self, tag): + # to tell checkpoint services if all files are ready. + pass diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..e834bf0d22d790781cf0be7de6054d22e15a7f3f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py @@ -0,0 +1,107 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch +import torch_nebula + +from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ + CheckpointEngine +from deepspeed.utils import logger, log_dist +from deepspeed.nebula.constants import * + + +def _get_tag_from_path(path): + return os.path.basename(os.path.dirname(path)) + + +class NebulaCheckpointEngine(CheckpointEngine): + + def __init__(self, config_params=None): + super().__init__(config_params) + self.checkpoint = None + self.tag_flag = None + self.enable_nebula_load = config_params.enable_nebula_load + self.nebula_load_path = config_params.load_path + if self.nebula_load_path is None: + self.nebula_load_path = config_params.persistent_storage_path + + nebula_config_params = { + NEBULA_PERSISTENT_STORAGE_PATH: config_params.persistent_storage_path, + NEBULA_PERSISTENT_TIME_INTERVAL: config_params.persistent_time_interval, + NEBULA_NUM_OF_VERSION_IN_RETENTION: config_params.num_of_version_in_retention, + } + torch_nebula.init(**nebula_config_params) + + def create(self, tag): + log_dist(f"[Nebula] Start Checkpoint for tag:{tag}", ranks=[0]) + # -2 means: customer needs to explicitly tell nebula + # current checkpoint is complete by commit method. + self.checkpoint = torch_nebula.Checkpoint(tag, -2) + + def save(self, state_dict, path: str): + log_dist(f"[Nebula] Create dummy files for loading.") + torch.save("", path) + + tag = _get_tag_from_path(path) + partition_name = os.path.basename(path) + logger.info(f"[Nebula] Saving {partition_name} under tag {tag}...") + self.checkpoint.save(partition_name, state_dict) + logger.info(f"[Nebula] Saved {partition_name} under tag {tag}.") + return None + + def load(self, path: str, map_location=None): + tag = _get_tag_from_path(path) + first_load_flag = self.tag_flag is None or self.tag_flag == tag + if not self.enable_nebula_load and first_load_flag: + self.tag_flag = tag + logger.info(f"[Nebula] Disable nebula load. Loading checkpoint from {path} ...") + partition = torch.load(path, map_location=map_location, weights_only=False) + logger.info(f"[Nebula] Disable nebula load. Loaded checkpoint from {path} .") + return partition + + partition_name = os.path.basename(path) + logger.info(f"[Nebula] Loading {path} under tag {tag} from nebula path {self.nebula_load_path}...") + + checkpoint = None + if tag in (None, 'latest', 'latest_universal'): + # In some cases, there is the inconsistent tag between deepspeed metadata (latest file) + # and nebula metadata, will lead to the failure on loading with deepspeed tag. Then we + # will try to load the valid latest checkpoint from nebula(tier3 > tier1). So, in summary + # when met failure loading for given tag, the loading priority would be like: + # nebula tier3 latest > nebula tier1 latest. + checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path) + else: + checkpoint = torch_nebula.get_checkpoint(tag=tag, persist_path=self.nebula_load_path) + + if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''): + logger.info( + f"Unable to find valid checkpoint tag:{tag} from Nebula, try to get latest checkpoint again from nebula {self.nebula_load_path} path!" + ) + # nebula tier3 latest + checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path) + if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''): + logger.info( + f"Unable to find latest checkpoint from Nebula tier3, try to get latest checkpoint again from nebula tier1 path!" + ) + # nebula tier1 latest + checkpoint = torch_nebula.get_latest_checkpoint() + logger.warning(f"Unable to find valid checkpoint from Nebula under tag:{tag}.") + return None + + tag = checkpoint.tag + self.tag_flag = -1 + partition = checkpoint.load(partition_name, map_location=map_location) + logger.info(f"[Nebula] Loaded {path} under tag {tag} from {self.nebula_load_path}.") + return partition + + def commit(self, tag): + # nebula commit will be call when all files under give tag are ready to be persisted in the async way. + logger.info(f"[Nebula] all files for {tag} are saved in tier1. It is ready to start persisting") + commit_rls = self.checkpoint.commit() + if not commit_rls: + logger.error(f"[Nebula] failed to commit the checkpoint, please check the log.") + return False + return commit_rls diff --git a/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..076c638532ad37faea448f2423b5f319743ed38c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.utils import logger, log_dist +from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ + CheckpointEngine + + +class TorchCheckpointEngine(CheckpointEngine): + + def __init__(self, config_params=None): + super().__init__(config_params) + + def create(self, tag): + log_dist(f"[Torch] Checkpoint {tag} is about to be saved!", ranks=[0]) + + def save(self, state_dict, path: str): + logger.info(f"[Torch] Saving {path}...") + torch.save(state_dict, path) + logger.info(f"[Torch] Saved {path}.") + return None + + def load(self, path: str, map_location=None): + logger.info(f"[Torch] Loading checkpoint from {path}...") + partition = torch.load(path, map_location=map_location, weights_only=False) + logger.info(f"[Torch] Loaded checkpoint from {path}.") + return partition + + def commit(self, tag): + logger.info(f"[Torch] Checkpoint {tag} is ready now!") + return True diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/comm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/comm/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ff995c52276861552c99fa6e44f3995786616c3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d8de96909e0c68772b26d5fc2093fd70ba04dd6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8cd3dff46060d1e34be33457a547ae2af715e56 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8db5dbdd588bf72d91fb4d723bf16c84f3942003 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d388cef1b34077858f27cb9fb05274f57dae1ef7 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b772c1db87854880ecc5a2f2383dcedf9e75b6e7 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/coalesced_collectives.py b/lib/python3.12/site-packages/deepspeed/runtime/comm/coalesced_collectives.py new file mode 100644 index 0000000000000000000000000000000000000000..2fadce52222cda680600253ae840f84b89bda7ed --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/comm/coalesced_collectives.py @@ -0,0 +1,218 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +batched collective operations for overhead amortization and better +bandwidth utilization +""" + +import math +from typing import List, Any +import torch +from torch import Tensor +from deepspeed import comm as dist +from deepspeed.comm import ProcessGroup, all_to_all_single +from deepspeed.accelerator import get_accelerator +from deepspeed.utils import instrument_w_nvtx +from deepspeed.ops import op_builder +from deepspeed.utils import logger + + +def _torch_reduce_scatter_fn(input_tensor: Tensor, output_tensor: Tensor, group=None, async_op=False, prof=False): + return instrument_w_nvtx(dist.reduce_scatter_fn)(output_tensor, input_tensor, group=group, async_op=False) + + +quantizer_module = None + + +@instrument_w_nvtx +@torch.no_grad() +def all_to_all_quant_reduce(tensors: List[Tensor], groups: {}) -> List[Tensor]: + global quantizer_module + if quantizer_module is None: + quantizer_module = op_builder.QuantizerBuilder().load() + local_world_size = get_accelerator().device_count() + global_world_size = dist.get_world_size() + num_nodes = global_world_size // local_world_size + this_rank = dist.get_rank() + intra_idx = int(this_rank / local_world_size) + inter_idx = this_rank % local_world_size + output_lst: List[Tensor] = [None] * len(tensors) + for idx, tensor in enumerate(tensors): + if tensor.dim() == 1: + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + elif tensor.numel() % (2 * global_world_size) != 0: + # Due to the constraint of 2-stage all-to-all, the input tensor must be divisible by 2 * global_world_size + # Otherwise, all-to-all cannot be performed because of shape mismatch. + # See more at https://github.com/deepspeedai/DeepSpeed/pull/5056 + logger.warning( + f"qgZ falls back to reduce_scatter because tensor size = {tensor.numel()} is not divisible by (2 * global_world_size) = {2 * global_world_size}. Please consider allocating a new world to enable qgZ" + ) + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + else: + intra_quant_group = max(tensor.shape[0], tensor.shape[1], global_world_size) + + inter_quant_group = intra_quant_group // local_world_size + intra_quant_int4, intra_q_scales = quantizer_module.swizzle_quant(tensor, intra_quant_group, 4, + quantizer_module.Symmetric, 1, num_nodes, + local_world_size) + local_output = torch.empty_like(intra_quant_int4) + scale_output = torch.empty_like(intra_q_scales) + all_to_all_single(local_output, intra_quant_int4, group=groups[f'local_{intra_idx}']) + all_to_all_single(scale_output, intra_q_scales, group=groups[f'local_{intra_idx}']) + global_input_tensor, global_scales = quantizer_module.quantized_reduction( + local_output, scale_output, intra_quant_group, inter_quant_group, 4, quantizer_module.Symmetric, + local_world_size) + global_output = torch.empty_like(global_input_tensor) + global_scale_output = torch.empty_like(global_scales) + all_to_all_single(global_output, global_input_tensor, group=groups[f'global_{inter_idx}']) + all_to_all_single(global_scale_output, global_scales, group=groups[f'global_{inter_idx}']) + final_output = quantizer_module.dequantize(global_output, global_scale_output, global_scale_output.numel(), + 4, quantizer_module.Symmetric) + assert final_output.numel( + ) % num_nodes == 0, f"final_output.numel()={final_output.numel()} is not divisible by num_nodes={num_nodes}" + output_lst[idx] = (sum(list(final_output.chunk(num_nodes))) / num_nodes).view(-1) + return output_lst + + +@instrument_w_nvtx +@torch.no_grad() +def all_to_all_loco_quant_reduce( + params: List[Tensor], + groups: {}, + loco_param: Any = None, +) -> List[Tensor]: + global quantizer_module + global loco_idx + if quantizer_module is None: + quantizer_module = op_builder.QuantizerBuilder().load() + local_world_size = get_accelerator().device_count() + global_world_size = dist.get_world_size() + num_nodes = global_world_size // local_world_size + this_rank = dist.get_rank() + intra_idx = int(this_rank / local_world_size) + inter_idx = this_rank % local_world_size + output_lst: List[Tensor] = [None] * len(params) + for idx, p in enumerate(params): + tensor = p.grad + if tensor.dim() == 1: + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + elif tensor.numel() % (2 * global_world_size) != 0: + # Due to the constraint of 2-stage all-to-all, the input tensor must be divisible by 2 * global_world_size + # Otherwise, all-to-all cannot be performed because of shape mismatch. + # See more at https://github.com/deepspeedai/DeepSpeed/pull/5056 + logger.warning( + f"qgZ falls back to reduce_scatter because tensor size = {tensor.numel()} is not divisible by (2 * global_world_size) = {2 * global_world_size}. Please consider allocating a new world to enable qgZ" + ) + output_lst[idx] = reduce_scatter_coalesced([tensor])[0] + else: + err_beta = loco_param['err_beta'] + reset_T = loco_param['reset_T'] + if not hasattr(p, 'intra_ef_buf') or loco_idx > reset_T: + loco_idx = 0 + intra_err = torch.zeros_like(p.grad) + inter_err = torch.zeros(tensor.numel() // local_world_size, device=tensor.device, dtype=tensor.dtype) + else: + intra_err = quantizer_module.dequantize(p.intra_ef_buf[0], p.intra_ef_buf[1], + p.intra_ef_buf[1].numel(), 8, quantizer_module.Symmetric) + inter_err = quantizer_module.dequantize(p.inter_ef_buf[0], p.inter_ef_buf[1], + p.inter_ef_buf[1].numel(), 8, quantizer_module.Symmetric) + + intra_quant_group = max(tensor.shape[0], tensor.shape[1], global_world_size) + inter_quant_group = intra_quant_group // local_world_size + intra_quant_int4, intra_q_scales = quantizer_module.loco_swizzle_quant(tensor, intra_err, err_beta, + intra_quant_group, 4, + quantizer_module.Symmetric, 1, + num_nodes, local_world_size) + local_output = torch.empty_like(intra_quant_int4) + scale_output = torch.empty_like(intra_q_scales) + all_to_all_single(local_output, intra_quant_int4, group=groups[f'local_{intra_idx}']) + all_to_all_single(scale_output, intra_q_scales, group=groups[f'local_{intra_idx}']) + + p.intra_ef_buf = quantizer_module.quantize(intra_err, intra_quant_group, 8, quantizer_module.Symmetric) + + global_input_tensor, global_scales = quantizer_module.loco_quantized_reduction( + local_output, scale_output, inter_err, err_beta, intra_quant_group, inter_quant_group, 4, + quantizer_module.Symmetric, local_world_size) + + global_output = torch.empty_like(global_input_tensor) + global_scale_output = torch.empty_like(global_scales) + all_to_all_single(global_output, global_input_tensor, group=groups[f'global_{inter_idx}']) + all_to_all_single(global_scale_output, global_scales, group=groups[f'global_{inter_idx}']) + + p.inter_ef_buf = quantizer_module.quantize(inter_err, inter_quant_group, 8, quantizer_module.Symmetric) + + final_output = quantizer_module.dequantize(global_output, global_scale_output, global_scale_output.numel(), + 4, quantizer_module.Symmetric) + assert final_output.numel( + ) % num_nodes == 0, f"final_output.numel()={final_output.numel()} is not divisible by num_nodes={num_nodes}" + output_lst[idx] = (sum(list(final_output.chunk(num_nodes))) / num_nodes).view(-1) + loco_idx += 1 + + return output_lst + + +@instrument_w_nvtx +@torch.no_grad() +def reduce_scatter_coalesced( + tensors: List[Tensor], + group: ProcessGroup = None, +) -> List[Tensor]: + """simultaneously reduce-scatter a list of tensors - this can be done more + efficiently than individual reduce scatter calls + TODO. see if PyTorch team wants a c++ version of this for ProcessGroupNCCL + """ + this_rank = dist.get_rank(group) + world_sz = dist.get_world_size(group) + + partition_lst_for_each_tensor = [None] * len(tensors) + for tensor_idx, tensor in enumerate(tensors): + flattened_tensor = tensor.view(-1) + chunk_sz = math.ceil(tensor.numel() / world_sz) + partition_lst_for_each_tensor[tensor_idx] = [ + flattened_tensor[rank * chunk_sz:rank * chunk_sz + chunk_sz] for rank in range(0, world_sz) + ] + + padded_partition_sz_for_each_tensor = tuple(math.ceil(t.numel() / world_sz) for t in tensors) + + if len(tensors) == 1 and tensors[0].numel() % world_sz == 0: + # if there's only one tensor being reduced and we don't need to pad + # we have an opportunity to avoid a memory allocation + tensor_partition_flat_buffer = tensors[0].view(-1) + else: + # interleave tensor partitions such that the correct reduced partitions of each tensor + # end up at each rank + tensor_partitions_lst_with_padding = [] + for rank in range(world_sz): + for tensor_idx in range(len(tensors)): + # add tensor content + tensor_chunk = partition_lst_for_each_tensor[tensor_idx][rank] + tensor_partitions_lst_with_padding.append(tensor_chunk) + + # add padding if necessary + padding_sz = padded_partition_sz_for_each_tensor[tensor_idx] - tensor_chunk.numel() + if padding_sz > 0: + tensor_partitions_lst_with_padding.append( + torch.empty(padding_sz, dtype=tensor_chunk.dtype, device=tensor_chunk.device)) + + tensor_partition_flat_buffer = instrument_w_nvtx(torch.cat)(tensor_partitions_lst_with_padding) + + tensor_partition_flat_buffer.div_(world_sz) # pre-divide + tensor_partition_buffer_for_each_rank: List[Tensor] = torch.chunk(tensor_partition_flat_buffer, world_sz) + + # batched reduce-scatter call + _torch_reduce_scatter_fn(tensor_partition_flat_buffer, + tensor_partition_buffer_for_each_rank[this_rank], + group=group) + + # reverse procedure of the interleaving done previously, done on the + # result of the batched reduce-scatter + output_lst: List[Tensor] = [None] * len(tensors) + offset = 0 + for tensor_idx in range(len(tensors)): + output_lst[tensor_idx] = tensor_partition_buffer_for_each_rank[this_rank].narrow( + 0, offset, partition_lst_for_each_tensor[tensor_idx][this_rank].numel()) + + offset += padded_partition_sz_for_each_tensor[tensor_idx] + return output_lst diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/compressed.py b/lib/python3.12/site-packages/deepspeed/runtime/comm/compressed.py new file mode 100644 index 0000000000000000000000000000000000000000..2c5482eb1ad769bb3a0ecd87b4a3489d31ce28db --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/comm/compressed.py @@ -0,0 +1,137 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import numpy as np +import torch +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import PackbitsBuilder + + +class CompressedBackend(object): + + def __init__(self, mpu=None): + if mpu is None: + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + else: + self.mpu = mpu + self.world_group = self.mpu.get_data_parallel_group() + self.size = dist.get_world_size(group=self.world_group) + self.rank = dist.get_rank(group=self.world_group) + self.packer = PackbitsBuilder().load() + + def my_igather(self, rank, size, group, sendbuf, recvbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) + else: + recvbuf[rank] = sendbuf + else: + req.append(dist.isend(sendbuf, group=group, dst=root)) + return req + + def my_gather(self, rank, size, group, sendbuf, recvbuf, root): + if rank == root: + for idx in range(size): + if idx != rank: + dist.recv(recvbuf[idx], src=idx, group=group) + else: + recvbuf[rank] = sendbuf + else: + dist.send(sendbuf, group=group, dst=root) + + def pack(self, buffer, size): + # pack float tensor into uint8 tensor + packed = self.packer.packbits(buffer.float(), buffer.numel(), self.rank) + return packed.reshape(size, -1) + + def unpack(self, buffer, size, dtype): + # unpack uint8 to float tensor + unpacked = self.packer.unpackbits(buffer, buffer.numel(), self.rank) + return unpacked.reshape(size, -1).to(dtype) + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + + # align size of original_buffer and error + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) + + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + sign_list_packed_tmp = self.pack(buffer_m, self.size).type(torch.int8) + + recvbuf_sign = torch.zeros([self.size, len(sign_list_packed_tmp[self.rank])], + dtype=sign_list_packed_tmp[0].dtype, + device=sign_list_packed_tmp.device) + + sign_list_packed = [sign_list_packed_tmp[idx] for idx in range(self.size)] + + recvbuf_scale = [ + torch.zeros(1, dtype=worker_scale.dtype, device=get_accelerator().current_device_name()) + for _ in range(self.size) + ] + + # communication phase 1 + # all to all for sign + dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) + # all gather for scale + dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) + + flattened_recvbuf_sign = recvbuf_sign.type(torch.uint8).flatten() + compensated_server_m = self.unpack(flattened_recvbuf_sign, self.size, torch.float32) \ + .mul_(torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) + + compensated_server_m.add_(server_error) + + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + server_sign_packed = self.pack(compensated_server_m, 1).type(torch.int8) + + # recvbuf_sign_server + recvbuf_sign_server_tmp = torch.zeros([self.size, len(server_sign_packed[0])], + dtype=recvbuf_sign.dtype, + device=server_sign_packed.device) + + recvbuf_sign_server = [recvbuf_sign_server_tmp[idx] for idx in range(self.size)] + + # recvbuf_scale_server + recvbuf_scale_server_tmp = torch.zeros([self.size, 1], + dtype=worker_scale.dtype, + device=server_sign_packed.device) + + recvbuf_scale_server = [recvbuf_scale_server_tmp[idx] for idx in range(self.size)] + + # communication Phase 2 + dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) + dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) + + recvbuf_sign_server = torch.stack(recvbuf_sign_server) + + flattened_recvbuf_sign_server = recvbuf_sign_server.type(torch.uint8).flatten() + + buffer_m.data.copy_( + self.unpack(flattened_recvbuf_sign_server, self.size, + torch.float32).mul_(recvbuf_scale_server_tmp).flatten().data) + + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + return buffer_m diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/hccl.py b/lib/python3.12/site-packages/deepspeed/runtime/comm/hccl.py new file mode 100644 index 0000000000000000000000000000000000000000..b8639c7da4c99327e4a6afe4e226d6c4a224886b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/comm/hccl.py @@ -0,0 +1,124 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import numpy as np +import torch +import torch_npu +import deepspeed.comm as dist + + +class HcclBackend(object): + + def __init__(self, mpu=None): + if mpu is None: + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + else: + self.mpu = mpu + self.world_group = self.mpu.get_data_parallel_group() + self.size = dist.get_world_size(group=self.world_group) + self.rank = dist.get_rank(group=self.world_group) + + def my_igather(self, rank, size, group, sendbuf, recvbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) + else: + recvbuf[rank] = sendbuf + else: + req.append(dist.isend(sendbuf, group=group, dst=root)) + return req + + def my_gather(self, rank, size, group, sendbuf, recvbuf, root): + if rank == root: + for idx in range(size): + if idx != rank: + dist.recv(recvbuf[idx], src=idx, group=group) + else: + recvbuf[rank] = sendbuf + else: + dist.send(sendbuf, group=group, dst=root) + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + + # align size of original_buffer and error + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) + + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + sign_list_packed_tmp = torch_npu.npu_sign_bits_pack(buffer_m, self.size).type(torch.int8) + + recvbuf_sign = torch.zeros([self.size, len(sign_list_packed_tmp[self.rank])], + dtype=sign_list_packed_tmp[0].dtype, + device=sign_list_packed_tmp.device) + + sign_list_packed = [sign_list_packed_tmp[idx] for idx in range(self.size)] + + recvbuf_scale = [ + torch.zeros(1, dtype=worker_scale.dtype, device=torch.device(local_rank)) for _ in range(self.size) + ] + + # communication phase 1 + # all to all for sign + dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) + # all gather for scale + dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) + + flattened_recvbuf_sign = recvbuf_sign.type(torch.uint8).flatten() + compensated_server_m = torch_npu.npu_sign_bits_unpack(flattened_recvbuf_sign, self.size, torch.float32) \ + .mul_(torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) + + compensated_server_m.add_(server_error) + + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + server_sign_packed = torch_npu.npu_sign_bits_pack(compensated_server_m, 1).type(torch.int8) + + # recvbuf_sign_server + recvbuf_sign_server_tmp = torch.zeros([self.size, len(server_sign_packed[0])], + dtype=recvbuf_sign.dtype, + device=server_sign_packed.device) + + recvbuf_sign_server = [recvbuf_sign_server_tmp[idx] for idx in range(self.size)] + + # recvbuf_scale_server + recvbuf_scale_server_tmp = torch.zeros([self.size, 1], + dtype=worker_scale.dtype, + device=server_sign_packed.device) + + recvbuf_scale_server = [recvbuf_scale_server_tmp[idx] for idx in range(self.size)] + + # communication Phase 2 + dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) + dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) + + recvbuf_sign_server = torch.stack(recvbuf_sign_server) + + flattened_recvbuf_sign_server = recvbuf_sign_server.type(torch.uint8).flatten() + + buffer_m.data.copy_( + torch_npu.npu_sign_bits_unpack(flattened_recvbuf_sign_server, self.size, + torch.float32).mul_(recvbuf_scale_server_tmp).flatten().data) + + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + return buffer_m diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/mpi.py b/lib/python3.12/site-packages/deepspeed/runtime/comm/mpi.py new file mode 100644 index 0000000000000000000000000000000000000000..bc544787aa7a7de5181d83897e725991f572e5ce --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/comm/mpi.py @@ -0,0 +1,215 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import cupy +import time +import numpy as np +from mpi4py import MPI + +from deepspeed.runtime.compression.cupy import CupyBackend + + +class MpiBackend(object): + + def __init__(self, cuda_aware): + self.comm = MPI.COMM_WORLD + self.rank = self.comm.Get_rank() + self.size = self.comm.Get_size() + self.cuda_aware = cuda_aware + self.compression_backend = CupyBackend() + + def my_igather(self, rank, size, comm, sendbuf, recbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(comm.Irecv(recbuf[idx], source=idx)) + else: + recbuf[rank] = sendbuf + else: + req.append(comm.Isend(sendbuf, dest=root)) + return req + + def gather_cuda(self, rank, world_size, comm, cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, + cupy_recvbuf_scale): + # We do in-place operations on cupy buffers so we do not return any buffers + requests = [] + for idx in range(world_size): + req_sign = self.my_igather(rank, world_size, comm, cupy_sign_list_packed[idx], cupy_recvbuf_sign, root=idx) + requests += req_sign + + for idx in range(world_size): + req_scale = self.my_igather(rank, world_size, comm, cupy_worker_scale, cupy_recvbuf_scale, root=idx) + requests += req_scale + + MPI.Request.Waitall(requests) + + def gather_host(self, rank, world_size, comm, cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, + cupy_recvbuf_scale): + + # In-place operations are not possible for newly created cupy arrays + # so we need to return the new buffers + numpy_recvbuf_sign = np.zeros([world_size, cupy_sign_list_packed[rank].size], + dtype=cupy_sign_list_packed[0].dtype) + numpy_recvbuf_scale = np.zeros([world_size, 1], dtype=cupy_worker_scale.dtype) + + # 1. convert from cupy to numpy + numpy_sign_list_packed = cupy_sign_list_packed + + for idx in range(world_size): + numpy_sign_list_packed[idx] = cupy.asnumpy(cupy_sign_list_packed[idx]) + + numpy_worker_scale = cupy.asnumpy(cupy_worker_scale) + numpy_recvbuf_scale = cupy.asnumpy(cupy_recvbuf_scale) + + cupy.cuda.get_current_stream().synchronize() + + # 2. use numpy buffers for communication + requests = [] + + for idx in range(world_size): + req_sign = self.my_igather(rank, + world_size, + comm, + numpy_sign_list_packed[idx], + numpy_recvbuf_sign, + root=idx) + requests += req_sign + + for idx in range(world_size): + req_scale = self.my_igather(rank, world_size, comm, numpy_worker_scale, numpy_recvbuf_scale, root=idx) + requests += req_scale + + MPI.Request.Waitall(requests) + + # 3. Convert back from numpy to cupy + cupy_recvbuf_sign = cupy.asarray(numpy_recvbuf_sign) + for idx in range(world_size): + cupy_sign_list_packed[idx] = cupy.asarray(numpy_sign_list_packed[idx]) + + cupy_worker_scale = cupy.asarray(numpy_worker_scale) + cupy_recvbuf_scale = cupy.asarray(numpy_recvbuf_scale) + cupy.cuda.get_current_stream().synchronize() + + return cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, cupy_recvbuf_scale + + def allgather_cuda(self, comm, cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server): + comm.Allgather(cupy_server_sign_packed, cupy_recvbuf_sign_server) + comm.Allgather(cupy_server_scale, cupy_recvbuf_scale_server) + + def allgather_host(self, comm, cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server): + + # 1. Convert cupy to numpy + numpy_recvbuf_sign_server = np.zeros([comm.Get_size(), cupy_server_sign_packed.size], + dtype=cupy_server_sign_packed.dtype) + numpy_recvbuf_scale_server = np.zeros([comm.Get_size(), 1], dtype=cupy_server_scale.dtype) + + numpy_server_sign_packed = cupy.asnumpy(cupy_server_sign_packed) + numpy_recvbuf_sign_server = cupy.asnumpy(cupy_recvbuf_sign_server) + numpy_server_scale = cupy.asnumpy(cupy_server_scale) + numpy_recvbuf_scale_server = cupy.asnumpy(cupy_recvbuf_scale_server) + cupy.cuda.get_current_stream().synchronize() + + # 2. Communicate numpy buffers + comm.Allgather(numpy_server_sign_packed, numpy_recvbuf_sign_server) + comm.Allgather(numpy_server_scale, numpy_recvbuf_scale_server) + comm.Barrier() + + # 3. Convert numpy back to cupy + cupy_server_sign_packed = cupy.asarray(numpy_server_sign_packed) + cupy_recvbuf_sign_server = cupy.asarray(numpy_recvbuf_sign_server) + cupy_server_scale = cupy.asarray(numpy_server_scale) + cupy_recvbuf_scale_server = cupy.asarray(numpy_recvbuf_scale_server) + cupy.cuda.get_current_stream().synchronize() + + return cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, cupy_recvbuf_scale_server + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + + all_start_time = time.time() + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + cupy.cuda.Device(local_rank).use() + + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + cupy_sign_list_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool()), self.size) + cupy_worker_scale = self.compression_backend.torch2cupy(worker_scale) + + cupy_recvbuf_sign = cupy.zeros([self.size, cupy_sign_list_packed[self.rank].size], + dtype=cupy_sign_list_packed[0].dtype) + cupy_recvbuf_scale = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) + + # Communication Phase 1 + gather_start = time.time() + if self.cuda_aware: + self.gather_cuda(self.rank, self.size, self.comm, cupy_sign_list_packed, cupy_recvbuf_sign, + cupy_worker_scale, cupy_recvbuf_scale) + else: + _, cupy_recvbuf_sign, _, cupy_recvbuf_scale = self.gather_host(self.rank, self.size, self.comm, + cupy_sign_list_packed, cupy_recvbuf_sign, + cupy_worker_scale, cupy_recvbuf_scale) + gather_end = time.time() + + # cupy_sign_list_packed, cupy_worker_scale, worker_scale = None, None, None + cupy_sign_list_packed = None + + compensated_server_m = self.compression_backend.cupy2torch( + (cupy.unpackbits(cupy_recvbuf_sign.flatten())).reshape(self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + self.compression_backend.cupy2torch(cupy_recvbuf_scale).mul_(1 / self.size)).sum(0) + compensated_server_m.add_(server_error) + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + cupy_server_scale = self.compression_backend.torch2cupy(server_scale) + + cupy_server_sign_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool()), 1) + compensated_server_m = None + + cupy_recvbuf_sign_server = cupy.zeros([self.size, cupy_server_sign_packed[0].size], + dtype=cupy_recvbuf_sign.dtype) + cupy_recvbuf_scale_server = cupy.zeros([self.size, 1], dtype=cupy_recvbuf_scale.dtype) + # cupy_recvbuf_sign, cupy_recvbuf_scale = None, None + cupy_recvbuf_sign = None + + # Communication Phase 2 + if self.cuda_aware: + self.allgather_cuda(self.comm, cupy_server_sign_packed[0], cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server) + else: + _, cupy_recvbuf_sign_server, _, cupy_recvbuf_scale_server = self.allgather_host( + self.comm, cupy_server_sign_packed[0], cupy_recvbuf_sign_server, cupy_server_scale, + cupy_recvbuf_scale_server) + + # cupy_server_sign_packed, cupy_server_scale, server_scale = None, None, None + cupy_server_sign_packed = None + + buffer_m.data.copy_( + self.compression_backend.cupy2torch((cupy.unpackbits(cupy_recvbuf_sign_server.flatten())).reshape( + self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + self.compression_backend.cupy2torch(cupy_recvbuf_scale_server)).flatten().data) + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + # cupy_recvbuf_sign_server, cupy_recvbuf_scale_server = None, None + + return buffer_m diff --git a/lib/python3.12/site-packages/deepspeed/runtime/comm/nccl.py b/lib/python3.12/site-packages/deepspeed/runtime/comm/nccl.py new file mode 100644 index 0000000000000000000000000000000000000000..a57b7519a295e584e77561843c52fa025bfaa66d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/comm/nccl.py @@ -0,0 +1,166 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed import comm as dist +import cupy +import numpy as np + +from deepspeed.runtime.compression.cupy import CupyBackend +from deepspeed.utils.torch import required_torch_version +from deepspeed.accelerator import get_accelerator + + +class NcclBackend(object): + + def __init__(self, mpu=None): + if mpu is None: + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + else: + self.mpu = mpu + self.world_group = self.mpu.get_data_parallel_group() + self.rank = dist.get_rank(group=self.world_group) + self.size = dist.get_world_size(group=self.world_group) + self.compression_backend = CupyBackend() + self.bool_not_supported = required_torch_version(min_version=1.10) + + def my_igather(self, rank, size, group, sendbuf, recvbuf, root): + req = [] + if rank == root: + for idx in range(size): + if idx != rank: + req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) + else: + recvbuf[rank] = sendbuf + else: + req.append(dist.isend(sendbuf, group=group, dst=root)) + return req + + def my_gather(self, rank, size, group, sendbuf, recvbuf, root): + if rank == root: + for idx in range(size): + if idx != rank: + dist.recv(recvbuf[idx], src=idx, group=group) + else: + recvbuf[rank] = sendbuf + else: + dist.send(sendbuf, group=group, dst=root) + + def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): + + # all_start_time = time.time() + original_shape = buffer_m.size() + if len(original_shape) > 1: + buffer_m = torch.flatten(buffer_m) + original_size = buffer_m.numel() + worker_error_size = worker_error.numel() + cupy.cuda.Device(local_rank).use() + + if original_size != worker_error_size: + empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) + buffer_m = torch.cat([buffer_m, empty_tensor]) + + buffer_m.add_(worker_error) + worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(buffer_m.numel()) + worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + if self.bool_not_supported: + cupy_sign_list_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool().to(dtype=torch.uint8)), self.size) + else: + cupy_sign_list_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool()), self.size) + cupy_worker_scale = self.compression_backend.torch2cupy(worker_scale) + + cupy_recvbuf_sign = cupy.zeros([self.size, cupy_sign_list_packed[self.rank].size], + dtype=cupy_sign_list_packed[0].dtype) + # cupy_recvbuf_scale = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) + + sign_list_packed = [ + self.compression_backend.cupy2torch(cupy_sign_list_packed[idx]) for idx in range(self.size) + ] + + # worker_scale = self.compression_backend.cupy2torch(cupy_worker_scale) + recvbuf_sign = self.compression_backend.cupy2torch(cupy_recvbuf_sign) + #recvbuf_scale = self.compression_backend.cupy2torch(cupy_recvbuf_scale) + recvbuf_scale = [ + torch.zeros(1, dtype=worker_scale.dtype, device=torch.device(get_accelerator().device_name(local_rank))) + for i in range(self.size) + ] + + # communication phase 1 + # gather_start = time.time() + # Alltoall for sign + dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) + # Allgather for scale + dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) + + # gather_end = time.time() + + # cupy_sign_list_packed, sign_list_packed, cupy_worker_scale, worker_scale = None, None, None, None + cupy_sign_list_packed = None + + cupy_recvbuf_sign = self.compression_backend.torch2cupy(recvbuf_sign) + #cupy_recvbuf_scale = self.compression_backend.torch2cupy(torch.stack(recvbuf_scale)) + + compensated_server_m = self.compression_backend.cupy2torch( + (cupy.unpackbits(cupy_recvbuf_sign.flatten())).reshape(self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) + compensated_server_m.add_(server_error) + server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) + server_error.set_(compensated_server_m - + server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) + + # cupy_server_scale = self.compression_backend.torch2cupy(server_scale) + + if self.bool_not_supported: + cupy_server_sign_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool().to(dtype=torch.uint8)), + 1) + else: + cupy_server_sign_packed = self.compression_backend.compress_by_chunk( + self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool()), 1) + compensated_server_m = None + + cupy_recvbuf_sign_server = cupy.zeros([self.size, cupy_server_sign_packed[0].size], + dtype=cupy_recvbuf_sign.dtype) + # cupy_recvbuf_sign, recvbuf_sign = None, None + cupy_recvbuf_sign = None + + server_sign_packed = [self.compression_backend.cupy2torch(cupy_server_sign_packed[0])] + recvbuf_sign_server = [ + self.compression_backend.cupy2torch(cupy_recvbuf_sign_server[idx]) for idx in range(self.size) + ] + + # server_scale = self.compression_backend.cupy2torch(cupy_server_scale) + cupy_recvbuf_scale_server = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) + # cupy_recvbuf_scale, recvbuf_scale = None, None + + recvbuf_scale_server = [ + self.compression_backend.cupy2torch(cupy_recvbuf_scale_server[idx]) for idx in range(self.size) + ] + + # Communication Phase 2 + dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) + dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) + + cupy_server_sign_packed = None + + # need to convert from a tensor list to a single tensor + # dist.all_gather only provides a tensor list as the recv/output buffer + recvbuf_sign_server = torch.stack(recvbuf_sign_server) + + cupy_recvbuf_sign_server = self.compression_backend.torch2cupy(recvbuf_sign_server) + + buffer_m.data.copy_( + self.compression_backend.cupy2torch((cupy.unpackbits(cupy_recvbuf_sign_server.flatten())).reshape( + self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( + self.compression_backend.cupy2torch(cupy_recvbuf_scale_server)).flatten().data) + if original_size != worker_error_size: + buffer_m = buffer_m[0:original_size] + if len(original_shape) > 1: + buffer_m = buffer_m.reshape(original_shape) + + return buffer_m diff --git a/lib/python3.12/site-packages/deepspeed/runtime/domino/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/domino/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..208299fb8c50f73468d293b6fa5dca71649d62e7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/domino/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00136d2fa6f3122c29fa843899b1ded87899caaa Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4285f9bd6103c30b3d4a2f02cb7d2c221d78c7d9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f542ef6162baff70637b7cef6c2c4fec29209de3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/domino/async_linear.py b/lib/python3.12/site-packages/deepspeed/runtime/domino/async_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..8e01da500409cd603c17af646e902cfe2a82525b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/domino/async_linear.py @@ -0,0 +1,140 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Adapted from https://github.com/NVIDIA/Megatron-LM/blob/23.08/megatron/core/tensor_parallel/layers.py + +import torch +from torch.nn.parameter import Parameter +import torch.nn.functional as F +from deepspeed.accelerator import get_accelerator +import deepspeed.comm as dist +from typing import Callable + +TP_group = None + + +class DominoAsyncColumnParallelLinearImpl(torch.autograd.Function): + + @staticmethod + def forward(ctx, inp, weight, bias, handle_dic, h_id): # inp: (b, s, k), weight: (m, k), bias (m) + ctx.save_for_backward(inp, weight, bias) + ctx.handle_dic = handle_dic + ctx.h_id = h_id + output = torch.matmul(inp, weight.t()) # (b, s, k) @ (k, m) -> (b, s, m) + if bias is not None: # bias (m) + output = output + bias + return output + + @staticmethod + def backward(ctx, grad_output): + inp, weight, bias = ctx.saved_tensors + grad_input = grad_weight = grad_bias = None + grad_input = torch.matmul(grad_output, weight) # (b, s, m) @ (m, k) -> (b, s, k) + handle = dist.all_reduce(grad_input, group=TP_group, async_op=True) + ctx.handle_dic[ctx.h_id] = handle + grad_output = grad_output.view(grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2]) # (b*s, m) + + inp = inp.view(inp.shape[0] * inp.shape[1], inp.shape[2]) # (b*s, k) + grad_weight = torch.matmul(grad_output.t(), inp) # (m, b*s) @ (b*s, k) -> (m, k) + + if bias is not None: + grad_bias = grad_output.sum(dim=0) # (b*s, m) -> (m) + return grad_input, grad_weight, grad_bias, None, None + + +class DominoAsyncColumnParallelLinear(torch.nn.Module): + + def __init__(self, + input_size, + output_size, + _tp_group, + config, + init_method: Callable, + bias=True, + skip_bias_add=False): + super(DominoAsyncColumnParallelLinear, self).__init__() + + self.skip_bias_add = skip_bias_add + + global TP_group + if TP_group == None: + TP_group = _tp_group + + self.weight = Parameter( + torch.empty( + output_size, + input_size, + device=get_accelerator().current_device_name(), + dtype=config.params_dtype, + )) + if config.perform_initialization: + init_method(self.weight) + + if bias: + self.bias = Parameter( + torch.empty(output_size, device=get_accelerator().current_device_name(), dtype=config.params_dtype)) + + if config.perform_initialization: + with torch.no_grad(): + self.bias.zero_() + else: + self.register_parameter('bias', None) + + def forward(self, input_: torch.Tensor, handle_dic, h_id): + + bias = self.bias if not self.skip_bias_add else None + + output = DominoAsyncColumnParallelLinearImpl.apply(input_, self.weight, bias, handle_dic, h_id) + + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + +class RowParallelLinearNoComm(torch.nn.Module): + + def __init__( + self, + input_size: int, + output_size: int, + config, + init_method: Callable, + bias: bool = True, + stride: int = 1, + skip_bias_add: bool = False, + ): + super(RowParallelLinearNoComm, self).__init__() + + self.skip_bias_add = skip_bias_add + + self.weight = Parameter( + torch.empty( + output_size, + input_size, + device=get_accelerator().current_device_name(), + dtype=config.params_dtype, + )) + if config.perform_initialization: + init_method(self.weight) + if bias: + self.bias = Parameter( + torch.empty( + output_size, + device=get_accelerator().current_device_name(), + dtype=config.params_dtype, + )) + + if config.perform_initialization: + with torch.no_grad(): + self.bias.zero_() + else: + self.register_parameter('bias', None) + + def forward(self, input_): + bias = self.bias if not self.skip_bias_add else None + + output = F.linear(input_, self.weight, bias) + + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias diff --git a/lib/python3.12/site-packages/deepspeed/runtime/domino/transformer.py b/lib/python3.12/site-packages/deepspeed/runtime/domino/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..3dfb133373b5aa7686f1dd9aac903add10b4e84d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/domino/transformer.py @@ -0,0 +1,605 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import torch.nn.functional as F +import enum +import deepspeed.comm as dist + +from .async_linear import DominoAsyncColumnParallelLinear, RowParallelLinearNoComm + + +class LayerType(enum.Enum): + encoder = 1 + decoder = 2 + + +class AttnType(enum.Enum): + self_attn = 1 + cross_attn = 2 + + +class AttnMaskType(enum.Enum): + padding = 1 + causal = 2 + + +class ModelType(enum.Enum): + encoder_or_decoder = 1 + encoder_and_decoder = 2 + + +class DominoUtil: + + BATCH_0 = "BATCH0" + + BATCH_1 = "BATCH1" + + HANDLE_DIC = {"BATCH0": None, "BATCH1": None} + + +class DominoModule(torch.nn.Module): + """extensions of torch Module.""" + + def __init__(self, ): + super(DominoModule, self).__init__() + + +def _Wait_bwd_comm(input_, dic_, h_id): + return NoOper.apply(input_, dic_, h_id) + + +class NoOper(torch.autograd.Function): + + @staticmethod + def symbolic(graph, input_, handle_dic, h_id): + return input_ + + @staticmethod + def forward(ctx, input_, handle_dic, h_id): + ctx.handle_dic = handle_dic + ctx.h_id = h_id + return input_ + + @staticmethod + def backward(ctx, grad_output): + handle = ctx.handle_dic[ctx.h_id] + handle.wait() + return grad_output, None, None + + +class CoreAttention(DominoModule): + + def __init__(self, config, tp_world_size, attn_mask_type=AttnMaskType.causal): + super(CoreAttention, self).__init__() + + self.attn_mask_type = attn_mask_type + + projection_size = config.kv_channels * config.num_attention_heads + + # Per attention head and per partition values. + assert projection_size % tp_world_size == 0, f"projection size {projection_size} should be multiple of TP world size {tp_world_size}" + self.hidden_size_per_partition = projection_size // tp_world_size + self.attention_dropout_rate = config.attention_dropout + + def forward(self, query_layer, key_layer, value_layer, attention_mask): + + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, + key_layer, + value_layer, + attn_mask=None, + dropout_p=self.attention_dropout_rate, + is_causal=True, + scale=None) + + # [b, np, sq, hn] --> [sq, b, np, hn] + context_layer = context_layer.permute(2, 0, 1, 3).contiguous() + + # [sq, b, np, hn] --> [sq, b, hp] + new_context_layer_shape = context_layer.size()[:-2] + \ + (self.hidden_size_per_partition,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer + + +class ShardedAttention(DominoModule): + """Sharded self-attention layer class. + Only support self attention and causal attention mask for now. + """ + + def __init__(self, + config, + mpu, + apply_rotary_pos_emb, + layer_number, + attention_type=AttnType.self_attn, + attn_mask_type=AttnMaskType.causal): + super(ShardedAttention, self).__init__() + + assert attention_type == AttnType.self_attn, "Only support self_attn for now!" + + self.layer_number = max(1, layer_number) + self.attention_type = attention_type + self.attn_mask_type = attn_mask_type + self.params_dtype = config.params_dtype + self.apply_rotary_pos_emb = apply_rotary_pos_emb + + query_projection_size = config.kv_channels * config.num_attention_heads + kv_projection_size = config.kv_channels * config.num_attention_heads + + tp_world_size = mpu.get_tensor_model_parallel_world_size() + self.hidden_size_per_attention_head = query_projection_size // config.num_attention_heads + self.num_attention_heads_per_partition = config.num_attention_heads // tp_world_size + + qkv_projection_per_partition = (query_projection_size + 2 * kv_projection_size) // tp_world_size + + self.query_key_value = DominoAsyncColumnParallelLinear(config.hidden_size, + qkv_projection_per_partition, + mpu.get_tensor_model_parallel_group(), + config=config, + init_method=config.init_method, + bias=config.add_bias_linear) + + self.core_attention = CoreAttention(config, tp_world_size, self.attn_mask_type) + + query_projection_size_per_partition = query_projection_size // tp_world_size + + # Output. + self.dense = RowParallelLinearNoComm(query_projection_size_per_partition, + config.hidden_size, + config=config, + init_method=config.output_layer_init_method, + bias=config.add_bias_linear, + skip_bias_add=True) + + def forward(self, hidden_states, attention_mask, micro_batch_num, rotary_pos_emb=None): + # hidden_states: [sq, b, h] + + mixed_x_layer, _ = self.query_key_value(hidden_states, DominoUtil.HANDLE_DIC, micro_batch_num) + + new_tensor_shape = mixed_x_layer.size()[:-1] + ( + self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head, + ) + + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + mixed_x_layer = mixed_x_layer.permute(1, 2, 0, 3).contiguous() + + (query_layer, key_layer, value_layer) = torch.split(mixed_x_layer, [ + self.hidden_size_per_attention_head, self.hidden_size_per_attention_head, + self.hidden_size_per_attention_head + ], + dim=3) + + query_layer = query_layer.view(query_layer.size(0), query_layer.size(1), -1, + self.hidden_size_per_attention_head) + + if rotary_pos_emb is not None: + if isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = rotary_pos_emb + else: + rotary_pos_emb = ((rotary_pos_emb, ) * 2) + q_pos_emb, k_pos_emb = rotary_pos_emb + query_layer = self.apply_rotary_pos_emb(query_layer, q_pos_emb) + key_layer = self.apply_rotary_pos_emb(key_layer, k_pos_emb) + + context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) + + output, bias = self.dense(context_layer) + return output, bias + + def domino_core_attention_forward(self, mixed_x_layer, attention_mask, rotary_pos_emb=None): + # hidden_states: [sq, b, h] + + # To illustrate the difference between intra-layer overlap and inter-layer overlap + # mixed_x_layer, _ = self.query_key_value(hidden_states, handle_dic, micro_batch_num) + + new_tensor_shape = mixed_x_layer.size()[:-1] + ( + self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head, + ) + + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + mixed_x_layer = mixed_x_layer.permute(1, 2, 0, 3).contiguous() + + (query_layer, key_layer, value_layer) = torch.split(mixed_x_layer, [ + self.hidden_size_per_attention_head, self.hidden_size_per_attention_head, + self.hidden_size_per_attention_head + ], + dim=3) + + query_layer = query_layer.view(query_layer.size(0), query_layer.size(1), -1, + self.hidden_size_per_attention_head) + + if rotary_pos_emb is not None: + if isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = rotary_pos_emb + else: + rotary_pos_emb = ((rotary_pos_emb, ) * 2) + q_pos_emb, k_pos_emb = rotary_pos_emb + query_layer = self.apply_rotary_pos_emb(query_layer, q_pos_emb) + key_layer = self.apply_rotary_pos_emb(key_layer, k_pos_emb) + + context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) + + # output, bias = self.dense(context_layer) + # return output, bias + + return context_layer + + +class bias_dropout_add(torch.nn.Module): + + def __init__(self, prob: float): + super(bias_dropout_add, self).__init__() + self.dropout = torch.nn.Dropout(prob) + + def forward(self, x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + if bias is not None: + x = x + bias + out = self.dropout(x) + out = out + residual + return out + + +class DominoTransformerLayer(DominoModule): + """A domino single transformer layer. + [s, b, h] -> [s, b, h] + """ + + def __init__(self, + config, + mpu, + apply_rotary_pos_emb, + layer_number, + layer_type=LayerType.encoder, + self_attn_mask_type=AttnMaskType.causal, + drop_path_rate=0.): + + super(DominoTransformerLayer, self).__init__() + self.layer_number = layer_number + self.layer_type = layer_type + + self.apply_residual_connection_post_layernorm \ + = config.apply_residual_connection_post_layernorm + + self.llama_model = False + + self.input_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) + + # Self attention. + self.self_attention = ShardedAttention(config, + mpu, + apply_rotary_pos_emb, + layer_number, + attention_type=AttnType.self_attn, + attn_mask_type=self_attn_mask_type) + + self.hidden_dropout = config.hidden_dropout + + self.post_attention_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) + + # MLP + ffn_hidden_size = config.ffn_hidden_size + if config.gated_linear_unit: + ffn_hidden_size *= 2 + + self.output_size_c = config.ffn_hidden_size + self.input_size_c = config.hidden_size + self.input_size_r = config.ffn_hidden_size + self.output_size_r = self.input_size_c + + tp_world_size = mpu.get_tensor_model_parallel_world_size() + self.TP_group = mpu.get_tensor_model_parallel_group() + self.output_size_per_partition = self.output_size_c // tp_world_size + self.input_size_per_partition = self.input_size_r // tp_world_size + + self.linear_fc1 = DominoAsyncColumnParallelLinear(self.input_size_c, + self.output_size_per_partition, + mpu.get_tensor_model_parallel_group(), + config=config, + init_method=config.init_method, + bias=config.add_bias_linear) + + self.mlp_activation_func = F.gelu + + self.linear_fc2 = RowParallelLinearNoComm(self.input_size_per_partition, + self.output_size_r, + config=config, + init_method=config.output_layer_init_method, + bias=config.add_bias_linear, + skip_bias_add=True) + + self.bias_dropout_add_func = bias_dropout_add(self.hidden_dropout) + + def forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + + hidden_states0, hidden_states1 = hidden_states + + layernorm_output0 = self.input_layernorm(hidden_states0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + # Micro batch 0: attention + attention_output0, attention_bias0 = self.self_attention(layernorm_output0, + attention_mask, + DominoUtil.BATCH_0, + rotary_pos_emb=rotary_pos_emb) + + fwd_handle0 = dist.all_reduce(attention_output0, group=self.TP_group, async_op=True) + # End of Micro batch 0: attention + + # Micro batch 1: attention + layernorm_output1 = self.input_layernorm(hidden_states1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + + attention_output1, attention_bias1 = self.self_attention(layernorm_output1, + attention_mask, + DominoUtil.BATCH_1, + rotary_pos_emb=rotary_pos_emb) + fwd_handle1 = dist.all_reduce(attention_output1, group=self.TP_group, async_op=True) + + # Micro batch 0: Residual connection. + fwd_handle0.wait() + if self.apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = hidden_states0 + + layernorm_input0 = self.bias_dropout_add_func(attention_output0, attention_bias0, residual0) + + layernorm_output0 = self.post_attention_layernorm(layernorm_input0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + if self.apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = layernorm_input0 + # End of Micro batch 0: Residual connection. + + # ------------ MLP ------------ + # Micro batch 0: MLP + output0, _ = self.linear_fc1(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + output0 = self.mlp_activation_func(output0) + + # Micro batch 1: Residual connection. + fwd_handle1.wait() + if self.apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = hidden_states1 + + layernorm_input1 = self.bias_dropout_add_func(attention_output1, attention_bias1, residual1) + + layernorm_output1 = self.post_attention_layernorm(layernorm_input1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + + if self.apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = layernorm_input1 + # End of Micro batch 1: Residual connection. + + hidden_states0, last_mlp_bias = self.linear_fc2(output0) + fwd_handle0 = dist.all_reduce(hidden_states0, group=self.TP_group, async_op=True) + # End of Micro batch 0: MLP + + # Micro batch 1: MLP + output1, _ = self.linear_fc1(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + output1 = self.mlp_activation_func(output1) + + hidden_states1, last_mlp_bias = self.linear_fc2(output1) + + fwd_handle1 = dist.all_reduce(hidden_states1, group=self.TP_group, async_op=True) + # End of Micro batch 1: MLP + + # ------------ End of MLP ------------ + + fwd_handle0.wait() + hidden_states0 = self.bias_dropout_add_func(hidden_states0, last_mlp_bias, residual0) + + fwd_handle1.wait() + hidden_states1 = self.bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) + + return hidden_states0, hidden_states1 + + +class DominoTransformer(DominoModule): + """Transformer class.""" + + def __init__(self, + config, + mpu, + apply_rotary_pos_emb, + model_type, + layer_type=LayerType.encoder, + self_attn_mask_type=AttnMaskType.causal, + post_layer_norm=True, + pre_process=True, + post_process=True, + drop_path_rate=0.0): + super(DominoTransformer, self).__init__() + + self.layer_type = layer_type + self.model_type = model_type + self.post_layer_norm = post_layer_norm + self.post_process = post_process + self.input_tensor = None + self.drop_path_rate = drop_path_rate + self.TP_group = mpu.get_tensor_model_parallel_group() + + if not dist.is_initialized(): + dist.init_distributed() + assert dist.is_initialized(), "deepspeed.comm failed to initialize!" + + self.num_layers = config.num_layers + + self.drop_path_rates = [rate.item() for rate in torch.linspace(0, self.drop_path_rate, config.num_layers)] + + def build_layer(layer_number): + + current_layer_type = layer_type + return DominoTransformerLayer(config, + mpu, + apply_rotary_pos_emb, + layer_number, + layer_type=current_layer_type, + self_attn_mask_type=self_attn_mask_type, + drop_path_rate=self.drop_path_rates[layer_number - 1]) + + self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) + + if self.post_process and self.post_layer_norm: + self.final_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) + + self._forward_impl = self.inter_layer_overlap_forward + if config.domino_intra_layer_overlap: + self._forward_impl = self.intra_layer_overlap_forward + + def forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + + return self._forward_impl(hidden_states, attention_mask, rotary_pos_emb) + + def inter_layer_overlap_forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + # hidden_states: [s, b, h] + + hidden_states0, hidden_states1 = torch.chunk(hidden_states, chunks=2, dim=1) + + last_mlp_bias = None + fwd_handle0, fwd_handle1 = None, None + residual0, residual1 = None, None + + layernorm_output0 = self.layers[0].input_layernorm(hidden_states0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + for index in range(self.num_layers): + + # Micro batch 0: attention + attention_output0, _ = self.layers[index].self_attention.query_key_value( + layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + attention_output0 = self.layers[index].self_attention.domino_core_attention_forward( + attention_output0, attention_mask, rotary_pos_emb=rotary_pos_emb) + + # Micro batch 1: Residual connection + if index > 0: + fwd_handle1.wait() + hidden_states1 = self.layers[index - 1].bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) + + layernorm_output1 = self.layers[index].input_layernorm(hidden_states1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + # End of Micro batch 1: Residual connection + + attention_output0, attention_bias0 = self.layers[index].self_attention.dense(attention_output0) + + fwd_handle0 = dist.all_reduce(attention_output0, group=self.TP_group, async_op=True) + # End of Micro batch 0: attention + + # Micro batch 1: attention + attention_output1, _ = self.layers[index].self_attention.query_key_value( + layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + attention_output1 = self.layers[index].self_attention.domino_core_attention_forward( + attention_output1, attention_mask, rotary_pos_emb=rotary_pos_emb) + + # Micro batch 0: Residual connection. + fwd_handle0.wait() + if self.layers[index].apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = hidden_states0 + + layernorm_input0 = self.layers[index].bias_dropout_add_func(attention_output0, attention_bias0, residual0) + + layernorm_output0 = self.layers[index].post_attention_layernorm(layernorm_input0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + + if self.layers[index].apply_residual_connection_post_layernorm: + residual0 = layernorm_output0 + else: + residual0 = layernorm_input0 + # End of Micro batch 0: Residual connection. + + attention_output1, attention_bias1 = self.layers[index].self_attention.dense(attention_output1) + fwd_handle1 = dist.all_reduce(attention_output1, group=self.TP_group, async_op=True) + # End of Micro batch 1: attention + + # ------------ MLP ------------ + # Micro batch 0: MLP + output0, _ = self.layers[index].linear_fc1(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + output0 = self.layers[index].mlp_activation_func(output0) + + # Micro batch 1: Residual connection. + fwd_handle1.wait() + if self.layers[index].apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = hidden_states1 + + layernorm_input1 = self.layers[index].bias_dropout_add_func(attention_output1, attention_bias1, residual1) + + layernorm_output1 = self.layers[index].post_attention_layernorm(layernorm_input1) + layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + + if self.layers[index].apply_residual_connection_post_layernorm: + residual1 = layernorm_output1 + else: + residual1 = layernorm_input1 + # End of Micro batch 1: Residual connection. + + hidden_states0, last_mlp_bias = self.layers[index].linear_fc2(output0) + fwd_handle0 = dist.all_reduce(hidden_states0, group=self.TP_group, async_op=True) + # End of Micro batch 0: MLP + + # Micro batch 1: MLP + output1, _ = self.layers[index].linear_fc1(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) + output1 = self.layers[index].mlp_activation_func(output1) + + # Micro batch 0: Residual connection. + fwd_handle0.wait() + hidden_states0 = self.layers[index].bias_dropout_add_func(hidden_states0, last_mlp_bias, residual0) + + if index < self.num_layers - 1: + layernorm_output0 = self.layers[index + 1].input_layernorm(hidden_states0) + layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) + # End of Micro batch 0: Residual connection. + + hidden_states1, last_mlp_bias = self.layers[index].linear_fc2(output1) + + fwd_handle1 = dist.all_reduce(hidden_states1, group=self.TP_group, async_op=True) + # End of Micro batch 1: MLP + + # ------------ End of MLP ------------ + + if self.post_process and self.post_layer_norm: + hidden_states0 = self.final_layernorm(hidden_states0) + + index = self.num_layers - 1 + + fwd_handle1.wait() + hidden_states1 = self.layers[index].bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) + + if self.post_process and self.post_layer_norm: + hidden_states1 = self.final_layernorm(hidden_states1) + + hidden_states = torch.cat([hidden_states0, hidden_states1], dim=1) + + return hidden_states + + def intra_layer_overlap_forward(self, hidden_states, attention_mask, rotary_pos_emb=None): + + hidden_states = torch.chunk(hidden_states, chunks=2, dim=1) + + for index in range(self.num_layers): + layer = self.layers[index] + hidden_states = layer(hidden_states, attention_mask, rotary_pos_emb) + + hidden_states0, hidden_states1 = hidden_states + if self.post_process and self.post_layer_norm: + hidden_states0 = self.final_layernorm(hidden_states0) + hidden_states1 = self.final_layernorm(hidden_states1) + + hidden_states = torch.cat([hidden_states0, hidden_states1], dim=1) + return hidden_states diff --git a/lib/python3.12/site-packages/deepspeed/runtime/engine.py b/lib/python3.12/site-packages/deepspeed/runtime/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..df033fe6d32ac523e72f52327e12c5e5f3ee587b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/engine.py @@ -0,0 +1,3999 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import re +import stat +import torch +import hashlib +from collections import defaultdict, OrderedDict, deque +from shutil import copyfile +import gc + +from torch.nn.modules import Module +from torch.nn.parameter import Parameter +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from contextlib import contextmanager + +from typing import Callable, Dict, Union, Iterable, Container + +import deepspeed + +from deepspeed import comm as dist +from deepspeed.runtime.utils import see_memory_usage, DummyOptim +from .zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum +from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer +from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus +from deepspeed.runtime.zero.utils import is_zero_supported_optimizer, ZeRORuntimeException +from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload +from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION + +from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer +from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer +from deepspeed.runtime.bf16_optimizer import BF16_Optimizer + +from deepspeed.linear.optimized_linear import LoRAOptimizedLinear +from deepspeed.module_inject.layers import GatherReplacedLayerParams, configure_tensor_parallel_runtime +from deepspeed.runtime.config import DEEPSPEED_OPTIMIZERS, \ + ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, \ + TORCH_ADAM_PARAM, ADAM_W_MODE, ADAM_W_MODE_DEFAULT, ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, \ + MUSGD_OPTIMIZER, LION_OPTIMIZER + +from deepspeed.runtime.dataloader import DeepSpeedDataLoader +from deepspeed.runtime.constants import \ + ROUTE_TRAIN, ROUTE_PREDICT, ROUTE_EVAL, \ + PLD_THETA, PLD_GAMMA, BFLOAT16, FP16, AMP, GRADIENT_ACCUMULATION_STEPS, \ + DATA_PARALLEL_GROUP, GLOBAL_RANK +from deepspeed.runtime.zero.config import ZeroStageEnum +from deepspeed.compression import compression_scheduler +from deepspeed.compression.constants import \ + WEIGHT_QUANTIZE_IN_FORWARD_ENABLED, \ + WEIGHT_QUANTIZATION, SHARED_PARAMETERS, \ + WEIGHT_QUANTIZE_ENABLED, \ + WEIGHT_QUANTIZE_GROUPS, \ + WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE, \ + WEIGHT_QUANTIZE_CHANGE_RATIO, \ + WEIGHT_QUANTIZE_TYPE, \ + WEIGHT_QUANTIZE_ROUNDING, \ + WEIGHT_QUANTIZE_VERBOSE, \ + WEIGHT_QUANTIZE_KERNEL +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FROZEN_PARAM_FRAGMENTS +from deepspeed.runtime.sparse_tensor import SparseTensor + +from deepspeed.runtime import lr_schedules +from deepspeed.utils import groups +from deepspeed.utils import logger, log_dist, instrument_w_nvtx +from deepspeed.utils.timer import NoopTimer, ThroughputTimer, SynchronizedWallClockTimer, \ + FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, \ + STEP_MICRO_TIMER, \ + FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \ + STEP_GLOBAL_TIMER +from deepspeed.utils.debug import debug_extract_module_and_param_names, debug_clear_module_and_param_names +from deepspeed.monitor.monitor import MonitorMaster +from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop +from deepspeed.runtime.utils import clip_grad_norm_, compare_tensors_in_structures +from deepspeed.runtime.eigenvalue import Eigenvalue +from deepspeed.runtime.data_pipeline.constants import DATA_SAMPLING, \ + DATA_ROUTING, DATA_SAMPLING_ENABLED, CURRICULUM_LEARNING, \ + CURRICULUM_LEARNING_ENABLED, DATA_SAMPLING_NUM_WORKERS, RANDOM_LTD, \ + RANDOM_LTD_ENABLED, RANDOM_LTD_LAYER_ID, RANDOM_LTD_LAYER_NUM, \ + RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE, RANDOM_LTD_LAYER_TOKEN_LR_ENABLED, \ + RANDOM_LTD_GLOBAL_BATCH_SIZE, RANDOM_LTD_MICRO_BATCH_SIZE, DATA_EFFICIENCY +from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler +from deepspeed.runtime.data_pipeline.data_routing.scheduler import RandomLTDScheduler +from deepspeed.runtime.data_pipeline.data_routing.helper import remove_random_ltd_state_dict +from deepspeed.runtime.data_pipeline.data_routing.basic_layer import RandomLayerTokenDrop + +from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine +from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint + +from .pipe.module import PipelineModule +from .utils import get_ma_status +from .compiler import is_compile_supported +from ..ops.adam import FusedAdam +from ..moe.sharded_moe import TopKGate, MOELayer +from ..moe.layer import MoE +from ..moe.utils import is_moe_param, configure_moe_param_groups +from ..git_version_info import version + +from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler +from deepspeed.utils.logging import print_json_dist, print_configuration + +from deepspeed.accelerator import get_accelerator + +from deepspeed.runtime.config import DtypeEnum + +from deepspeed.compile.util import is_deepcompile_supported, get_deepcompile_handle, deepcompile_backward_prologue +from deepspeed.compile.backend import register_compile_pass, opt_passes +from deepspeed.compile.passes import zero3_compile, prefetch, selective_gather, offload_adam_states +from deepspeed.compile.init_z1 import init_z1 +from deepspeed.compile.init_z3 import init_z3 + +MEMORY_OPT_ALLREDUCE_SIZE = 500000000 + +DeepSpeedOptimizerCallable = \ + Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer] +DeepSpeedSchedulerCallable = Callable[[Optimizer], _LRScheduler] + +try: + import apex + from apex import amp + APEX_INSTALLED = True +except ImportError: + # Fail silently so we don't spam logs unnecessarily if user isn't using amp + APEX_INSTALLED = False + + +def split_half_float_double_sparse(tensors): + device_type = get_accelerator().device_name() + supported_types = get_accelerator().supported_dtypes() + + for t in tensors: + assert t.dtype in supported_types, f"attempting to reduce an unsupported grad type: {t.dtype}" + + sparse_tensor_buckets, dense_tensor_buckets = [], [] + for i, dtype in enumerate(supported_types): + sparse_bucket, dense_bucket = [], [] + for t in tensors: + if t.dtype == dtype: + if isinstance(t, SparseTensor): + sparse_bucket.append(t) + else: + dense_bucket.append(t) + if sparse_bucket: + sparse_tensor_buckets.append((dtype, sparse_bucket)) + if dense_bucket: + dense_tensor_buckets.append((dtype, dense_bucket)) + return sparse_tensor_buckets, dense_tensor_buckets + + +class EngineTimers(object): + r"""Wallclock timers for DeepSpeedEngine""" + + def __init__(self, enable_micro_timers, enable_global_timers): + self.forward_timers = [] + self.backward_timers = [] + self.backward_inner_timers = [] + self.backward_reduce_timers = [] + self.step_timers = [] + self.global_timers = [] + self.micro_timers = [] + + if enable_micro_timers: + self.forward_timers += [FORWARD_MICRO_TIMER] + self.backward_timers += [BACKWARD_MICRO_TIMER] + self.backward_inner_timers += [BACKWARD_INNER_MICRO_TIMER] + self.backward_reduce_timers += [BACKWARD_REDUCE_MICRO_TIMER] + self.step_timers += [STEP_MICRO_TIMER] + self.micro_timers += [ + FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, + STEP_MICRO_TIMER + ] + + if enable_global_timers: + self.forward_timers += [FORWARD_GLOBAL_TIMER] + self.backward_timers += [BACKWARD_GLOBAL_TIMER] + self.backward_inner_timers += [BACKWARD_INNER_GLOBAL_TIMER] + self.backward_reduce_timers += [BACKWARD_REDUCE_GLOBAL_TIMER] + self.step_timers += [STEP_GLOBAL_TIMER] + self.global_timers += [ + FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, + STEP_GLOBAL_TIMER + ] + + +class DeepSpeedEngine(Module): + r"""DeepSpeed engine for training.""" + + def __init__(self, + args, + model, + optimizer=None, + model_parameters=None, + training_data=None, + lr_scheduler=None, + mpu=None, + dist_init_required=None, + collate_fn=None, + config=None, + config_class=None, + mesh_device=None, + dont_change_device=False): + super(DeepSpeedEngine, self).__init__() + self.dont_change_device = dont_change_device + self.client_optimizer = optimizer + self.client_lr_scheduler = lr_scheduler + self.training_data = training_data + self.collate_fn = collate_fn + self.mpu = mpu + self.all_to_all_group = None + self.data_parallel_group = None + self.global_steps = 0 + self.global_samples = 0 + self.micro_steps = 0 + self.skipped_steps = 0 + self.gradient_average = True + self.warn_unscaled_loss = True + self.config = config + self._config = config_class + self.loaded_checkpoint_mp_world_size = None + self.loaded_checkpoint_dp_world_size = None + self.enable_backward_allreduce = True + self.inside_no_sync_ctxt = False + self.progressive_layer_drop = None + self.eigenvalue = None + self.block_eigenvalue = None + self.gas_boundary_ctr = 0 + self.dist_backend = get_accelerator().communication_backend_name() + self.has_moe_layers = False + self.num_experts = [] + self.gate_modules = [] + self.moe_layers = [] + self._step_applied = False + self._global_grad_norm = None + self.use_ds_comm = False # False --> Use torch.dist, True --> Use ds.comm backend. + self.checkpoint_engine = None + + self._is_gradient_accumulation_boundary = None + self.scale_wrt_gas = None + self.losses = None + self.mesh_device = mesh_device + + # for debug purposes - can then debug print: debug_get_module_name(module) + debug_extract_module_and_param_names(model) + + if self.mesh_device: + groups.mesh_device = self.mesh_device + + self._do_args_sanity_check(args) + self._configure_with_arguments(args, mpu) + self._do_sanity_check() + if self.autotp_size() > 1: + self._configure_tensor_parallel(model, self.tensor_parallel_config()) + see_memory_usage(f"DeepSpeed Engine: After args sanity test", force=self.memory_breakdown()) + if mpu is not None: + if self.elasticity_enabled(): + if not self.is_elastic_model_parallel_supported(): + assert not self.elasticity_enabled(), ("Elasticity is not currently supported" + " with model parallelism.") + + self._set_distributed_vars(args) + + dist.configure(self._config) + + self.monitor = MonitorMaster(self._config.monitor_config) + + see_memory_usage( + f"DeepSpeed Engine: Before configure distributed model", + force=self.memory_breakdown(), + ) + + self.pipeline_parallelism = isinstance(model, PipelineModule) + + # Configure distributed model + self._configure_distributed_model(model) + + if not self.is_deepcompile_enabled(): + self.module_forward_pre_hook = self._create_module_forward_pre_hook() + self.module_forward_post_hook = self._create_module_forward_post_hook() + + # needed for zero_to_fp32 weights reconstruction to remap nameless data to state_dict + self.param_names = {param: name for name, param in model.named_parameters()} + + self._get_model_parameters() + + see_memory_usage(f"DeepSpeed Engine: After configure distributed model") + + # Configure wall clock timers + self.timers = SynchronizedWallClockTimer() + # Throughput timer + self.tput_timer = ThroughputTimer(self._config.timers_config, + batch_size=self.train_batch_size(), + steps_per_output=self.steps_per_print(), + monitor_memory=False) + + log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0]) + + if self.flops_profiler_enabled(): + self.flops_profiler = FlopsProfiler(self.module, self, self.flops_profiler_recompute_fwd_factor()) + + if training_data: + self.training_dataloader = self.deepspeed_io(training_data) + else: + self.training_dataloader = None + + # Configure optimizer and scheduler + self.optimizer = None + self.basic_optimizer = None + self.lr_scheduler = None + has_optimizer = False + + if optimizer or self.optimizer_name(): + has_optimizer = True + # If no parameters given by init default to module parameters + if model_parameters is None: + model_parameters = self.module.parameters() + + # Convert model parameters from generator to list + if not isinstance(model_parameters, list): + model_parameters = list(model_parameters) + + if has_optimizer: + self._configure_optimizer(optimizer, model_parameters) + self._configure_lr_scheduler() + self._report_progress(0) + elif self.zero_optimization(): + # no optim selected but zero is enabled + self.optimizer = self._configure_zero_optimizer(optimizer=None) + elif self.bfloat16_enabled(): + self.optimizer = self._configure_bf16_optimizer(optimizer=None) + + # Hook optimizer for snip_momentum pruning + if hasattr(model, 'pruners'): + from ..compression.helper import rewrite_optimizer_step + self.optimizer.pruners = model.pruners + rewrite_optimizer_step(self.optimizer) + + # Bookkeeping for sparse support + self.sparse_tensor_module_names = set() + # if self.sparse_gradients_enabled(): + for name, module in self.module.named_modules(): + if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)) and self.sparse_gradients_enabled(): + self.sparse_tensor_module_names.add(name + ".weight") + logger.info("Will convert {} to sparse tensor during training".format(name)) + + self._optimized_linear_offload_setup() + + self.save_non_zero_checkpoint = False + self.save_zero_checkpoint = False + if not isinstance(self.optimizer, DeepSpeedZeRoOffload): + self._configure_checkpointing(dist_init_required) + + if self.eigenvalue_enabled(): + self.eigenvalue = self._configure_eigenvalue() + + if self.pld_enabled(): + self.progressive_layer_drop = self._configure_progressive_layer_drop() + + if self.curriculum_enabled_legacy(): + self.curriculum_scheduler_legacy = self._configure_curriculum_scheduler_legacy() + + if self.random_ltd_enabled(): + random_ltd_config = self.random_ltd_config() + random_ltd_config[RANDOM_LTD_GLOBAL_BATCH_SIZE] = self.train_batch_size() + random_ltd_config[RANDOM_LTD_MICRO_BATCH_SIZE] = self.train_micro_batch_size_per_gpu() + self.random_ltd_scheduler = self._configure_random_ltd_scheduler(random_ltd_config) + + # Engine timers + + self.engine_timers = EngineTimers(enable_micro_timers=self.wall_clock_breakdown(), + enable_global_timers=self.wall_clock_breakdown() + or self.flops_profiler_enabled()) + + if self.global_rank == 0: + self._config.print("DeepSpeedEngine configuration") + if self.dump_state(): + print_configuration(self, "DeepSpeedEngine") + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + + self._is_compiled = False + if is_deepcompile_supported(): + # Predefined compile passes + self.register_compile_pass(zero3_compile.NAME, zero3_compile.add_z3_gather_release) + self.register_compile_pass(prefetch.NAME, prefetch.schedule_prefetch) + self.register_compile_pass(selective_gather.NAME, selective_gather.selective_gather) + self.register_compile_pass(offload_adam_states.NAME, offload_adam_states.move_opt_states) + + def _optimized_linear_offload_setup(self): + self.optimized_linear_base_weight_sharding = False + self.optimized_linear_lora_enabled = False + offload_ratio = None + for _, module in self.module.named_modules(): + if isinstance(module, LoRAOptimizedLinear): + self.optimized_linear_lora_enabled = True + offload_ratio = None + if offload_ratio is not None: + assert offload_ratio == module.lora_config.offload_ratio, \ + "all lora_config offload ratios should be the same across the model" + offload_ratio = module.lora_config.offload_ratio + if module.zero_shards > 1: + # set attr so checkpoint saving can handle BWS properly + self.optimized_linear_base_weight_sharding = True + + if offload_ratio is None: + # Nothing enabled, do nothing + return + + total_params = 0 + for _, p in self.module.named_parameters(): + if hasattr(p, 'ds_optim_param'): + total_params += p.numel() + + offload_limit = total_params * offload_ratio + logger.info(f'offloading {offload_ratio*100}% of eligible params, specifically {offload_limit} params') + total_offloaded = 0 + for _, p in self.module.named_parameters(): + if hasattr(p, 'ds_optim_param'): + if total_offloaded < offload_limit: + total_offloaded += p.numel() + p.ds_offload = True + p.offload() + else: + p.ds_offload = False + + def _configure_tensor_parallel(self, model, tp_config): + self._configure_tensor_parallel_states(model) + configure_tensor_parallel_runtime(tp_config) + + def _configure_tensor_parallel_states(self, model): + """ + Configures the tensor parallel states for the model. + This includes setting up the tensor parallel groups, initializing the TP mesh, + and registering a pre-hook to ensure that the Dataloader inputs are consistent across ranks. + """ + self._set_client_model(model) + # sanity check + # currently, the compatibility between 'autotp' and 'zero > 1' has not been validated + assert self.zero_optimization_stage( + ) <= 2, "Currently, the compatibility between 'autotp' and 'zero_stage = 3' has not been validated" + + self.mpu = groups + self.mpu._init_tp_mesh_device(tensor_model_parallel_size=self.autotp_size()) + + self.first_dataloader_check = None + + def check_dataloader_inputs_same_across_ranks(module, args, kwargs): + + def broadcast_and_check(args, bcast_rank, bcast_group): + if isinstance(args, tuple): + args = list(args) + if len(args) > 0: + if self.mpu.get_tensor_model_parallel_rank() == 0: + _src_args = [args] + dist.broadcast_object_list(object_list=_src_args, + src=bcast_rank, + group=bcast_group, + device=get_accelerator().current_device()) + # Rank 0 does not need to compare with itself + is_equal = True + else: + _src_args = [None] + dist.broadcast_object_list(object_list=_src_args, + src=bcast_rank, + group=bcast_group, + device=get_accelerator().current_device()) + + is_equal = compare_tensors_in_structures(args, _src_args[0]) + + equal_tensor = torch.tensor(is_equal, + dtype=self.communication_data_type, + device=get_accelerator().current_device()) + dist.all_reduce(equal_tensor, group=bcast_group) + assert torch.equal( + equal_tensor, + torch.tensor(groups.get_tensor_model_parallel_world_size(), + dtype=self.communication_data_type, + device=get_accelerator().current_device()) + ), "Data inconsistency within the TP group. Please check the Dataloader implementation to ensure consistency." + + bcast_rank = self.mpu.get_tensor_model_parallel_src_rank() + bcast_group = self.mpu.get_tensor_model_parallel_group() + + broadcast_and_check(args, bcast_rank, bcast_group) + broadcast_and_check(kwargs, bcast_rank, bcast_group) + + logger.info(f":The Dataloader has passed the TP group consistency check.") + self.first_dataloader_check.remove() + + self.first_dataloader_check = self.module.register_forward_pre_hook(check_dataloader_inputs_same_across_ranks, + prepend=True, + with_kwargs=True) + + def destroy(self): + if self.optimizer is not None and hasattr(self.optimizer, 'destroy'): + self.optimizer.destroy() + if self.is_deepcompile_enabled(): + get_deepcompile_handle().cleanup() + debug_clear_module_and_param_names() + + def _get_model_parameters(self): + if self.autotuning_profile_model_info(): + self.autotuning_model_info = {} + num_params = 0 + trainable_num_params = 0 + + for p in self.module.parameters(): + # since user code might call deepspeed.zero.Init() before deepspeed.initialize(), need to check the attribute to check if the parameter is partitioned in zero 3 already or not + n = 0 + if hasattr(p, "ds_tensor"): # if the parameter is partitioned in zero 3 + n += p.ds_numel + else: # if the parameter is not partitioned in zero 3 yet + n += p.numel() + num_params += n + if p.requires_grad: + trainable_num_params += n + if self.global_rank == 0: + self.autotuning_model_info["num_params"] = num_params * self.mp_world_size + self.autotuning_model_info["trainable_num_params"] = trainable_num_params * self.mp_world_size + + logger.info(f"model parameter = {num_params}") + + def get_batch_info(self): + """Get all training batch related settings. + Returns: + train_batch_size (int): The effective training batch size. This is the amount of data + samples that leads to one step of model update. + train_micro_batch_size_per_gpu (int): Batch size to be processed by one GPU in one + step (without gradient accumulation). + gradient_accumulation_steps (int): Number of training steps to accumulate gradients + before averaging and applying them. + """ + return ( + self.train_batch_size, + self.train_micro_batch_size_per_gpu, + self.gradient_accumulation_steps, + ) + + def set_train_batch_size(self, train_batch_size): + """Adjust the global batch size by increasing or decreasing the number of + micro-batches (i.e., gradient accumulation steps). The size of each micro-batch + (i.e., ``train_micro_batch_size_per_gpu``) is not changed. + Args: + train_batch_size (int): The new global batch size for training. + Raises: + ValueError: if ``train_batch_size`` is not divisible by the + configured micro-batch size and data parallelism. + """ + if train_batch_size % (self.train_micro_batch_size_per_gpu() * self.dp_world_size) != 0: + #print(f'{train_batch_size=} {self.train_micro_batch_size_per_gpu()=} {self.dp_world_size=}') + raise ValueError(f'Train batch size must be divisible by micro-batch data parallelism') + new_gas = train_batch_size // (self.train_micro_batch_size_per_gpu() * self.dp_world_size) + # overwrite config + self._config.train_batch_size = train_batch_size + self._config.gradient_accumulation_steps = new_gas + + def set_train_micro_batch_size(self, micro_batch_size): + """Adjust the micro batch size(i.e., the micro batch size in every data parallel group), + while keep the gradient accumulation steps the same. + Args: + micro_batch_size (int): The new micro batch size for training. + """ + # overwrite config + new_global_batch_size = micro_batch_size * self._config.gradient_accumulation_steps * self.dp_world_size + self._config.train_batch_size = new_global_batch_size + self._config.train_micro_batch_size_per_gpu = micro_batch_size + + def set_data_post_process_func(self, post_process_func): + if self.training_dataloader is not None: + self.training_dataloader.post_process_func = post_process_func + + def set_custom_curriculum_learning_schedule(self, schedule_func_dict): + if self.training_dataloader is not None and self.curriculum_learning_enabled(): + self.training_dataloader.data_sampler.set_custom_curriculum_learning_schedule(schedule_func_dict) + + def get_global_grad_norm(self) -> float: + """Return the 2-norm of all gradients. If there is model parallelism, + the norm will be global. + The computed norm will be cached and reused until the next step() pass. + .. note:: + In the presence of model parallelism, this is a collective call + and acts as a barrier among ``mpu.get_model_parallel_group()``. + Returns: + float: norm + """ + return self._global_grad_norm + + def __getattr__(self, name): + """ + Pass through attributes defined in the model if they are not overridden by ds-engine. + """ + + _module = {} + if "module" in self.__dict__: + _module = self.__dict__['module'] + if name in dir(self): + return getattr(self, name) + elif name in dir(_module): + return getattr(_module, name) + else: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") + + def checkpoint_tag_validation_enabled(self): + return self._config.checkpoint_tag_validation_enabled + + def checkpoint_tag_validation_fail(self): + return self._config.checkpoint_tag_validation_fail + + def elasticity_enabled(self): + return self._config.elasticity_enabled + + def is_elastic_model_parallel_supported(self): + if self.elasticity_enabled(): + # Add code for finding number of GPUs per node automatically + if self._config.num_gpus_per_node % self._config.elastic_model_parallel_size == 0: + return True + else: + return False + + def pld_enabled(self): + return self._config.pld_enabled + + def pld_params(self): + return self._config.pld_params + + def pld_theta(self): + return self.pld_params()[PLD_THETA] + + def pld_gamma(self): + return self.pld_params()[PLD_GAMMA] + + def eigenvalue_enabled(self): + return self._config.eigenvalue_enabled + + def eigenvalue_verbose(self): + return self._config.eigenvalue_verbose + + def eigenvalue_max_iter(self): + return self._config.eigenvalue_max_iter + + def eigenvalue_tol(self): + return self._config.eigenvalue_tol + + def eigenvalue_stability(self): + return self._config.eigenvalue_stability + + def eigenvalue_gas_boundary_resolution(self): + return self._config.eigenvalue_gas_boundary_resolution + + def eigenvalue_layer_name(self): + return self._config.eigenvalue_layer_name + + def eigenvalue_layer_num(self): + return self._config.eigenvalue_layer_num + + def curriculum_enabled_legacy(self): + return self._config.curriculum_enabled_legacy + + def curriculum_params_legacy(self): + return self._config.curriculum_params_legacy + + def data_efficiency_enabled(self): + return self._config.data_efficiency_enabled + + def data_efficiency_config(self): + return self._config.data_efficiency_config + + def data_sampling_enabled(self): + return self._config.data_efficiency_config[DATA_SAMPLING][DATA_SAMPLING_ENABLED] + + def data_sampling_config(self): + return self._config.data_efficiency_config[DATA_SAMPLING] + + def curriculum_learning_enabled(self): + return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED] + + def curriculum_learning_config(self): + return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING] + + def random_ltd_enabled(self): + return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD][RANDOM_LTD_ENABLED] + + def random_ltd_config(self): + return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD] + + def random_ltd_initialize(self): + assert self.random_ltd_enabled() + random_ltd_config = self.random_ltd_config() + random_ltd_queue = deque([x for x in sorted(random_ltd_config[RANDOM_LTD_LAYER_ID])]) + count = 0 + for name, layer in self.module.named_modules(): + if isinstance(layer, RandomLayerTokenDrop): + if len(random_ltd_queue) != 0 and str(random_ltd_queue[0]) in name: ###[1,2,3] + layer.init_config(random_ltd_config, self.random_ltd_scheduler, count) + random_ltd_queue.popleft() + count += 1 + + if random_ltd_config[RANDOM_LTD_LAYER_NUM] != count: + raise ValueError(f'random_ltd_layer_num {random_ltd_config[RANDOM_LTD_LAYER_NUM]} must be \ + equivalent to the len of random_ltd_layer_id {count}') + + if random_ltd_config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]: + assert self.client_lr_scheduler is None + raise ValueError(f'not yet support') + #self.lr_scheduler = lr_schedules.WarmupLayerTokenDecayLR(self.optimizer, self.random_ltd_scheduler) + + def get_sequence_parallel_group(self): + return self.seq_parallel_group + + def wall_clock_breakdown(self): + return self._config.wall_clock_breakdown + + def flops_profiler_enabled(self): + return self._config.flops_profiler_config.enabled or self.autotuning_enabled() + + def flops_profiler_recompute_fwd_factor(self): + return self._config.flops_profiler_config.recompute_fwd_factor + + def flops_profiler_profile_step(self): + step = self._config.flops_profiler_config.profile_step + if self._config.autotuning_config.enabled: + step = self.autotuning_start_profile_step() + return step + + def flops_profiler_module_depth(self): + return self._config.flops_profiler_config.module_depth + + def flops_profiler_top_modules(self): + return self._config.flops_profiler_config.top_modules + + def flops_profiler_detailed(self): + if self._config.autotuning_config.enabled: + return False + return self._config.flops_profiler_config.detailed + + def flops_profiler_output_file(self): + return self._config.flops_profiler_config.output_file + + def memory_breakdown(self): + return self._config.memory_breakdown + + def autotuning_enabled(self): + return self._config.autotuning_config.enabled + + def autotuning_start_profile_step(self): + return self._config.autotuning_config.start_profile_step + + def autotuning_end_profile_step(self): + return self._config.autotuning_config.end_profile_step + + def autotuning_metric_path(self): + path = self._config.autotuning_config.metric_path + if not path: + path = os.path.join(os.getcwd(), "autotuning_metric.json") + return path + + def autotuning_model_info_path(self): + path = self._config.autotuning_config.model_info_path + if not path: + path = os.path.join(os.getcwd(), "autotuning_model_info.json") + return path + + def autotuning_metric(self): + return self._config.autotuning_config.metric + + def autotuning_profile_model_info(self): + return self.autotuning_enabled( + ) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.get( + "profile", False) + + def sparse_gradients_enabled(self): + return self._config.sparse_gradients_enabled + + def train_batch_size(self): + return self._config.train_batch_size + + def train_micro_batch_size_per_gpu(self): + return self._config.train_micro_batch_size_per_gpu + + def optimizer_name(self): + return (self.client_optimizer.__class__.__name__ if self.client_optimizer else self._config.optimizer_name) + + def optimizer_params(self): + return self._config.optimizer_params + + def optimizer_legacy_fusion(self): + return self._config.optimizer_legacy_fusion + + def scheduler_name(self): + return self._config.scheduler_name + + def scheduler_params(self): + return self._config.scheduler_params + + def quantize_training(self): + return ( + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] + [WEIGHT_QUANTIZE_IN_FORWARD_ENABLED], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ENABLED], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_GROUPS], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] + [WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_CHANGE_RATIO], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_TYPE], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ROUNDING], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_VERBOSE], + self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_KERNEL], + ) + + def zero_optimization(self): + return self._config.zero_enabled + + def zero_allow_untested_optimizer(self): + return self._config.zero_allow_untested_optimizer + + def zero_force_ds_cpu_optimizer(self): + return self._config.zero_force_ds_cpu_optimizer + + def zero_reduce_scatter(self): + return self._config.zero_config.reduce_scatter + + def zero_overlap_comm(self): + return self._config.zero_config.overlap_comm + + def zero_offload_optimizer(self): + return self._config.zero_config.offload_optimizer + + def zero_offload_param(self): + return self._config.zero_config.offload_param + + def zero_use_cpu_optimizer(self): + if self._config.zero_config.offload_optimizer is not None: + return self._config.zero_config.offload_optimizer.device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme] + return False + + def zero_cpu_offload(self): + if self._config.zero_config.offload_optimizer is not None: + return self._config.zero_config.offload_optimizer.device == OffloadDeviceEnum.cpu + return False + + def zero_partial_offload(self): + return getattr(self._config.zero_config.offload_optimizer, "ratio", 1.0) + + def zero_sub_group_size(self): + return self._config.zero_config.sub_group_size + + def zero_optimization_stage(self): + return self._config.zero_optimization_stage + + def mics_shard_size(self): + return self._config.mics_shard_size + + def zero_reduce_bucket_size(self): + return self._config.zero_config.reduce_bucket_size + + def zero_multi_rank_bucket_allreduce(self): + return self._config.zero_config.use_multi_rank_bucket_allreduce + + def zero_allgather_bucket_size(self): + return self._config.zero_config.allgather_bucket_size + + def zero_optimization_partition_gradients(self): + return self.zero_optimization_stage() >= ZeroStageEnum.gradients + + def zero_optimization_partition_weights(self): + return self.zero_optimization_stage() >= ZeroStageEnum.weights + + def is_first_weights_partition_group(self): + ret = True if self.mics_shard_size() < 0 \ + and self.zero_optimization_partition_weights() else False + if self.mics_shard_size() > 0 and self.global_rank < self.mics_shard_size(): + ret = True + return ret + + def zero_contiguous_gradients(self): + return self._config.zero_config.contiguous_gradients + + def zero_load_from_fp32_weights(self): + return self._config.zero_config.load_from_fp32_weights + + def zero_elastic_checkpoint(self): + return self._config.zero_config.elastic_checkpoint + + def zero_nvme_offload_optimizer(self): + return getattr(self.optimizer, "swap_optimizer", False) + + def zero_max_live_parameters(self): + return self._config.zero_config.max_live_parameters + + def zero_max_reuse_distance(self): + return self._config.zero_config.max_reuse_distance + + def zero_prefetch_bucket_size(self): + return self._config.zero_config.prefetch_bucket_size + + def zero_module_granularity_threshold(self): + return self._config.zero_config.module_granularity_threshold + + def zero_param_persistence_threshold(self): + return self._config.zero_config.param_persistence_threshold + + def zero_model_persistence_threshold(self): + return self._config.zero_config.model_persistence_threshold + + def zero_gather_16bit_weights_on_model_save(self): + return self._config.zero_config.gather_16bit_weights_on_model_save + + def zero_grad_hooks(self): + return self._config.zero_config.grad_hooks + + def zero_legacy_stage1(self): + return self._config.zero_config.legacy_stage1 + + def zero_ignore_unused_parameters(self): + return self._config.zero_config.ignore_unused_parameters + + def tensor_parallel_config(self): + return self._config.tensor_parallel_config + + def autotp_size(self): + return self._config.tensor_parallel_config.autotp_size + + def graph_harvesting(self): + return self._config.graph_harvesting + + def fp16_enabled(self): + return self._config.fp16_enabled + + def bfloat16_enabled(self): + return self._config.bfloat16_enabled + + def fp16_master_weights_and_gradients(self): + return self._config.fp16_master_weights_and_gradients + + def amp_enabled(self): + return self._config.amp_enabled + + def amp_params(self): + return self._config.amp_params + + def fp16_auto_cast(self): + return self._config.fp16_auto_cast + + def loss_scale(self): + return self._config.loss_scale + + def gradient_accumulation_steps(self): + return self._config.gradient_accumulation_steps + + def use_node_local_storage(self): + return self._config.use_node_local_storage + + def load_universal_checkpoint(self): + return self._config.load_universal_checkpoint + + @property + def communication_data_type(self): + res = self._config.communication_data_type + if res is not None: + return res + + if self.fp16_enabled(): + return torch.float16 + + if self.bfloat16_enabled(): + return torch.bfloat16 + + return torch.float32 + + @communication_data_type.setter + def communication_data_type(self, value): + self._config.communication_data_type = value + + def postscale_gradients(self): + return not self._config.prescale_gradients + + def gradient_predivide_factor(self): + return self._config.gradient_predivide_factor + + def steps_per_print(self): + return self._config.steps_per_print + + def zero_allgather_partitions(self): + return self._config.zero_config.allgather_partitions + + def zero_round_robin_gradients(self): + return self._config.zero_config.round_robin_gradients + + def zero_hpz_partition_size(self): + return self._config.zero_config.zero_hpz_partition_size + + def zero_quantized_weights(self): + return self._config.zero_config.zero_quantized_weights + + def zero_quantized_nontrainable_weights(self): + return self._config.zero_config.zero_quantized_nontrainable_weights + + def zero_quantized_gradients(self): + return self._config.zero_config.zero_quantized_gradients + + def zeropp_loco_param(self): + return self._config.zero_config.zeropp_loco_param + + def zero_log_trace_cache_warnings(self): + return self._config.zero_config.log_trace_cache_warnings + + def dump_state(self): + return self._config.dump_state + + def gradient_clipping(self): + return self._config.gradient_clipping + + def dynamic_loss_scale(self): + return self._config.loss_scale == 0 + + def initial_dynamic_scale(self): + return self._config.initial_dynamic_scale + + def dynamic_loss_scale_args(self): + return self._config.dynamic_loss_scale_args + + def swap_tensor_config(self): + return self._config.swap_tensor_config + + def aio_config(self): + return self._config.aio_config + + def get_data_types(self): + model_dtype = torch.float32 + if self.fp16_enabled(): + model_dtype = torch.float16 + elif self.bfloat16_enabled(): + model_dtype = torch.bfloat16 + + if self._config.grad_accum_dtype is None: + if model_dtype == torch.bfloat16 and not self.zero_optimization(): + grad_accum_dtype = torch.float32 + else: + grad_accum_dtype = model_dtype + else: + grad_accum_dtype = DtypeEnum(self._config.grad_accum_dtype).value + + return (model_dtype, grad_accum_dtype) + + def _optimizer_has_ckpt_event_prologue(self): + return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_prologue') + + def _optimizer_has_ckpt_event_epilogue(self): + return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_epilogue') + + def _configure_lr_scheduler(self): + if self.client_lr_scheduler: + if isinstance(self.client_lr_scheduler, Callable): + log_dist('DeepSpeed using client callable to create LR scheduler', ranks=[0]) + self.lr_scheduler = self.client_lr_scheduler(self.basic_optimizer) + else: + log_dist('DeepSpeed using client LR scheduler', ranks=[0]) + self.lr_scheduler = self.client_lr_scheduler + else: + # load lr scheduler from json configuration if lr scheduler is not defined and passed in + lr_scheduler = self._scheduler_from_config(self.optimizer) + log_dist(f"DeepSpeed using configured LR scheduler = {self.scheduler_name()}", ranks=[0]) + self.lr_scheduler = lr_scheduler + + log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0]) + + def _configure_checkpointing(self, dist_init_required): + self.checkpoint_engine = TorchCheckpointEngine() + + if self._config is not None and self._config.nebula_config.enabled: + try: + from deepspeed.runtime.checkpoint_engine.nebula_checkpoint_engine import \ + NebulaCheckpointEngine + self.checkpoint_engine = NebulaCheckpointEngine(config_params=self._config.nebula_config) + except ImportError as err: + logger.error(f"No torch_nebula was found! Will fall back to torch.save. Details: {err}") + self.checkpoint_engine = TorchCheckpointEngine() + + dp_rank = groups._get_sequence_data_parallel_rank() + + rank = self.local_rank if self.use_node_local_storage() else dp_rank + + # only the first data parallel process needs to store the model checkpoint + # if you want to use node local storage this must be done by rank 0 on each + # node + self.save_non_zero_checkpoint = (rank == 0) or (self.zero_optimization_partition_weights() + and self.is_first_weights_partition_group()) + + if self.zero_optimization() or self.bfloat16_enabled(): + param_rank = dist.get_rank(group=self.optimizer.dp_process_group) + + # Only the first parameter parallel process needs to store the + # optimizer state checkpoints for zero + self.save_zero_checkpoint = param_rank == dp_rank + + def _scheduler_from_config(self, optimizer): + scheduler_name = self.scheduler_name() + if scheduler_name is not None: + if hasattr(lr_schedules, scheduler_name): + scheduler = getattr(lr_schedules, scheduler_name) + else: + assert hasattr(torch.optim.lr_scheduler, + scheduler_name), f"DeepSpeed does not recognize LR scheduler {scheduler_name}" + + scheduler = getattr(torch.optim.lr_scheduler, scheduler_name) + + scheduler_params = self.scheduler_params() + instantiated_scheduler = scheduler(optimizer, **scheduler_params) + return instantiated_scheduler + else: + return None + + def _set_distributed_vars(self, args): + device_rank = args.device_rank if args is not None and hasattr(args, 'device_rank') else self.local_rank + if device_rank >= 0: + get_accelerator().set_device(device_rank) + self.device = torch.device(get_accelerator().device_name(device_rank)) + self.world_size = dist.get_world_size() + self.global_rank = dist.get_rank() + else: + self.world_size = 1 + self.global_rank = 0 + self.device = get_accelerator().device() + + # Configure based on command line arguments + def _configure_with_arguments(self, args, mpu): + # After the distributed backend is initialized we are guaranteed the LOCAL_RANK + # environment variable is set. We must align args.local_rank to this value for + # backwards compatibility with scripts relying on [args|self].local_rank containing + # the correct local rank info. _do_args_sanity_check will ensure this is the case. + + if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ: + ompi_local_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK") + local_rank = os.environ.get('LOCAL_RANK', ompi_local_rank) + assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({ompi_local_rank}), " \ + "not sure how to proceed as we're seeing conflicting local rank info." + os.environ['LOCAL_RANK'] = local_rank + + self.local_rank = int(os.environ['LOCAL_RANK']) + if hasattr(args, 'local_rank'): + args.local_rank = self.local_rank + + # Validate command line arguments + def _do_args_sanity_check(self, args): + assert "LOCAL_RANK" in os.environ or "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ, "DeepSpeed requires the LOCAL_RANK environment " \ + "variable, it is set by the deepspeed launcher, deepspeed.init_distributed, or the torch's launcher. If using a " \ + "different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed." + + if hasattr(args, 'local_rank') and args.local_rank is not None: + assert isinstance(args.local_rank, + int), f"args.local_rank of {args.local_rank} is an unknown type {type(args.local_rank)}" + if args.local_rank >= 0: + env_local_rank = int(os.environ.get("LOCAL_RANK")) + assert ( + env_local_rank == args.local_rank + ), f"Mismatch in local rank setting, args.local_rank={args.local_rank} but env['LOCAL_RANK']={env_local_rank}." + + def _is_supported_optimizer(self, optimizer_name): + return (optimizer_name in DEEPSPEED_OPTIMIZERS or getattr(torch.optim, optimizer_name, None) is not None) + + def _supported_optims(self): + FairseqOptimizer = None + try: + from fairseq.optim.fairseq_optimizer import FairseqOptimizer + except ImportError: + pass + + expected_optim_types = [Optimizer] + if FairseqOptimizer: + # fairseq optims are not torch.optim objects + expected_optim_types.append(FairseqOptimizer) + return expected_optim_types + + # Validate configuration based on command line arguments + def _do_sanity_check(self): + if self.fp16_enabled() and not get_accelerator().is_fp16_supported(): + raise ValueError("Type fp16 is not supported on your device.") + + if self.bfloat16_enabled() and not get_accelerator().is_bf16_supported(): + raise ValueError("Type bf16 is not supported on your device.") + + expected_optim_types = self._supported_optims() + expected_optim_types += [type(None), Callable] + assert isinstance(self.client_optimizer, tuple(expected_optim_types)), \ + f'Client Optimizer is of unexpected type {type(self.client_optimizer)}' + + if not self.client_optimizer: + if self.optimizer_name() is not None: + assert self._is_supported_optimizer( + self.optimizer_name()), "{} is not a supported DeepSpeed Optimizer".format(self.optimizer_name()) + + if (self.optimizer_name() == LAMB_OPTIMIZER or self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER): + assert (self.dynamic_loss_scale()), "DeepSpeed {} optimizer requires dynamic loss scaling".format( + self.optimizer_name()) + + # Detect invalid combinations of client optimizer and client scheduler + if isinstance(self.client_lr_scheduler, _LRScheduler): + assert isinstance(self.client_optimizer, Optimizer), \ + f'Client Optimizer (type = {type(self.client_optimizer)} is not instantiated but Client LR Scheduler is instantiated' + + def _broadcast_model(self): + + def is_replicated(p): + if hasattr(p, "ds_status") and p.ds_status is not ZeroParamStatus.AVAILABLE: + return False + elif hasattr(p, 'ds_optim_param'): + # do not broadcast OptimizedLinear parameters, they are unique per base weight shard + return False + return True + + for n, p in self.module.named_parameters(): + # Broadcast the model for different parameters + if is_moe_param(p): + if torch.is_tensor(p) and is_replicated(p): + dist.broadcast(p.data, + groups._get_expert_broadcast_src_rank(p.group_name), + group=self.expert_data_parallel_group[p.group_name]) + else: + if torch.is_tensor(p) and is_replicated(p): + dist.broadcast(p.data, groups._get_broadcast_src_rank(), group=self.seq_data_parallel_group) + + @staticmethod + def __check_params(model: Module, dtype: torch.dtype) -> None: + return + if not all(param.dtype == dtype for param in model.parameters()) and dist.get_rank() == 0: + raise ValueError(f"{dtype} is enabled but the following parameters have dtype that is " + f"not {dtype}: " + f"{[(n, p.dtype) for n, p in model.named_parameters() if p.dtype != dtype]}") + + def _set_client_model(self, model): + # register client model in _modules so that nn.module methods work correctly + modules = self.__dict__.get('_modules') + modules['module'] = model + # register module attribute in engine but avoid getattr + self.__dict__['module'] = model + + def _configure_distributed_model(self, model): + self._set_client_model(model) + is_zero_init_model = self.zero_optimization_partition_weights() and any( + [hasattr(param, "ds_id") for param in self.module.parameters()]) + + if self.fp16_enabled(): + if is_zero_init_model: + self.__check_params(self.module, torch.half) + self.module.half() + elif self.bfloat16_enabled(): + if is_zero_init_model: + self.__check_params(self.module, torch.bfloat16) + self.module.bfloat16() + else: + self.__check_params(self.module, torch.float) + + # zero.Init() handles device placement of model + if not (self.dont_change_device or is_zero_init_model): + self.module.to(self.device) + + # MoE related initialization + for _, module in self.module.named_modules(): + if isinstance(module, MoE): + self.has_moe_layers = True + self.num_experts.append(module.num_experts) + + if self.has_moe_layers: + for _, module in self.module.named_modules(): + if isinstance(module, TopKGate): + self.gate_modules.append(module) + if self.wall_clock_breakdown(): + module.wall_clock_breakdown = True + if isinstance(module, MOELayer): + self.moe_layers.append(module) + if self.wall_clock_breakdown(): + module.wall_clock_breakdown = True + + # Pass the mpu from here to groups. For subsequent use, just query groups + if self.mpu is not None: + groups.mpu = self.mpu + + # Set deepspeed parallelism spec. for the model including expert parallelism + for _, module in self.module.named_modules(): + if hasattr(module, 'set_deepspeed_parallelism'): + module.set_deepspeed_parallelism(self._config.use_data_before_expert_parallel_) + + # Query the groups module to get information about various parallel groups + self.local_all_to_all_group = None + if self.zero_quantized_gradients(): + message = "Using LoCo quantized gradients" if self.zeropp_loco_param() else "Using quantized gradients" + log_dist(message, ranks=[0]) + self.local_all_to_all_group = groups._get_local_all_to_all_group() + self.data_parallel_group = groups._get_data_parallel_group() + self.dp_world_size = groups._get_data_parallel_world_size() + self.seq_data_parallel_group = groups._get_sequence_data_parallel_group() + self.seq_dp_world_size = groups._get_sequence_data_parallel_world_size() + self.mp_world_size = groups._get_model_parallel_world_size() + self.expert_parallel_group = groups._get_expert_parallel_group_dict() + self.expert_data_parallel_group = groups._get_expert_data_parallel_group_dict() + self.sequence_parallel_size = groups._get_sequence_parallel_world_size() + if self.sequence_parallel_size > 1: + self.communication_data_type = self._config.seq_parallel_communication_data_type + self.seq_parallel_group = groups._get_sequence_parallel_group() + + if dist.get_rank() == 0: + summary = "********** distributed groups summary **********\n" + summary += f"\t {self.dp_world_size=}\n" + summary += f"\t {self.mp_world_size=}\n" + summary += f"\t {self.seq_dp_world_size=}\n" + summary += f"\t {self.sequence_parallel_size=}\n" + summary += "***********************************************" + logger.info(summary) + + if not (self.amp_enabled() or is_zero_init_model): + self._broadcast_model() + + # check if parameters are duplicated in optimizer param_groups + def _check_for_duplicates(self, optimizer): + for name, param in self.module.named_parameters(): + param_id = id(param) + + def ids_list(group): + return [id(param) for param in group] + + occurrence = sum([ + ids_list(group['params']).count(param_id) if param_id in ids_list(group['params']) else 0 + for group in optimizer.param_groups + ]) + assert occurrence <= 1, f"Parameter with name: {name} occurs multiple times in optimizer.param_groups. Make sure it only appears once to prevent undefined behavior." + + def _do_optimizer_sanity_check(self, basic_optimizer): + model_dtype, grad_accum_dtype = self.get_data_types() + zero_enabled = self.zero_optimization() + amp_enabled = self.amp_enabled() + # config based assertions + assert ( + not (amp_enabled and zero_enabled) + ), "Amp and ZeRO are not currently compatible, please use (legacy) fp16 mode which performs similar to amp opt_mode=O2" + if zero_enabled: + if not is_zero_supported_optimizer(basic_optimizer): + assert ( + self.zero_allow_untested_optimizer() + ), 'You are using an untested ZeRO Optimizer. Please add <"zero_allow_untested_optimizer": true> in the configuration file to use it.' + + if self.global_rank == 0: + logger.warning("**** You are using ZeRO with an untested optimizer, proceed with caution *****") + if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage( + ) == 1 and not self.zero_cpu_offload(): + return BFLOAT16 + return ZERO_OPTIMIZATION + elif amp_enabled: + if model_dtype != grad_accum_dtype: + raise NotImplementedError( + "Model data type and gradient accumulation data type must be equal to use Amp") + if model_dtype == torch.bfloat16 or model_dtype == torch.float16: + raise NotImplementedError("Cannot enable both amp with (legacy) fp16 or bfloat16 mode") + try: + logger.info("Initializing Apex amp from: {}".format(amp.__path__)) + except NameError: + # If apex/amp is available it will be imported above + raise RuntimeError("Unable to import apex/amp, please make sure it is installed") + return AMP + # data type checks + elif model_dtype == grad_accum_dtype: + if model_dtype == torch.bfloat16: + if self.pipeline_parallelism: + logger.warning( + "**** BF16 gradient accumulation is not safe numerically with large number of accumulation steps, proceed with caution *****" + ) + return BFLOAT16 + else: + raise NotImplementedError( + "Bfloat16 wrapper must use a gradient accumulation type of fp32, enable ZeRO to use Bfloat16 gradient accumulation" + ) + if model_dtype == torch.float16: + return FP16 + # else optimizer_wrapper = None + elif model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32: + return BFLOAT16 + else: + raise NotImplementedError("unsupported mix of model dtype and gradient accumulation type") + + return None + + # Configure optimizer + def _configure_optimizer(self, client_optimizer, model_parameters): + if client_optimizer is None: + if self.has_moe_layers: + model_parameters = configure_moe_param_groups(model_parameters) + basic_optimizer = self._configure_basic_optimizer(model_parameters) + log_dist(f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", ranks=[0]) + else: + if isinstance(client_optimizer, tuple(self._supported_optims())): + basic_optimizer = client_optimizer + log_dist('Using client Optimizer as basic optimizer', ranks=[0]) + else: + basic_optimizer = client_optimizer(model_parameters) + log_dist('Using client callable to create basic optimizer', ranks=[0]) + + if self.zero_use_cpu_optimizer() and not isinstance(basic_optimizer, deepspeed.ops.adam.DeepSpeedCPUAdam): + if self.zero_force_ds_cpu_optimizer(): + msg = f'You are using ZeRO-Offload with a client provided optimizer ({type(basic_optimizer)}) which in most cases will yield poor performance. Please either use deepspeed.ops.adam.DeepSpeedCPUAdam or set an optimizer in your ds-config (https://www.deepspeed.ai/docs/config-json/#optimizer-parameters). If you really want to use a custom optimizer w. ZeRO-Offload and understand the performance impacts you can also set <"zero_force_ds_cpu_optimizer": false> in your configuration file.' + raise ZeRORuntimeException(msg) + + basic_optimizer.param_groups[:] = [pg for pg in basic_optimizer.param_groups if len(pg["params"]) != 0] + log_dist("Removing param_group that has no 'params' in the basic Optimizer", ranks=[0]) + + self._check_for_duplicates(basic_optimizer) + + self.basic_optimizer = basic_optimizer + log_dist("DeepSpeed Basic Optimizer = {}".format(basic_optimizer.__class__.__name__), ranks=[0]) + + optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer) + + if optimizer_wrapper == ZERO_OPTIMIZATION: + self.optimizer = self._configure_zero_optimizer(basic_optimizer) + elif optimizer_wrapper == AMP: + amp_params = self.amp_params() + log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0]) + model, self.optimizer = amp.initialize(self.module, basic_optimizer, **amp_params) + self._set_client_model(model) + self._broadcast_model() + # TODO: maybe need to broadcast experts differently? + elif optimizer_wrapper == FP16: + self.optimizer = self._configure_fp16_optimizer(basic_optimizer) + elif optimizer_wrapper == BFLOAT16: + self.optimizer = self._configure_bf16_optimizer(basic_optimizer) + else: + self.optimizer = basic_optimizer + + log_dist("DeepSpeed Final Optimizer = {}".format(self.optimizer.__class__.__name__), ranks=[0]) + + self.compression_scheduler = self._configure_compression_scheduler() + self.quantizer = self._configure_quantization() + + def _configure_basic_optimizer(self, model_parameters): + optimizer_parameters = self.optimizer_params() + if optimizer_parameters is None: + optimizer_parameters = {} + # print(optimizer_parameters.keys()) + if "max_grad_norm" in optimizer_parameters.keys(): + raise ValueError( + "'max_grad_norm' is not supported as an optimizer parameter, please switch to using the deepspeed parameter 'gradient_clipping' see: https://www.deepspeed.ai/docs/config-json/#gradient-clipping for more details" + ) + + if self.optimizer_name() in [ADAM_OPTIMIZER, ADAMW_OPTIMIZER]: + torch_adam = optimizer_parameters.pop(TORCH_ADAM_PARAM, False) + adam_w_mode = optimizer_parameters.pop(ADAM_W_MODE, ADAM_W_MODE_DEFAULT) + + # Optimizer name of Adam forces AdamW logic unless adam_w_mode is explicitly set + effective_adam_w_mode = self.optimizer_name() == ADAMW_OPTIMIZER or adam_w_mode + + if torch_adam: + if not effective_adam_w_mode: + optimizer = torch.optim.Adam(model_parameters, **optimizer_parameters) + else: + optimizer = torch.optim.AdamW(model_parameters, **optimizer_parameters) + else: + if self.zero_use_cpu_optimizer(): + from deepspeed.ops.adam import DeepSpeedCPUAdam + optimizer = DeepSpeedCPUAdam(model_parameters, + **optimizer_parameters, + adamw_mode=effective_adam_w_mode) + else: + from deepspeed.ops.adam import FusedAdam + + optimizer = FusedAdam( + model_parameters, + **optimizer_parameters, + adam_w_mode=effective_adam_w_mode, + ) + + elif self.optimizer_name() == ADAGRAD_OPTIMIZER: + if self.zero_use_cpu_optimizer(): + from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad + optimizer = DeepSpeedCPUAdagrad(model_parameters, **optimizer_parameters) + else: + optimizer = torch.optim.Adagrad(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == LAMB_OPTIMIZER: + from deepspeed.ops.lamb import FusedLamb + + optimizer = FusedLamb(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == ONEBIT_ADAM_OPTIMIZER: + assert not self.zero_optimization(), "1bit-Adam is not compatible with ZeRO" + from deepspeed.runtime.fp16.onebit.adam import OnebitAdam + + optimizer = OnebitAdam(model_parameters, self, **optimizer_parameters) + if not self.fp16_enabled(): + logger.warning(f"Currently the convergence of 1-bit Adam is only verified under FP16") + elif self.optimizer_name() == ZERO_ONE_ADAM_OPTIMIZER: + assert not self.zero_optimization(), "0/1 Adam is not compatible with ZeRO" + from deepspeed.runtime.fp16.onebit.zoadam import ZeroOneAdam + + optimizer = ZeroOneAdam(model_parameters, self, **optimizer_parameters) + if not self.fp16_enabled(): + logger.warning(f'Currently the convergence of 0/1 Adam is only verified under FP16') + elif self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER: + assert not self.zero_optimization(), "1bit-Lamb is not compatible with ZeRO" + from deepspeed.runtime.fp16.onebit.lamb import OnebitLamb + + optimizer = OnebitLamb(model_parameters, self, **optimizer_parameters) + if not self.fp16_enabled(): + logger.warning(f"Currently the convergence of 1-bit Lamb is only verified under FP16") + elif self.optimizer_name() == LION_OPTIMIZER: + if self.zero_use_cpu_optimizer(): + from deepspeed.ops.lion import DeepSpeedCPULion + optimizer = DeepSpeedCPULion(model_parameters, **optimizer_parameters) + else: + from deepspeed.ops.lion import FusedLion + optimizer = FusedLion(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == MUADAM_OPTIMIZER: + try: + from mup import MuAdam + except ImportError: + logger.error(f"Install mup to use MuAdam optimizer") + optimizer = MuAdam(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == MUADAMW_OPTIMIZER: + try: + from mup import MuAdamW + except ImportError: + logger.error(f"Install mup to use MuAdamW optimizer") + optimizer = MuAdamW(model_parameters, **optimizer_parameters) + elif self.optimizer_name() == MUSGD_OPTIMIZER: + try: + from mup import MuSGD + except ImportError: + logger.error(f"Install mup to use MuSGD optimizer") + optimizer = MuSGD(model_parameters, **optimizer_parameters) + else: + torch_optimizer = getattr(torch.optim, self.optimizer_name()) + optimizer = torch_optimizer(model_parameters, **optimizer_parameters) + return optimizer + + def _configure_compression_scheduler(self): + return compression_scheduler(self.module, self._config.compression_config) + + def _configure_random_ltd_scheduler(self, configs): + return RandomLTDScheduler(configs) + + def _configure_quantization(self): + ( + quantize_weight_in_forward, + quantize_enabled, + q_groups, + q_mixed_fp16, + q_change_ratio, + q_type, + q_rounding, + q_verbose, + use_quantizer_kernel, + ) = self.quantize_training() + if quantize_enabled and not quantize_weight_in_forward: + assert self.fp16_enabled( + ), "MoQ (quantize in optimization step) weight quantization is only supported for FP16" + quantizer = None + if quantize_enabled and not quantize_weight_in_forward: + from deepspeed.runtime.quantize import Quantizer + + quantizer = Quantizer( + q_groups, + q_mixed_fp16, + q_change_ratio, + q_type, + q_rounding, + q_verbose, + self.eigenvalue_enabled(), + use_quantizer_kernel, + self.eigenvalue_layer_num() if self.eigenvalue_enabled() else 0, + ) + return quantizer + + def _configure_fp16_optimizer(self, optimizer): + initial_dynamic_scale = self.initial_dynamic_scale() + dynamic_loss_args = self.dynamic_loss_scale_args() + clip_grad = self.gradient_clipping() + if APEX_INSTALLED: + fused_opts = (apex.optimizers.FusedAdam, FusedAdam) + else: + fused_opts = FusedAdam + if isinstance(optimizer, fused_opts) \ + or self.optimizer_name() in [ONEBIT_ADAM_OPTIMIZER, ZERO_ONE_ADAM_OPTIMIZER]: + if self.dynamic_loss_scale(): + log_dist(f'Creating fp16 optimizer with dynamic loss scale', ranks=[0]) + timers = self.timers if self.wall_clock_breakdown() else NoopTimer() + optimizer = FP16_Optimizer( + optimizer, + deepspeed=self, + dynamic_loss_scale=True, + initial_dynamic_scale=initial_dynamic_scale, + dynamic_loss_args=dynamic_loss_args, + mpu=self.mpu, + clip_grad=clip_grad, + fused_adam_legacy=self.optimizer_legacy_fusion(), + timers=timers, + has_moe_layers=self.has_moe_layers, + ) + else: + log_dist(f'Creating fp16 optimizer with static loss scale: {self.loss_scale()}', ranks=[0]) + optimizer = FP16_Optimizer( + optimizer, + deepspeed=self, + static_loss_scale=self.loss_scale(), + mpu=self.mpu, + clip_grad=clip_grad, + fused_adam_legacy=self.optimizer_legacy_fusion(), + has_moe_layers=self.has_moe_layers, + ) + else: + log_dist(f'Creating fp16 unfused optimizer with dynamic loss scale', ranks=[0]) + optimizer = FP16_UnfusedOptimizer( + optimizer, + deepspeed=self, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=dynamic_loss_args, + mpu=self.mpu, + clip_grad=clip_grad, + fused_lamb_legacy=self.optimizer_name() == LAMB_OPTIMIZER, + ) + + return optimizer + + def _configure_bf16_optimizer(self, optimizer): + clip_grad = self.gradient_clipping() + + if optimizer is None: + optimizer = DummyOptim(list(self.module.parameters())) + + log_dist('Creating BF16 optimizer', ranks=[0]) + + timers = self.timers if self.wall_clock_breakdown() else NoopTimer() + optimizer = BF16_Optimizer(optimizer, + self.param_names, + mpu=self.mpu, + clip_grad=clip_grad, + allgather_bucket_size=self.zero_allgather_bucket_size(), + dp_process_group=self.seq_data_parallel_group, + timers=timers, + grad_acc_dtype=self.get_data_types()[1], + graph_harvesting=self.graph_harvesting(), + immediate_grad_update=self._config.bfloat16_immediate_grad_update, + has_moe_layers=self.has_moe_layers) + + return optimizer + + def _configure_zero_optimizer(self, optimizer): + zero_stage = self.zero_optimization_stage() + + mics_shard_size = self.mics_shard_size() + model_dtype, gradient_accumulation_dtype = self.get_data_types() + + timers = self.timers if self.wall_clock_breakdown() else NoopTimer() + + if optimizer is None: + optimizer = DummyOptim(list(self.module.parameters())) + + if self.zero_legacy_stage1(): + raise Exception( + "The deprecated version of ZeRO Stage 1 is not supported in deepspeed >= 0.5.9. Please downgrade to a version less than 0.5.9 if you need to use this deprecated version of ZeRO." + ) + + if zero_stage <= ZeroStageEnum.gradients: + overlap_comm = self.zero_overlap_comm() + contiguous_gradients = self.zero_contiguous_gradients() + round_robin_gradients = self.zero_round_robin_gradients() + assert not isinstance(optimizer, DummyOptim), "zero stage {} requires an optimizer".format(zero_stage) + + log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) + + if isinstance(self.module, PipelineModule): + if overlap_comm: + logger.warning("Pipeline parallelism does not support overlapped communication, will be disabled.") + overlap_comm = False + optimizer = DeepSpeedZeroOptimizer( + optimizer, + self.param_names, + timers=timers, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=self.dynamic_loss_scale_args(), + clip_grad=self.gradient_clipping(), + contiguous_gradients=contiguous_gradients, + reduce_bucket_size=self.zero_reduce_bucket_size(), + use_multi_rank_bucket_allreduce=self.zero_multi_rank_bucket_allreduce(), + allgather_bucket_size=self.zero_allgather_bucket_size(), + dp_process_group=self.seq_data_parallel_group, + expert_parallel_group=self.expert_parallel_group if self.has_moe_layers else None, + expert_data_parallel_group=self.expert_data_parallel_group if self.has_moe_layers else None, + reduce_scatter=self.zero_reduce_scatter(), + overlap_comm=overlap_comm, + offload_optimizer_config=self.zero_offload_optimizer(), + mpu=self.mpu, + postscale_gradients=self.postscale_gradients(), + gradient_predivide_factor=self.gradient_predivide_factor(), + gradient_accumulation_steps=self.gradient_accumulation_steps(), + ignore_unused_parameters=self.zero_ignore_unused_parameters(), + partition_grads=zero_stage == ZeroStageEnum.gradients, + round_robin_gradients=round_robin_gradients, + has_moe_layers=self.has_moe_layers, + fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(), + gradient_accumulation_dtype=gradient_accumulation_dtype, + communication_data_type=self.communication_data_type, + elastic_checkpoint=self.zero_elastic_checkpoint()) + + elif zero_stage == ZeroStageEnum.weights: + assert not self.has_moe_layers, "MoE not supported with Stage 3" + if isinstance(optimizer, DummyOptim): + log_dist("Creating ZeRO Offload", ranks=[0]) + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + if self.zero_hpz_partition_size() > 1 and zero_param_parallel_group is None: + self._set_zero_group_parallelism() + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + optimizer = DeepSpeedZeRoOffload( + self.module, + timers=timers, + ds_config=self.config, + overlap_comm=self.zero_overlap_comm(), + prefetch_bucket_size=self.zero_prefetch_bucket_size(), + max_reuse_distance=self.zero_max_reuse_distance(), + max_live_parameters=self.zero_max_live_parameters(), + param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), + offload_param_config=self.zero_offload_param(), + mpu=self.mpu, + zero_param_parallel_group=zero_param_parallel_group, + zero_quantized_weights=self.zero_quantized_weights(), + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), + zero_module_granularity_threshold=self.zero_module_granularity_threshold(), + log_trace_cache_warnings=self.zero_log_trace_cache_warnings(), + ) + else: + log_dist( + f'Creating fp16 ZeRO stage {zero_stage} optimizer,' + f' MiCS is enabled {mics_shard_size>0},' + f' Hierarchical params gather {self._config.mics_hierarchial_params_gather}', + ranks=[0]) + if mics_shard_size > 0: + return self._return_mics_optimizer(optimizer, timers) + + log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) + from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 + optimizer = DeepSpeedZeroOptimizer_Stage3( + self.module, + optimizer, + timers=timers, + ds_config=self.config, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=self.dynamic_loss_scale_args(), + clip_grad=self.gradient_clipping(), + contiguous_gradients=self.zero_contiguous_gradients(), + reduce_bucket_size=self.zero_reduce_bucket_size(), + prefetch_bucket_size=self.zero_prefetch_bucket_size(), + max_reuse_distance=self.zero_max_reuse_distance(), + max_live_parameters=self.zero_max_live_parameters(), + param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), + dp_process_group=self.seq_data_parallel_group, + all2all_process_group=self.local_all_to_all_group, + reduce_scatter=self.zero_reduce_scatter(), + overlap_comm=self.zero_overlap_comm(), + offload_optimizer_config=self.zero_offload_optimizer(), + offload_param_config=self.zero_offload_param(), + sub_group_size=self.zero_sub_group_size(), + offload_ratio=self.zero_partial_offload(), + mpu=self.mpu, + postscale_gradients=self.postscale_gradients(), + gradient_predivide_factor=self.gradient_predivide_factor(), + gradient_accumulation_steps=self.gradient_accumulation_steps(), + aio_config=self.aio_config(), + gradient_accumulation_dtype=gradient_accumulation_dtype, + communication_data_type=self.communication_data_type, + zero_hpz_partition_size=self.zero_hpz_partition_size(), + zero_quantized_weights=self.zero_quantized_weights(), + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), + zero_module_granularity_threshold=self.zero_module_granularity_threshold(), + zeropp_loco_param=self.zeropp_loco_param(), + log_trace_cache_warnings=self.zero_log_trace_cache_warnings(), + ) + + else: + raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage)) + + return optimizer + + def _return_mics_optimizer(self, basic_optimizer, timers): + from deepspeed.runtime.zero.mics import MiCS_Optimizer + model_dtype, gradient_accumulation_dtype = self.get_data_types() + optimizer = MiCS_Optimizer(self.module, + basic_optimizer, + timers=timers, + ds_config=self.config, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=self.dynamic_loss_scale_args(), + clip_grad=self.gradient_clipping(), + contiguous_gradients=self.zero_contiguous_gradients(), + reduce_bucket_size=self.zero_reduce_bucket_size(), + prefetch_bucket_size=self.zero_prefetch_bucket_size(), + max_reuse_distance=self.zero_max_reuse_distance(), + max_live_parameters=self.zero_max_live_parameters(), + param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), + dp_process_group=self.seq_data_parallel_group, + reduce_scatter=self.zero_reduce_scatter(), + overlap_comm=self.zero_overlap_comm(), + offload_optimizer_config=self.zero_offload_optimizer(), + offload_param_config=self.zero_offload_param(), + sub_group_size=self.zero_sub_group_size(), + mpu=self.mpu, + postscale_gradients=self.postscale_gradients(), + gradient_predivide_factor=self.gradient_predivide_factor(), + gradient_accumulation_steps=self.gradient_accumulation_steps(), + aio_config=self.aio_config(), + gradient_accumulation_dtype=gradient_accumulation_dtype, + communication_data_type=self.communication_data_type) + return optimizer + + def _configure_eigenvalue(self): + eigenvalue = Eigenvalue( + verbose=self.eigenvalue_verbose(), + max_iter=self.eigenvalue_max_iter(), + tol=self.eigenvalue_tol(), + stability=self.eigenvalue_stability(), + gas_boundary_resolution=self.eigenvalue_gas_boundary_resolution(), + layer_name=self.eigenvalue_layer_name(), + layer_num=self.eigenvalue_layer_num(), + ) + + return eigenvalue + + def _configure_progressive_layer_drop(self): + pld = ProgressiveLayerDrop(theta=self.pld_theta(), gamma=self.pld_gamma()) + + return pld + + def _configure_curriculum_scheduler_legacy(self): + scheduler = CurriculumScheduler(self.curriculum_params_legacy()) + return scheduler + + @staticmethod + def is_map_style_dataset(obj): + return hasattr(obj, "__getitem__") and hasattr(obj, "__len__") + + @staticmethod + def is_iterable_style_dataset(obj): + return isinstance(obj, torch.utils.data.IterableDataset) # hasattr(obj, "__iter__") should work as well + + def dataloader_drop_last(self): + return self._config.dataloader_drop_last + + def was_step_applied(self) -> bool: + """Returns True if the latest ``step()`` produced in parameter updates. + Note that a ``False`` return is not an error condition. Steps are frequently + no-ops, such as between gradient accumulation boundaries or when overflows + occur. + Returns: + bool: Whether the latest ``step()`` modified model parameters. + """ + return self._step_applied + + def deepspeed_io(self, + dataset, + batch_size=None, + route=ROUTE_TRAIN, + pin_memory=True, + data_sampler=None, + collate_fn=None, + num_local_io_workers=None): + if not (self.is_map_style_dataset(dataset) or self.is_iterable_style_dataset(dataset)): + raise ValueError("Training data must be a torch Dataset") + + if batch_size is None: + batch_size = self.train_micro_batch_size_per_gpu() + + if collate_fn is None: + collate_fn = self.collate_fn + + # Currently we only use timer in train route + deepspeed_io_timer = None + if route == ROUTE_TRAIN: + deepspeed_io_timer = self.tput_timer + + # If mpu is provided, forward world size and parallel rank to sampler. + data_parallel_world_size = self.dp_world_size + data_parallel_rank = self.global_rank + if self.mpu is not None: + data_parallel_world_size = self.mpu.get_data_parallel_world_size() + data_parallel_rank = self.mpu.get_data_parallel_rank() + + if data_sampler is None and (route == ROUTE_PREDICT or route == ROUTE_EVAL): + data_sampler = torch.utils.data.DistributedSampler( + dataset, + num_replicas=data_parallel_world_size, + rank=data_parallel_rank, + shuffle=False, + ) + + deepspeed_dataloader_config = {} + if self.curriculum_learning_enabled(): + deepspeed_dataloader_config = { + CURRICULUM_LEARNING: self.curriculum_learning_enabled(), + DATA_EFFICIENCY: self.data_efficiency_config(), + DATA_PARALLEL_GROUP: self.data_parallel_group, + GRADIENT_ACCUMULATION_STEPS: self.gradient_accumulation_steps(), + GLOBAL_RANK: self.global_rank, + DATA_SAMPLING_NUM_WORKERS: self.data_sampling_config()[DATA_SAMPLING_NUM_WORKERS] + } + return DeepSpeedDataLoader(dataset=dataset, + batch_size=batch_size, + pin_memory=pin_memory, + collate_fn=collate_fn, + local_rank=self.local_rank, + tput_timer=deepspeed_io_timer, + num_local_io_workers=num_local_io_workers, + data_sampler=data_sampler, + data_parallel_world_size=data_parallel_world_size, + data_parallel_rank=data_parallel_rank, + dataloader_drop_last=self.dataloader_drop_last(), + deepspeed_dataloader_config=deepspeed_dataloader_config) + + def train(self, mode=True): + r"""""" + + self.warn_unscaled_loss = True + self.module.train(mode) + + def eval(self): + r"""""" + + self.warn_unscaled_loss = True + self.module.train(False) + + def _scale_loss_by_gas(self, prescaled_loss, eval_micro_batches=None): + # In pipeline evaluation, there is an option to use different micro-bs, which creates different number of + # micro batches, thus the training gas, is not valid in this case. need to use the number of eval_micro_batches + scaling_factor = self.gradient_accumulation_steps() if eval_micro_batches is None else eval_micro_batches + if isinstance(prescaled_loss, torch.Tensor): + scaled_loss = prescaled_loss / scaling_factor + elif isinstance(prescaled_loss, tuple) or isinstance(prescaled_loss, list): + scaled_loss = [] + for l in prescaled_loss: + if isinstance(l, torch.Tensor): + scaled_loss.append(l / scaling_factor) + else: + scaled_loss.append(l) + else: + scaled_loss = prescaled_loss + if self.warn_unscaled_loss: + logger.warning(f"DeepSpeed unable to scale loss because of type: {type(prescaled_loss)}") + self.warn_unscaled_loss = False + + return scaled_loss + + def _create_module_forward_pre_hook(self): + + def _module_forward_pre_hook(module, inputs, kwargs): + return self._forward_prologue(inputs, kwargs) + + return self.module.register_forward_pre_hook(_module_forward_pre_hook, prepend=False, with_kwargs=True) + + def _create_module_forward_post_hook(self): + + def _module_forward_post_hook(module, input, output): + self._forward_epilogue() + + return self.module.register_forward_hook(_module_forward_post_hook) + + def _forward_prologue(self, inputs, kwargs): + return_modified = False + + if not self.autotuning_profile_model_info(): + see_memory_usage("Engine before forward", force=self.memory_breakdown()) + + flops_profiler_active = (self.flops_profiler_enabled() + and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0) + + # used to check quantization happens at step 0! + if self.global_steps == 0 and hasattr(self, "compression_scheduler"): + self.compression_scheduler.step(step_zero_check=True) + if self.quantizer: + tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( + ) == 2 else self.optimizer.fp16_groups + if self.compression_scheduler.weight_quantization_enabled: + self.quantizer.quantize( + tensor_to_quantize, + (self.optimizer.overflow if self.fp16_enabled() else False), + self.eigenvalue_enabled(), + None, + ) + return_modified = True + + if flops_profiler_active: + self.flops_profiler.start_profile(ignore_list=None) + + if kwargs is not None: + if self.module.training: + if self.progressive_layer_drop: + kwargs.update(self.progressive_layer_drop.get_state()) + + if self.__class__.__name__ != "PipelineEngine": + # TODO: The above if condition is a HACK since for PipelineEngine + # it's difficult to inject argument in forward pass. + if self.module.training and self.curriculum_enabled_legacy(): + self.curriculum_scheduler_legacy.update_difficulty(self.global_steps + 1) + if self.curriculum_params_legacy()["curriculum_type"] == "seqlen": + kwargs.update({"curriculum_seqlen": self.curriculum_scheduler_legacy.get_current_difficulty()}) + return_modified = True + + if self.module.training and self.random_ltd_enabled(): + self.random_ltd_scheduler.update_seq(self.global_steps) + + if self.training_dataloader is None: + self.tput_timer.start() + + self._start_timers(self.engine_timers.forward_timers) + + if self.zero_optimization_partition_weights(): + # Enable automated discovery of external parameters by indicating that + # we are in a forward pass. + for module in self.module.modules(): + module._parameters._in_forward = True + + if self.fp16_auto_cast(): + inputs = self._cast_inputs_half(inputs) + return_modified = True + + if return_modified: + return inputs, kwargs + + def _forward_epilogue(self): + if self.zero_optimization_partition_weights(): + # Disable automated discovery of external parameters + for module in self.module.modules(): + module._parameters._in_forward = False + + self._stop_timers(self.engine_timers.forward_timers) + + flops_profiler_active = (self.flops_profiler_enabled() + and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0) + + if flops_profiler_active: + self.flops_profiler.stop_profile() + + if not self.autotuning_profile_model_info(): + see_memory_usage("Engine after forward", force=self.memory_breakdown()) + + @instrument_w_nvtx + def forward(self, *inputs, **kwargs): + r"""Execute forward propagation + Arguments: + *inputs: Variable length input list + **kwargs: variable length keyword arguments + """ + if self.autotuning_profile_model_info(): + ma = get_ma_status() + + if self.is_deepcompile_enabled() and hasattr(self, "launch_compile_passes"): + # We can't have this in forward prologue as the compiler compiles hooks including the forward prologue. + self.launch_compile_passes(self.global_steps) + + loss = self.module(*inputs, **kwargs) + + if self.autotuning_profile_model_info(): + activation_mem = get_ma_status() - ma + self.autotuning_model_info["activation_mem_per_gpu"] = activation_mem + print_json_dist(self.autotuning_model_info, [0], path=self.autotuning_model_info_path()) + exit() + + return loss + + def _cast_inputs_half(self, inputs): + if isinstance(inputs, (list, tuple)): + new_inputs = [] + for v in inputs: + new_inputs.append(self._cast_inputs_half(v)) + return inputs.__class__(new_inputs) + elif isinstance(inputs, dict): + new_inputs = {} + for k, v in inputs.items(): + new_inputs[k] = self._cast_inputs_half(v) + return new_inputs + elif hasattr(inputs, 'half') and inputs.is_floating_point(): + return inputs.half() + else: + return inputs + + def print_forward_breakdown(self, fwd_time): + gate_time = 0.0 + moe_time = 0.0 + falltoall = 0.0 + salltoall = 0.0 + + for gate in self.gate_modules: + #logger.info(f"Individual TopK gate time: {gate.gate_time:.2f} ms") + gate_time += gate.gate_time + + for l in self.moe_layers: + #logger.info(f"MoE layer; total: {l.time_moe:.2f} ms, first alltoall: {l.time_falltoall:.2f}, second alltoall: {l.time_salltoall:.2f}") + moe_time += l.time_moe + falltoall += l.time_falltoall + salltoall += l.time_salltoall + + # TODO: Allreduce/average them across ranks for more accurate timing. + + # if deepspeed.comm.get_rank() == 0: + log_dist( + f"time (ms) | fwd: {fwd_time:.2f} (fwd_moe: {moe_time:.2f}, 1st_a2a: {falltoall:.2f}, 2nd_a2a: {salltoall:.2f}, top_k: {gate_time:.2f})", + ranks=[0]) + + @instrument_w_nvtx + def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE): + # Pass (PP) gas boundary flag to optimizer (required for zero) + self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() + # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well + if self.zero_optimization_partition_gradients(): + self.optimizer.overlapping_partition_gradients_reduce_epilogue() + + # Communicate only at gradient accumulation boundaries + elif self.is_gradient_accumulation_boundary(): + if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states and hasattr( + self.optimizer, 'reduce_gradients'): + self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism) + else: + grads = None + self.buffered_allreduce_fallback(grads=grads, elements_per_buffer=bucket_size) + + def _backward_prologue(self, loss, scale_wrt_gas=True): + see_memory_usage("Engine before backward", force=self.memory_breakdown()) + if self.scale_wrt_gas is not None: + scale_wrt_gas = self.scale_wrt_gas + + # scale loss w.r.t. gradient accumulation if reduction is not disabled + do_gradient_reduction = self.enable_backward_allreduce and not self.inside_no_sync_ctxt and not self.is_deepcompile_enabled( + ) + if do_gradient_reduction and self.gradient_accumulation_steps() > 1 and scale_wrt_gas: + loss = self._scale_loss_by_gas(loss.float()) + + # Log training loss + mean_loss = loss.mean().detach() + self.losses = mean_loss if self.losses is None else self.losses + mean_loss + if self.monitor.enabled: + if self.is_gradient_accumulation_boundary(): + if self.global_rank == 0: + self.summary_events = [( + f"Train/Samples/train_loss", + self.losses.item(), + self.global_samples, + )] + self.monitor.write_events(self.summary_events) + + if self.is_deepcompile_enabled(): + deepcompile_backward_prologue(self.is_gradient_accumulation_boundary()) + + return loss + + def _backward_epilogue(self): + self._start_timers(self.engine_timers.backward_reduce_timers) + if self.enable_backward_allreduce and not self.inside_no_sync_ctxt: + # Traditional code path that allreduces the module parameter grads + self.allreduce_gradients() + + self._stop_timers(self.engine_timers.backward_reduce_timers) + see_memory_usage("Engine after backward", force=self.memory_breakdown()) + + def _do_optimizer_backward(self, loss, retain_graph): + self._start_timers(self.engine_timers.backward_inner_timers) + if self.zero_optimization(): + self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() + self.optimizer.backward(loss, retain_graph=retain_graph) + elif self.amp_enabled(): + # AMP requires delaying unscale when inside gradient accumulation boundaries + # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations + delay_unscale = not self.is_gradient_accumulation_boundary() + with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss: + scaled_loss.backward(retain_graph=retain_graph) + elif self.fp16_enabled(): + if self.eigenvalue_enabled(): + self.optimizer.backward(loss, create_graph=True, retain_graph=True) + else: + self.optimizer.backward(loss, retain_graph=retain_graph) + elif self.bfloat16_enabled(): + self.optimizer.backward(loss, retain_graph=retain_graph) + else: + if self.eigenvalue_enabled(): + loss.backward(create_graph=True, retain_graph=True) + else: + loss.backward(retain_graph=retain_graph) + self._stop_timers(self.engine_timers.backward_inner_timers) + + @contextmanager + def no_sync(self): + r""" + Context manager to disable gradient reduction during backward pass. + This context manager has the following effects on other DeepSpeed features: + 1. Incompatible with ZeRO stage 2/3 which rely on reduction for gradient partitioning. + 2. It is illegal to call engine.step() within the context manager. + 3. Tracking of gradient accumulation steps is disabled. + """ + assert not self.zero_optimization_partition_gradients(), \ + f"no_sync context manager is incompatible with gradient partitioning logic of ZeRO stage {self.zero_optimization_stage()}" + + assert not self.inside_no_sync_ctxt, f"no_sync context manager reentry is unsupported" + + self.inside_no_sync_ctxt = True + try: + yield + finally: + self.inside_no_sync_ctxt = False + + @instrument_w_nvtx + def backward(self, loss, retain_graph=False, scale_wrt_gas=True): + r"""Execute backward pass on the loss + Arguments: + loss: Torch tensor on which to execute backward propagation + retain_graph: bool, default: false + forward on user defined choice of retain_graph + """ + assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ + "must provide optimizer during init in order to use backward" + + self._start_timers(self.engine_timers.backward_timers) + loss = self._backward_prologue(loss, scale_wrt_gas) + self._do_optimizer_backward(loss, retain_graph) + self._backward_epilogue() + self._stop_timers(self.engine_timers.backward_timers) + + return loss + + def is_gradient_accumulation_boundary(self): + """ + Query whether the current micro-batch is at the boundary of + gradient accumulation, and thus will trigger gradient reductions and + an optimizer step. + + Returns: + bool: if the current step is a gradient accumulation boundary. + + """ + if self._is_gradient_accumulation_boundary is None: + return (self.micro_steps + 1) % \ + self.gradient_accumulation_steps() == 0 + else: + return self._is_gradient_accumulation_boundary + + def set_gradient_accumulation_boundary(self, is_boundary): + """ + Manually overrides the DeepSpeed engine's gradient accumulation boundary state, this is an optional + feature and should be used with care. The state should be set before to the intended + value before each forward/backward. The final forward/backward should have the + boundary state set to True. This style allows client code to only call engine.step() once after all + the gradient accumulation passes are complete. See example below: + .. code-block:: python + engine.set_gradient_accumulation_boundary(False) + for _ in range(gradient_accumulation_steps - 1): + micro_batch = next(data_loader) + loss = engine(micro_batch) + engine.backward(loss) + engine.set_gradient_accumulation_boundary(True) + micro_batch = next(data_loader) + loss = engine(micro_batch) + engine.backward(loss) + engine.step() + Arguments: + is_boundary (bool): are we at a gradient accumulation boundary or not? + """ + self._is_gradient_accumulation_boundary = is_boundary + self.optimizer.is_gradient_accumulation_boundary = is_boundary + + def zero_grad(self): + """ + Zero parameter grads. + """ + for param_name, param in self.module.named_parameters(): + param.grad = None + + def clip_fp32_gradients(self): + clip_grad_norm_(parameters=self.module.parameters(), max_norm=self.gradient_clipping(), mpu=self.mpu) + + def _take_model_step(self, lr_kwargs, block_eigenvalue={}): + if self.gradient_clipping() > 0.0: + if not (self.fp16_enabled() or self.bfloat16_enabled() or self.amp_enabled() or self.zero_optimization()): + self.clip_fp32_gradients() + elif self.amp_enabled(): + # AMP's recommended way of doing clipping + # https://nvidia.github.io/apex/advanced.html#gradient-clipping + master_params = amp.master_params(self.optimizer) + clip_grad_norm_(parameters=master_params, max_norm=self.gradient_clipping(), mpu=self.mpu) + self.optimizer.step() + + if hasattr(self.optimizer, '_global_grad_norm'): + self._global_grad_norm = self.optimizer._global_grad_norm + + # Quantize the updated parameter if there is no overflow + if self.quantizer: + tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( + ) == 2 else self.optimizer.fp16_groups + if self.compression_scheduler.weight_quantization_enabled: + self.quantizer.quantize( + tensor_to_quantize, + (self.optimizer.overflow if self.fp16_enabled() else False), + self.eigenvalue_enabled(), + block_eigenvalue, + ) + # zero grad in basic optimizer could be unreliable and may not exhibit + # the behavior that we want + if self.bfloat16_enabled(): + # TODO: Temporary until bf16_optimizer and zero_optimizer are integrated + if self.zero_optimization() and hasattr(self.optimizer, "zero_grad"): + self.optimizer.zero_grad() + else: + pass + elif self.zero_optimization() or self.fp16_enabled() or self.amp_enabled(): + self.optimizer.zero_grad() + else: + self.zero_grad() + + # Check overflow here since in DS fp16 optimizer, the overflow is updated in above step() function. + overflow = False + if hasattr(self.optimizer, "overflow"): + overflow = self.optimizer.overflow + self._step_applied = not overflow + + if overflow: + self.skipped_steps += 1 + else: + self.compression_scheduler.step() + if self.lr_scheduler is not None: + try: + self.lr_scheduler.step(**(lr_kwargs or {})) + except TypeError: + # XXX Hack to work with Megatron 2.0 and DeepSpeed pipelines. + # We don't currently have a way to specify lr_kwargs from + # pipe_engine.train_batch() + self.lr_scheduler.step(self.train_batch_size()) + + if self.steps_per_print() is not None: + report_progress = self.global_rank == 0 if self.global_rank else True + if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0: + self._report_progress(self.global_steps + 1) + + self.losses = None + self.global_steps += 1 + self.global_samples += self.train_batch_size() + + def step(self, lr_kwargs=None): + r"""Execute the weight update step after forward and backward propagation + on effective_train_batch. + """ + assert not self.inside_no_sync_ctxt, \ + "It is illegal to call Engine.step() inside no_sync context manager" + + see_memory_usage("Engine before step", force=self.memory_breakdown()) + + # Check early because self.global_steps is incremented at some point here. + # TODO: Delay self.global_steps increment until very end of this function. + flops_profiler_active = self.flops_profiler_enabled( + ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 + + self._start_timers(self.engine_timers.step_timers) + + assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ + "must provide optimizer during init in order to use step" + + report_progress = False + + self._step_applied = False # assume False, will flip to True + + # Update the model when we reach gradient accumulation boundaries + if self.is_gradient_accumulation_boundary(): + self.gas_boundary_ctr += 1 + + if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) + and self.quantizer.any_precision_switch()): + log_dist(f"computing eigenvalue...", ranks=[0]) + self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, + self.optimizer.cur_scale) + + if self.progressive_layer_drop: + self.progressive_layer_drop.update_state(self.global_steps) + + if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() + and self.quantizer.any_precision_switch()): + self._take_model_step(lr_kwargs, self.block_eigenvalue) + else: + self._take_model_step(lr_kwargs) + + report_progress = self.global_rank == 0 if self.global_rank else True + + self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress) + + self._stop_timers(self.engine_timers.step_timers) + + # Log learning rate + if self.monitor.enabled: + if self.is_gradient_accumulation_boundary(): + if self.global_rank == 0: + self.summary_events = [(f"Train/Samples/lr", self.get_lr()[0], self.global_samples)] + + if self.fp16_enabled() and hasattr(self.optimizer, "cur_scale"): + self.summary_events.append(( + f"Train/Samples/loss_scale", + self.optimizer.cur_scale, + self.global_samples, + )) + + if (self.eigenvalue_enabled() + and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution()): + ev_values = self.block_eigenvalue.values() + for i in range(len(ev_values)): + self.summary_events.append(( + f"Train/Eigenvalues/ModelBlockParam_{i}", + self.ev_values[i][0], + self.global_samples, + )) + self.monitor.write_events(self.summary_events) + + # Check flops profiling + if flops_profiler_active: + if self.autotuning_enabled(): + self.flops = self.flops_profiler.get_total_flops() * 3 + self.fwd_duration = self.flops_profiler.get_total_duration() + else: + self.flops_profiler.print_model_profile( + profile_step=self.global_steps, + module_depth=self.flops_profiler_module_depth(), + top_modules=self.flops_profiler_top_modules(), + detailed=self.flops_profiler_detailed(), + output_file=self.flops_profiler_output_file(), + ) + self.flops_profiler.end_profile() + + if self.autotuning_enabled() and self.global_steps == (self.autotuning_end_profile_step() + 1): + self._autotuning_exit() + + if self.wall_clock_breakdown(): + # Log micro timing and reset + self.timers.log(names=self.engine_timers.micro_timers, memory_breakdown=self.memory_breakdown()) + + if self.wall_clock_breakdown() or self.flops_profiler_enabled(): + # Log global timing and reset + if self.is_gradient_accumulation_boundary(): + if self.monitor.enabled: + self._write_monitor() + + if self.has_moe_layers: + fwd_time = self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False) + self.print_forward_breakdown(fwd_time=fwd_time) + + self.timers.log(self.engine_timers.global_timers) + + self.micro_steps += 1 + see_memory_usage("Engine after step", force=self.memory_breakdown()) + + def _start_timers(self, timer_names): + for name in timer_names: + self.timers(name).start() + + def _stop_timers(self, timer_names): + record = self.is_gradient_accumulation_boundary() and \ + self.flops_profiler_enabled() and \ + (self.global_steps >= self.flops_profiler_profile_step()) + for name in timer_names: + self.timers(name).stop(record=record) + + def _autotuning_exit(self): + if self.global_rank == 0: + msg = self.timers.get_mean([ + FORWARD_GLOBAL_TIMER, + BACKWARD_GLOBAL_TIMER, + STEP_GLOBAL_TIMER, + ], reset=False) + titer = 0.0 + titer += msg[FORWARD_GLOBAL_TIMER] if FORWARD_GLOBAL_TIMER in msg else 0 + titer += msg[BACKWARD_GLOBAL_TIMER] if BACKWARD_GLOBAL_TIMER in msg else 0 + titer += msg[STEP_GLOBAL_TIMER] if STEP_GLOBAL_TIMER in msg else 0 + titer *= self.gradient_accumulation_steps() + msg["latency"] = titer + msg["FLOPS_per_gpu"] = self.flops * 1_000_000 * self.gradient_accumulation_steps() / titer + msg["throughput"] = self.train_batch_size() * 1_000_000 / \ + msg["latency"] + print_json_dist(msg, [0], path=self.autotuning_metric_path()) + log_dist( + f"Wrote metrics to {self.autotuning_metric_path()}, {os.path.abspath(self.autotuning_metric_path())}", + ranks=[0]) + import atexit + atexit.register(print, "Autotuning: done with running current ds config.") + exit() + + def _write_monitor(self): + if self.global_rank == 0: + self.summary_events = [ + ( + f"Train/Samples/elapsed_time_ms_forward", + self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_backward", + self.timers(BACKWARD_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_backward_inner", + self.timers(BACKWARD_INNER_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_backward_allreduce", + self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ( + f"Train/Samples/elapsed_time_ms_step", + self.timers(STEP_GLOBAL_TIMER).elapsed(reset=False), + self.global_samples, + ), + ] + self.monitor.write_events(self.summary_events) + + def _get_optimizer_param(self, param_name): + result = [] + if not self.optimizer: + return result + for group in self.optimizer.param_groups: + if param_name in group: + result.append(group[param_name]) + else: + result.append(0.0) + return result + + def get_lr(self): + return self._get_optimizer_param("lr") + + def get_type(self): + return self._get_optimizer_param("type") + + def get_mom(self): + if self.optimizer_name() in ["SGD", "RMSprop"]: + return self._get_optimizer_param("momentum") + else: + return self._get_optimizer_param("betas") + + def get_pld_theta(self): + if self.progressive_layer_drop: + return self.progressive_layer_drop.get_theta() + else: + return None + + def _report_progress(self, step): + lr = self.get_lr() + mom = self.get_mom() + log_dist(f"step={step}, skipped={self.skipped_steps}, lr={lr}, mom={mom}", ranks=[0]) + + def allreduce_bucket(self, bucket, dp_group, dp_world_size=None): + tensor = self.flatten(bucket) + + tensor_to_allreduce = tensor + + if self.communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(self.communication_data_type) + + if dp_world_size is None: + dp_world_size = dist.get_world_size(group=dp_group) + if self.postscale_gradients(): + if self.gradient_predivide_factor() != 1.0: + tensor_to_allreduce.mul_(1.0 / self.gradient_predivide_factor()) + + dist.all_reduce(tensor_to_allreduce, group=dp_group) + if self.gradient_average: + if self.gradient_predivide_factor() != dp_world_size: + tensor_to_allreduce.mul_(self.gradient_predivide_factor() / dp_world_size) + else: + tensor_to_allreduce.mul_(1. / dp_world_size) + dist.all_reduce(tensor_to_allreduce, group=dp_group) + + if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + tensor.copy_(tensor_to_allreduce) + + return tensor + + def allreduce_and_copy(self, small_bucket, dp_group, dp_world_size=None): + allreduced = self.allreduce_bucket(small_bucket, dp_group, dp_world_size) + for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): + buf.copy_(synced) + + def allreduce_no_retain(self, bucket, dp_group, numel_per_bucket=500000000, dp_world_size=None): + small_bucket = [] + numel = 0 + for tensor in bucket: + small_bucket.append(tensor) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy(small_bucket, dp_group, dp_world_size) + small_bucket = [] + numel = 0 + if len(small_bucket) > 0: + self.allreduce_and_copy(small_bucket, dp_group, dp_world_size) + + def _get_gradients_for_reduction(self): + non_expert_grads = [] + expert_grads = {} + if self.has_moe_layers: + for key in self.expert_data_parallel_group.keys(): + expert_grads[key] = [] + + for param_name, param in self.module.named_parameters(): + if not param.requires_grad: + continue + + if param.grad is None: + # In cases where there is an imbalance of empty grads across + # ranks we must create empty grads, this will ensure that every + # rank is reducing the same size. In some cases it may make + # sense in the future to support the ability to average not + # w.r.t. world size but with a different value. + param.grad = torch.zeros(param.size(), dtype=param.dtype, device=param.device) + + grad_data = param.grad.data + if param_name in self.sparse_tensor_module_names or grad_data.is_sparse: + # Call param.grad without data to avoid problem with setting of updated grads + grad_data = SparseTensor(param.grad) + + if is_moe_param(param): + expert_grads[param.group_name].append(grad_data) + else: + non_expert_grads.append(grad_data) + + return non_expert_grads, expert_grads + + def _reduce_non_expert_gradients(self, grads, elements_per_buffer): + split_sparse_tensor_buckets, split_dense_tensor_buckets = split_half_float_double_sparse(grads) + if self.pipeline_parallelism: + dp_group = self.mpu.get_data_parallel_group() + dp_world_size = dist.get_world_size(dp_group) + else: + dp_group = groups._get_sequence_data_parallel_group() + dp_world_size = dist.get_world_size(dp_group) / float(self.sequence_parallel_size) + for _, sparse_bucket_tuple in enumerate(split_sparse_tensor_buckets): + if sparse_bucket_tuple: + bucket_type, sparse_bucket = sparse_bucket_tuple + self.sparse_allreduce_no_retain(sparse_bucket, dp_group=dp_group, dp_world_size=dp_world_size) + + for _, dense_bucket_tuple in enumerate(split_dense_tensor_buckets): + if dense_bucket_tuple: + bucket_type, dense_bucket = dense_bucket_tuple + self.allreduce_no_retain(dense_bucket, + dp_group=dp_group, + numel_per_bucket=elements_per_buffer, + dp_world_size=dp_world_size) + + def _reduce_expert_gradients(self, expert_grads, elements_per_buffer): + # to maintain the gradients value unaffected by ep_size setting, + # utilize dp_world_size for allreduce average + dp_world_size = dist.get_world_size(groups._get_data_parallel_group()) + for ep_name, expert_grads_group in expert_grads.items(): + ep_dp_group = groups._get_expert_data_parallel_group(ep_name) + split_sparse_tensor_buckets, split_dense_tensor_buckets = split_half_float_double_sparse( + expert_grads_group) + + for _, sparse_bucket_tuple in enumerate(split_sparse_tensor_buckets): + if sparse_bucket_tuple: + bucket_type, sparse_bucket = sparse_bucket_tuple + self.sparse_allreduce_no_retain(sparse_bucket, dp_group=ep_dp_group, dp_world_size=dp_world_size) + + for _, dense_bucket_tuple in enumerate(split_dense_tensor_buckets): + if dense_bucket_tuple: + bucket_type, dense_bucket = dense_bucket_tuple + # Separate between diff groups + self.allreduce_no_retain(dense_bucket, + dp_group=ep_dp_group, + numel_per_bucket=elements_per_buffer, + dp_world_size=dp_world_size) + + def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000): + if grads is None: + if hasattr(self.optimizer, "get_grads_for_reduction"): + # This is currently for BF16 optimizer + non_expert_grads, expert_grads = self.optimizer.get_grads_for_reduction() + else: + non_expert_grads, expert_grads = self._get_gradients_for_reduction() + else: + assert not self.has_moe_layers, "attempting to reduce grads in unsupported way w.r.t. MoE" + non_expert_grads = grads + + self._reduce_non_expert_gradients(non_expert_grads, elements_per_buffer) + + if self.has_moe_layers: + self._reduce_expert_gradients(expert_grads, elements_per_buffer) + + def sparse_allreduce_no_retain(self, bucket, dp_group, dp_world_size=None): + allreduced_sparses = self.sparse_allreduce_bucket(bucket, dp_group, dp_world_size) + # Densify sparse tensor and copy back to original location + for tensor in allreduced_sparses: + if tensor.is_sparse: + tensor.orig_dense_tensor.data = tensor.to_coo_tensor() + else: + tensor.orig_dense_tensor.copy_(tensor.to_dense()) + + def sparse_allreduce_bucket(self, bucket, dp_group, dp_world_size=None): + sparse_list = [] + for sparse in bucket: + sparse_list.append(self.sparse_allreduce(sparse, dp_group, dp_world_size)) + return sparse_list + + def sparse_allreduce(self, sparse, dp_group, dp_world_size=None): + original_data_type = sparse.values.dtype + if self.communication_data_type != sparse.values.dtype: + if self.communication_data_type in (torch.float16, torch.bfloat16): + indices = sparse.indices.to(torch.int32) + else: + indices = sparse.indices + values = sparse.values.to(self.communication_data_type) + else: + indices = sparse.indices + values = sparse.values + + if dp_world_size is None: + dp_world_size = dist.get_world_size(group=dp_group) + if self.postscale_gradients(): + if self.gradient_average: + values.mul_(self.gradient_predivide_factor() / (dp_world_size)) + else: + values.mul_(1. / (dp_world_size)) + + indices_device_list = self.sparse_all_gather(indices, dp_group) + values_device_list = self.sparse_all_gather(values, dp_group) + + sparse.indices = torch.cat(indices_device_list).to(torch.long) + sparse.values = torch.cat(values_device_list).to(original_data_type) + return sparse + + def sparse_all_gather(self, value, dp_group): + my_size = torch.LongTensor([value.size()[0]]).to(self.device) + all_sizes = self.all_gather_scalar(my_size, dp_group) + max_size = torch.cat(all_sizes).max() + fill_size = max_size - my_size + + assert value.dim() in [1, 2] + if value.dim() == 1: + if fill_size > 0: + value = torch.cat([value, value.new_empty(fill_size)]) + tensor_list = [value.new_empty(max_size) for _ in range(dist.get_world_size(group=dp_group))] + else: + if fill_size > 0: + value = torch.cat([value, value.new_empty(fill_size, value.size()[1])]) + tensor_list = [ + value.new_empty(max_size, + value.size()[1]) for _ in range(dist.get_world_size(group=dp_group)) + ] + + dist.all_gather(tensor_list, value, group=dp_group) + tensors = [] + for dev_idx, t in enumerate(tensor_list): + size = all_sizes[dev_idx][0] + tensors.append(t.index_select(0, torch.arange(size, dtype=torch.long, device=self.device))) + + return tensors + + def all_gather_scalar(self, value, dp_group): + tensor_list = [value.new_zeros(value.size()) for _ in range(dist.get_world_size(group=dp_group))] + dist.all_gather(tensor_list, value, group=dp_group) + return tensor_list + + def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclude_frozen_parameters=False): + sd = self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) + + # Remove frozen parameter weights from state_dict if specified + if exclude_frozen_parameters: + for n, p in self.module.named_parameters(): + if not p.requires_grad and n in sd: + del sd[n] + + if self.random_ltd_enabled(): + sd = remove_random_ltd_state_dict(sd) + return sd + + @staticmethod + def load_moe_state_dict(checkpoint_path, + tag, + state_dict, + old_moe_load, + model=None, + mpu=None, + num_experts=1, + checkpoint_engine=TorchCheckpointEngine()): + if old_moe_load: + expp_rank = groups._get_expert_data_parallel_rank(groups._get_max_expert_size_name()) + + num_local_experts = max(num_experts) // groups._get_expert_parallel_world_size( + groups._get_max_expert_size_name()) + for local_expert_id in range(num_local_experts): + global_expert_id = expp_rank * num_local_experts + local_expert_id + expert_state_dict = checkpoint_engine.load( + DeepSpeedEngine._get_expert_ckpt_name( + checkpoint_path, + -1, # -1 means ignore layer_id + global_expert_id, + tag, + mpu), + map_location=torch.device('cpu')) + + # Updating global -> local expert ids + moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' + for key in list(expert_state_dict.keys()): + local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', + f'{moe_str_prefix}{local_expert_id}') + expert_state_dict[local_key] = expert_state_dict.pop(key) + state_dict.update(expert_state_dict) + + else: + moe_layer_id = 0 + for n_module, module in model.named_modules(): + if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: + group_name = module.expert_group_name + num_local_experts = module.num_local_experts + expp_rank = groups._get_expert_parallel_rank(group_name) + # loop all local_experts + for local_expert_id in range(num_local_experts): + global_expert_id = expp_rank * num_local_experts + local_expert_id + expert_state_dict = checkpoint_engine.load(DeepSpeedEngine._get_expert_ckpt_name( + checkpoint_path, moe_layer_id, global_expert_id, tag, mpu), + map_location=torch.device('cpu')) + # print(expert_state_dict.keys()) + # Updating global -> local expert ids + moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' + for key in list(expert_state_dict.keys()): + local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', + f'{moe_str_prefix}{local_expert_id}') + expert_state_dict[local_key] = expert_state_dict.pop(key) + state_dict.update(expert_state_dict) + moe_layer_id += 1 + + def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False): + if fetch_z3_params: + params_to_fetch = [ + p for p in self.module.parameters() + if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE + ] + else: + params_to_fetch = [] + + with deepspeed.zero.GatheredParameters(params_to_fetch, modifier_rank=0): + module_state_dict = checkpoint['module'] + if custom_load_fn: + custom_load_fn(src=module_state_dict, dst=self.module) + else: + self.module.load_state_dict( + module_state_dict, # TODO + strict=strict) + + if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None: + saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS] + for param in self.module.parameters(): + if param.requires_grad: + continue + if param not in self.param_names: + raise ValueError(f"failed to find frozen {param} in named params") + name = self.param_names[param] + if hasattr(param, 'ds_id'): + param.ds_tensor.data.copy_(saved_frozen_params[name].data) + else: + param.data.copy_(saved_frozen_params[name].data) + + def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode): + return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}' + + def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16_mode): + file_prefix = self._get_zero_ckpt_prefix(dp_rank, bf16_mode=bf16_mode) + zero_ckpt_name = os.path.join( + checkpoints_path, + str(tag), + f"{file_prefix}_mp_rank_{mp_rank:02d}_optim_states.pt", + ) + return zero_ckpt_name + + def _get_zero_ckpt_name(self, checkpoints_path, tag): + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + pp_rank = dist.get_rank(group=self.optimizer.dp_process_group) + bf16_mode = self.bfloat16_enabled() + return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode) + + def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None): + if mp_placeholder is not None: + mp_rank_str = mp_placeholder + else: + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + mp_rank_str = f"{mp_rank:02d}" + + if self.zero_optimization_partition_weights(): + if self.load_universal_checkpoint(): + filename = "zero_pp_rank_0" + else: + filename = "zero_pp_rank_{}".format(dist.get_rank(group=self.optimizer.dp_process_group)) + ckpt_name = os.path.join( + checkpoints_path, + str(tag), + f"{filename}_mp_rank_{mp_rank_str}_model_states.pt", + ) + else: + ckpt_name = os.path.join( + checkpoints_path, + str(tag), + "mp_rank_" + mp_rank_str + "_model_states.pt", + ) + return ckpt_name + + def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank): + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + ckpt_name = os.path.join(checkpoints_path, str(tag), + f'expp_rank_{expp_rank}_mp_rank_{mp_rank:02d}_optim_states.pt') + return ckpt_name + + @staticmethod + def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None): + mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank() + if layer_id <= -1: + # Used to support old checkpoint loading + ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), + f'expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') + else: + # Used to support new checkpoint loading + ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), + f'layer_{layer_id}_expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') + return ckpt_name + + def _get_all_ckpt_names(self, checkpoints_path, tag): + # It is required that (checkpoints_path, tag) are consistent among all ranks. + ckpt_file_pattern = self._get_ckpt_name(checkpoints_path, tag, mp_placeholder="*") + import glob + + ckpt_files = glob.glob(ckpt_file_pattern) + ckpt_files.sort() + return ckpt_files + + def load_checkpoint(self, + load_dir, + tag=None, + load_module_strict=True, + load_optimizer_states=True, + load_lr_scheduler_states=True, + load_module_only=False, + custom_load_fn=None): + """ + Load training checkpoint + + Arguments: + load_dir: Required. Directory to load the checkpoint from + tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file + load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match. + load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance + load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. + load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting. + custom_load_fn: Optional. Custom model load function. + + Returns: + A tuple of ``load_path`` and ``client_state``. + *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed. + *``client_state``: State dictionary used for loading required training states in the client code. + + Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right + after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and + ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine + before ``load_checkpoint()``. + + """ + + if tag is None: + latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest" + latest_path = os.path.join(load_dir, latest_tag) + if os.path.isfile(latest_path): + with open(latest_path, "r") as fd: + tag = fd.read().strip() + else: + if self.load_universal_checkpoint(): + raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist') + else: + logger.warning( + f"Unable to find latest file at {latest_path}, if trying to load latest " + "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint." + ) + return None, None + + if self._optimizer_has_ckpt_event_prologue(): + # Prepare for checkpoint load by ensuring all parameters are partitioned + self.optimizer.checkpoint_event_prologue() + + load_path, client_states = self._load_checkpoint(load_dir, + tag, + load_module_strict=load_module_strict, + load_optimizer_states=load_optimizer_states, + load_lr_scheduler_states=load_lr_scheduler_states, + load_module_only=load_module_only, + custom_load_fn=custom_load_fn) + + load_zero_checkpoint = load_path is not None and (self.zero_optimization() or self.bfloat16_enabled()) + if load_zero_checkpoint: + if (load_optimizer_states and not load_module_only) or self.load_universal_checkpoint(): + success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states) + else: + success = False + if not success: + self.optimizer._restore_from_bit16_weights() + + if self.zero_nvme_offload_optimizer(): + from shutil import copytree, disk_usage + offload_dir = self.optimizer.optimizer_swapper.swap_folder + offload_ckpt_dir = os.path.join(load_dir, tag, "offloaded_tensors") + _, _, free = disk_usage(offload_dir) + logger.info( + f"Copying NVMe offload checkpoint from {offload_ckpt_dir} to {offload_dir}, {free / 1e9:,.2f} GB free on target filesystem..." + ) + copytree(offload_ckpt_dir, offload_dir, dirs_exist_ok=True) + _, _, free = disk_usage(offload_dir) + logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem") + self.optimizer.reset_swap_buffers() + + if self._optimizer_has_ckpt_event_epilogue(): + self.optimizer.checkpoint_event_epilogue() + + if self.load_universal_checkpoint() and not self.zero_optimization_partition_weights(): + self.optimizer.update_lp_params() + + return load_path, client_states + + def _load_checkpoint(self, + load_dir, + tag, + load_module_strict=True, + load_optimizer_states=True, + load_lr_scheduler_states=True, + load_module_only=False, + custom_load_fn=None): + + from deepspeed.runtime.state_dict_factory import SDLoaderFactory + + ckpt_list = self._get_all_ckpt_names(load_dir, tag) + sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine) + + is_pipe_parallel = isinstance(self.module, PipelineModule) + + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel) + + if checkpoint is None: + return None, None + + fetch_z3_params = False + if self.zero_optimization_partition_weights() and not load_optimizer_states: + checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir) + fetch_z3_params = True + + if is_pipe_parallel: + # Pipeline parallelism uses this to load its own checkpoint files. + self._curr_ckpt_path = os.path.join(load_dir, tag) + + if self.has_moe_layers: + # print(checkpoint.keys()) + old_moe_load = False + if not isinstance(checkpoint['num_experts'], list): + old_moe_load = True + DeepSpeedEngine.load_moe_state_dict(load_dir, + tag, + state_dict=checkpoint['module'], + old_moe_load=old_moe_load, + model=self.module, + mpu=self.mpu, + num_experts=self.num_experts, + checkpoint_engine=self.checkpoint_engine) + if not self.load_universal_checkpoint(): + self.load_module_state_dict(checkpoint=checkpoint, + strict=load_module_strict, + custom_load_fn=custom_load_fn, + fetch_z3_params=fetch_z3_params) + + self.loaded_checkpoint_dp_world_size = checkpoint['dp_world_size'] + + optim_checkpoint = None + if load_module_only: + deepspeed_states = ['module'] + if self.optimizer is not None and hasattr(self.optimizer, 'refresh_fp32_params'): + self.optimizer.refresh_fp32_params() + else: + has_zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() + if load_optimizer_states and self.optimizer is not None and not has_zero_optimizer_state: + if self.has_moe_layers: + largest_group_name = groups._get_max_expert_size_name() + expp_rank = groups._get_expert_parallel_rank(largest_group_name) + optim_load_path = self._get_optimizer_ckpt_name(load_dir, tag, expp_rank) + optim_checkpoint = self.checkpoint_engine.load(optim_load_path, map_location=torch.device('cpu')) + else: + optim_checkpoint = checkpoint + + if self.fp16_enabled() or self.bfloat16_enabled(): + self.optimizer.load_state_dict(optim_checkpoint['optimizer'], + load_optimizer_states=load_optimizer_states) + else: + optim_checkpoint = checkpoint + + self.optimizer.load_state_dict(optim_checkpoint['optimizer']) + + if load_lr_scheduler_states and self.lr_scheduler is not None: + self.lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) + + if self.random_ltd_enabled() and self.random_ltd_scheduler is not None and 'random_ltd' in checkpoint: + self.random_ltd_scheduler.load_state_dict(checkpoint['random_ltd']) + + if self.training_dataloader is not None and self.curriculum_learning_enabled( + ) and 'data_sampler' in checkpoint: + self.training_dataloader.data_sampler.load_state_dict(checkpoint['data_sampler']) + + def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters, loaded_parameters): + result = set() + + for name in original_set: + if name in loaded_parameters and name not in loaded_set: + continue # parameter existed in previous model and was not sparse + result.add(name) + + for name in loaded_set: + if name in original_parameters: + result.add(name) # parameter exists in both configs and it was sparse + + return result + + if 'sparse_tensor_module_names' in checkpoint: + sparse_tensor_module_names = checkpoint['sparse_tensor_module_names'] + elif 'csr_tensor_module_names' in checkpoint: + sparse_tensor_module_names = checkpoint['csr_tensor_module_names'] + else: + sparse_tensor_module_names = None + if sparse_tensor_module_names is not None: + if load_module_strict: + self.sparse_tensor_module_names = sparse_tensor_module_names + else: + self.sparse_tensor_module_names = get_sparse_tensor_module_names( + self.sparse_tensor_module_names, sparse_tensor_module_names, + dict(self.module.named_parameters()), checkpoint["module"]) + + self.global_steps = checkpoint['global_steps'] + self.global_samples = checkpoint.get('global_samples', self.global_steps * self.train_batch_size()) + self.skipped_steps = checkpoint['skipped_steps'] + self.loaded_checkpoint_mp_world_size = checkpoint['mp_world_size'] + deepspeed_states = [ + 'module', 'sparse_tensor_module_names', 'skipped_steps', 'global_steps', 'dp_world_size', + 'mp_world_size', 'data_sampler', 'random_ltd' + ] + client_state = {} + + if load_lr_scheduler_states: + deepspeed_states.append('lr_scheduler') + if load_optimizer_states: + deepspeed_states.append('optimizer') + + client_state = {key: value for key, value in checkpoint.items() if not key in deepspeed_states} + + if optim_checkpoint is not None: + client_state['optimizer'] = optim_checkpoint['optimizer'] + + return load_path, client_state + + def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): + + load_serial = None + # When use loading checkpoint serial, checkpoint loading start from local rank 0, + # all other local rank would be paused, waiting for its rank-1 peer ready and its notification. + if self._config.zero_config.pipeline_loading_checkpoint: + assert self.zero_optimization_stage( + ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading" + load_serial = torch.zeros(1).to(self.device) + if dist.get_local_rank() != 0: + dist.recv(tensor=load_serial, src=dist.get_rank() - 1) + if self.load_universal_checkpoint(): + zero_sd_list = None + checkpoint_folder = f'{os.path.join(load_dir, tag)}' + else: + if load_optimizer_states and self.seq_dp_world_size != self.loaded_checkpoint_dp_world_size: + raise ZeRORuntimeException("The checkpoint being loaded used a DP " \ + f"world size of {self.loaded_checkpoint_dp_world_size} but the " \ + f"current world size is {self.seq_dp_world_size}. Automatic adjustment " \ + "of ZeRO's optimizer state partitioning with a new world size is not " \ + "currently supported.") + checkpoint_folder = None + zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) + if zero_sd_list is None: + return False + + param_shapes = self._get_zero_param_shapes() + self.optimizer.load_state_dict(state_dict_list=zero_sd_list, + load_optimizer_states=load_optimizer_states, + load_from_fp32_weights=self.zero_load_from_fp32_weights(), + checkpoint_folder=checkpoint_folder, + load_serial=load_serial, + param_shapes=param_shapes) + + if self.load_universal_checkpoint(): + logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}') + else: + logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}") + return True + + def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode): + zero_ckpt_names = [] + for dp_rank in range(dp_world_size): + ckpt_name = self._get_rank_zero_ckpt_name(checkpoints_path=load_dir, + tag=tag, + mp_rank=mp_rank, + dp_rank=dp_rank, + bf16_mode=bf16_mode) + zero_ckpt_names.append(ckpt_name) + + return zero_ckpt_names + + def _get_all_zero_checkpoint_names(self, load_dir, tag, bf16_mode): + mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() + zero_ckpt_names = self._get_mp_rank_zero_checkpoint_names(load_dir=load_dir, + tag=tag, + mp_rank=mp_rank, + dp_world_size=self.loaded_checkpoint_dp_world_size, + bf16_mode=bf16_mode) + for i, ckpt_name in enumerate(zero_ckpt_names): + if not os.path.exists(ckpt_name): + # transparently handle the old file pattern for optim_states + if "optim_states.pt" in ckpt_name: + ckpt_name_try = ckpt_name.replace("_optim_states.pt", "optim_states.pt") + if os.path.exists(ckpt_name_try): + zero_ckpt_names[i] = ckpt_name_try + continue + + return zero_ckpt_names + + def _get_all_zero_checkpoint_state_dicts(self, zero_ckpt_names): + zero_sd_list = [] + for i, ckpt_name in enumerate(zero_ckpt_names): + _state = None + if ckpt_name is None: + _state = {OPTIMIZER_STATE_DICT: None} + # Fully load state for current rank + elif self.zero_elastic_checkpoint() or dist.get_rank(group=self.optimizer.dp_process_group) == i: + _state = self.checkpoint_engine.load( + ckpt_name, + map_location='cpu', + ) + else: + _state = {OPTIMIZER_STATE_DICT: None} + zero_sd_list.append(_state) + + zero_optimizer_sd = [sd[OPTIMIZER_STATE_DICT] for sd in zero_sd_list] + logger.info(f"successfully read {len(zero_optimizer_sd)} ZeRO state_dicts for rank {self.global_rank}") + return zero_optimizer_sd + + def _get_all_zero_checkpoints(self, load_dir, tag): + for bf16_mode in [self.bfloat16_enabled(), not self.bfloat16_enabled()]: + zero_ckpt_names = self._get_all_zero_checkpoint_names(load_dir, tag, bf16_mode) + if zero_ckpt_names is not None: + # Warn if loading checkpoint of different bit16 type + if bf16_mode is not self.bfloat16_enabled(): + checkpoint_bit16 = BFLOAT16 if bf16_mode else FP16 + engine_bit16 = BFLOAT16 if self.bfloat16_enabled() else FP16 + logger.warning(f'Loading {checkpoint_bit16} zero checkpoints into {engine_bit16} training engine') + return self._get_all_zero_checkpoint_state_dicts(zero_ckpt_names) + + return None + + def _checkpoint_tag_validation(self, tag): + if self.checkpoint_tag_validation_enabled(): + s_hash = hashlib.sha1(tag.encode()) + bhash = torch.ByteTensor([s_hash.digest()]).flatten().to(self.device) + max_bhash = bhash.clone() + min_bhash = bhash.clone() + dist.all_reduce(max_bhash, op=dist.ReduceOp.MAX) + dist.all_reduce(min_bhash, op=dist.ReduceOp.MIN) + valid = all(min_bhash == bhash) and all(max_bhash == bhash) + msg = (f"[rank={dist.get_rank()}] The checkpoint tag name '{tag}' is not consistent across " + "all ranks. Including rank unique information in checkpoint tag could cause issues when " + "restoring with different world sizes.") + if self.checkpoint_tag_validation_fail(): + assert valid, msg + elif not valid: + logger.warning(msg) + + def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False): + """Save training checkpoint + + Arguments: + save_dir: Required. Directory for saving the checkpoint + tag: Optional. Checkpoint tag used as a unique identifier for the checkpoint, global step is + used if not provided. Tag name must be the same across all ranks. + client_state: Optional. State dictionary used for saving required training states in the client code. + save_latest: Optional. Save a file 'latest' pointing to the latest saved checkpoint. + exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. + Important: all processes must call this method and not just the process with rank 0. It is + because each process needs to save its master weights and scheduler+optimizer states. This + method will hang waiting to synchronize with other processes if it's called just for the + process with rank 0. + + """ + if self._optimizer_has_ckpt_event_prologue(): + # Custom preparation for checkpoint save, if applicable + self.optimizer.checkpoint_event_prologue() + + rank = self.local_rank if self.use_node_local_storage() else self.global_rank + + # This is to make sure the checkpoint names are created without collision + # There seems to be issue creating them in parallel + + # Ensure save_dir directory exists + if rank == 0: + self.checkpoint_engine.makedirs(save_dir, exist_ok=True) + dist.barrier() + + if tag is None: + tag = f"global_step{self.global_steps}" + + # Ensure tag is a string + tag = str(tag) + self.checkpoint_engine.create(tag) + + # Ensure checkpoint tag is consistent across ranks + self._checkpoint_tag_validation(tag) + + if self.has_moe_layers: + self.save_non_zero_checkpoint = False + self._create_checkpoint_file(save_dir, tag, False) + self._save_moe_checkpoint(save_dir, + tag, + client_state=client_state, + exclude_frozen_parameters=exclude_frozen_parameters) + + # We distribute the task of saving layer checkpoint files among + # data parallel instances, so all procs should call _save_checkpoint. + # All procs then call module_state_dict(), but only procs of data + # parallel rank 0 save the general model params. + if not self.has_moe_layers: + self._create_checkpoint_file(save_dir, tag, False) + self._save_checkpoint(save_dir, + tag, + client_state=client_state, + exclude_frozen_parameters=exclude_frozen_parameters) + + if self.save_zero_checkpoint: + self._create_zero_checkpoint_files(save_dir, tag) + self._save_zero_checkpoint(save_dir, tag) + + if self.zero_nvme_offload_optimizer(): + from shutil import copytree, disk_usage + offload_dir = self.optimizer.optimizer_swapper.swap_folder + offload_ckpt_dir = os.path.join(save_dir, tag, "offloaded_tensors") + _, _, free = disk_usage(save_dir) + logger.info( + f"Copying NVMe offload files from {offload_dir} to {offload_ckpt_dir}, {free / 1e9:,.2f} GB free on target filesystem..." + ) + copytree(offload_dir, + offload_ckpt_dir, + ignore=lambda _, dir_list: list(filter(lambda x: 'gradient' in x, dir_list)), + dirs_exist_ok=False) + _, _, free = disk_usage(save_dir) + logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem") + + if self._optimizer_has_ckpt_event_epilogue(): + self.optimizer.checkpoint_event_epilogue() + + # Save latest checkpoint tag + self.checkpoint_engine.commit(tag) + if save_latest and rank == 0: + with open(os.path.join(save_dir, 'latest'), 'w') as fd: + fd.write(tag) + + dist.barrier() + + return True + + def _get_non_moe_state_dict(self, full_state_dict): + """ + Get the state dict of the non-moe layers + """ + for key in list(full_state_dict.keys()): + if 'expert' in key and 'moe.gate.wg.weight' not in key: + full_state_dict.pop(key) + + return full_state_dict + + def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): + save_path = self._get_ckpt_name(save_dir, tag) + # A hack to save the checkpointing directory. Pipeline parallelism overrides + # module_state_dict() and uses this path to save the model. module_state_dict() + # then instead just returns None. + + # Using layer_#_export_# to save the model's expert state_dict + moe_layer_id = 0 + for n_module, module in self.module.named_modules(): + if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: + group_name = module.expert_group_name + num_local_experts = module.num_local_experts + expp_rank = groups._get_expert_parallel_rank(group_name) + exp_dp_rank = groups._get_expert_data_parallel_rank(group_name) + # print(expp_rank, exp_dp_rank) + if exp_dp_rank != 0: + moe_layer_id += 1 + continue + + # get all moe parameters + moe_state_dict = {} + for n, p in module.state_dict().items(): + if 'expert' in n and 'moe.gate.wg.weight' not in n: + moe_state_dict[n_module + '.' + n] = p + moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' + # print(moe_state_dict.keys()) # until now, everything is fine. So the bug happens at next few lines + # Reorder the moe name rank, so that each checkpoint only has one expert + experts_state_dict = defaultdict(dict) + for key in list(moe_state_dict.keys()): + m = re.match(f".*{moe_str_prefix}([0-9]+).*", key) + + local_expert_id = None + if not m: + logger.warning(f'No expert found in key {key}.') + else: + local_expert_id = m.group(1) + + global_expert_id = expp_rank * \ + num_local_experts + int(local_expert_id) + expert_key = key.replace(f'{moe_str_prefix}{local_expert_id}', + f'{moe_str_prefix}{global_expert_id}') + # truncating extra tensor (shared) storage + truncated = moe_state_dict.pop(key).clone().detach() + experts_state_dict[str(global_expert_id)][expert_key] = truncated + + # let save the moe parameters + for global_expert_id, expert_state_dict in experts_state_dict.items(): + # save the moe parameters + moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag, self.mpu) + if self.random_ltd_enabled(): + expert_state_dict = remove_random_ltd_state_dict(expert_state_dict) + self.checkpoint_engine.save(expert_state_dict, moe_save_path) + moe_layer_id += 1 + + self._curr_ckpt_path = os.path.join(save_dir, tag) + + largest_group_name = groups._get_max_expert_size_name() + expp_rank = groups._get_expert_parallel_rank(largest_group_name) + exp_dp_rank = groups._get_expert_data_parallel_rank(largest_group_name) + + # In the case of E + D parallelism, only the + # first expert parallel group should save the expert weights + # since each expert parallel group is a copy of the model's experts + if exp_dp_rank == 0: + # Save optimizer states. They are different across each exp parallel rank. + optimizer_state = { + 'optimizer': self.optimizer.state_dict() if self.optimizer and not self.zero_optimization() else None + } + # TODO: why use BufferedWriter not the path + file_path = self._get_optimizer_ckpt_name(save_dir, tag, expp_rank) + self.checkpoint_engine.save(optimizer_state, file_path) + + # Load flow uses below saved file for model parameters, RNG and more + if groups._get_data_parallel_rank() == 0: + # Get non-moe parameters + # Classes DeepSpeedEngine and PipelineEngine have different behavior for method module_state_dict. + # DeepSpeedEngine returns the state dict, where PipelineEngine saves the state dict and returns None. + # We need to get the state dict, therefore, call to DeepSpeedEngine (base class for PipelineEngine) + model_state_dict = self._get_non_moe_state_dict( + DeepSpeedEngine.module_state_dict(self, exclude_frozen_parameters=exclude_frozen_parameters)) + + # TODO: update num experts info,.. in checkpoint + state = { + 'module': + model_state_dict, + 'lr_scheduler': + self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, + 'data_sampler': + self.training_dataloader.data_sampler.state_dict() if + (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, + 'random_ltd': + self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, + 'sparse_tensor_module_names': + self.sparse_tensor_module_names, + 'skipped_steps': + self.skipped_steps, + 'global_steps': + self.global_steps, + 'global_samples': + self.global_samples, + 'dp_world_size': + self.dp_world_size, + 'mp_world_size': + self.mp_world_size, + 'num_experts': + self.num_experts + } + state.update(client_state) + logger.info(f'Saving model checkpoint: {save_path}') + self.checkpoint_engine.save(state, save_path) + + def _create_checkpoint_file(self, save_dir, tag, zero_checkpoint): + name_function = (self._get_zero_ckpt_name if zero_checkpoint else self._get_ckpt_name) + try: + checkpoint_name = name_function(save_dir, tag) + path = os.path.dirname(checkpoint_name) + self.checkpoint_engine.makedirs(path, exist_ok=True) + except: + logger.error(f"Failed saving model checkpoint to {save_dir} with tag {tag}") + return False + + return True + + def _create_zero_checkpoint_files(self, save_dir, tag): + success = True + # zero checkpoint files are created sequentially + for rank in range(dist.get_world_size(self.optimizer.dp_process_group)): + if rank == self.global_rank: + success = self._create_checkpoint_file(save_dir, tag, True) + + dist.barrier(group=self.optimizer.dp_process_group) + + return success + + def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): + + save_path = self._get_ckpt_name(save_dir, tag) + + zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() + + save_frozen_param = self.zero_optimization_partition_gradients() and not exclude_frozen_parameters + + # A hack to save the checkpointing directory. Pipeline parallelism overrides + # module_state_dict() and uses this path to save the model. module_state_dict() + # then instead just returns None. The module_state_dict() implementation in + # PipelineEngine expects the save path to be set in self._curr_ckpt_path. + self._curr_ckpt_path = os.path.join(save_dir, tag) + module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) + self._curr_ckpt_path = None + + state = dict(module=module, + buffer_names=self._get_buffer_names(), + optimizer=self.optimizer.state_dict() if self.optimizer and not zero_optimizer_state else None, + param_shapes=self._get_zero_param_shapes() if self.optimizer and zero_optimizer_state else None, + frozen_param_shapes=self._get_zero_frozen_param_attributes(self._get_param_shape_func) + if save_frozen_param else None, + shared_params=self._get_shared_params() if self.optimizer and zero_optimizer_state else None, + frozen_param_fragments=self._get_zero_frozen_param_attributes(self._get_param_fragment_func) + if save_frozen_param else None, + lr_scheduler=self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, + data_sampler=self.training_dataloader.data_sampler.state_dict() if + (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, + random_ltd=self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, + sparse_tensor_module_names=self.sparse_tensor_module_names, + skipped_steps=self.skipped_steps, + global_steps=self.global_steps, + global_samples=self.global_samples, + dp_world_size=self.seq_dp_world_size, + mp_world_size=self.mp_world_size, + ds_config=self.config, + ds_version=version) + state.update(client_state) + + if self.save_non_zero_checkpoint: + log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0, 1]) + self.checkpoint_engine.save(state, save_path) + + def _get_buffer_names(self): + buffer_names = [] + + # we save buffer names so that we could extract later the real buffers from the saved + # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they + # are intermixed with param placeholders + + # have to traverse the tree to be able to skip non-persistent buffers + def get_layer_named_buffers(module, prefix=""): + for name, buf in module.named_buffers(recurse=False): + if buf is not None and name not in module._non_persistent_buffers_set: + buffer_names.append(prefix + name) + + for name, child in module.named_children(): + if child is not None: + get_layer_named_buffers(child, prefix + name + ".") + + get_layer_named_buffers(self.module, prefix="") + + return buffer_names + + def _get_param_shape_func(self, param): + return param.ds_shape if hasattr(param, 'ds_id') else param.shape + + def _get_param_fragment_func(self, param): + return param.ds_tensor.detach().cpu() if hasattr(param, 'ds_id') else param.detach().cpu() + + def _get_zero_frozen_param_attributes(self, attr_func): + frozen_param_fragments = OrderedDict() + + for param in self.module.parameters(): + if param.requires_grad: + continue + if param not in self.param_names: + raise ValueError(f"failed to find frozen {param} in named params") + name = self.param_names[param] + frozen_param_fragments[name] = attr_func(param) + + return frozen_param_fragments + + def _get_zero_param_shapes(self): + """Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the + optimizer. the names are exactly as in state_dict. The order is absolutely important, since + the saved data is just flattened data with no identifiers and requires reconstruction in the + same order it was saved. + We can't rely on self.module.named_parameters() to get the saved tensors, as some params + will be missing and others unsaved and then it'd be impossible to reconstruct state_dict + from the flattened weights. + optimizer.bit16_groups seems to be the easiest to use as it's in all zeroX versions. + """ + param_group_shapes = [] + cnt = 0 + numel = 0 + + # zero2 started using a round_robin_bit16_groups which is a shuffled version of bit16_groups - + # if we don't use it, we get parameters ordered incorrectly + if hasattr(self.optimizer, "round_robin_bit16_groups"): + bit16_groups = self.optimizer.round_robin_bit16_groups + elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"): + bit16_groups = self.optimizer.bf16_groups + else: + bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage( + ) == 2 else self.optimizer.fp16_groups + + for bit16_group in bit16_groups: + param_shapes = OrderedDict() + for param in bit16_group: + cnt += 1 + numel += param.ds_numel if hasattr(param, "ds_numel") else param.numel() + shape = param.ds_shape if hasattr(param, "ds_shape") else param.shape + if param not in self.param_names: + raise ValueError(f"failed to find optimizer param in named params") + name = self.param_names[param] + param_shapes[name] = shape + + # uncomment to debug zero_to_fp32.py problems + # if self.global_rank == 0: print(f"saving param {name} {shape} (numel={shape.numel()})") + param_group_shapes.append(param_shapes) + # if self.global_rank == 0: print(f"Total saved {numel} numels in {cnt} params") + + return param_group_shapes + + def _get_shared_params(self): + """ + Returns a dict of shared params, which can later be used to reconstruct the original state dict, + e.g. in `zero_to_fp32`. Each dict entry is a pair of param names, where the key is the name + of the variable that isn't stored and the value is the actual param holding data. + """ + shared_index = {} + shared_params_by_full_name = {} + + is_zero3_model = (self.zero_optimization_partition_weights() + and any(hasattr(param, "ds_id") for param in self.module.parameters())) + + def get_layer_state_dict(module, prefix=""): + # handle params + for name, param in module.named_parameters(recurse=False): + if param is None or (is_zero3_model and not hasattr(param, "ds_id")): + continue + key = prefix + name + + # When weights are manged by stage 3, we can't rely on param.data_ptr() as it will be reused + # as weights get gathered and reduced, but param.ds_id is unique across all zero weights + # (and shared params will have the same param.ds_id) + param_id = param.ds_id if is_zero3_model else param.data_ptr() + + if param_id in shared_index: + # shared weights + #print(f"`{key}` is shared with `{shared_index[param_id]}`") + shared_params_by_full_name[key] = shared_index[param_id] + else: + shared_index[param_id] = key + + for name, child in module.named_children(): + if child is not None: + get_layer_state_dict(child, prefix + name + ".") + + if dist.get_rank() == 0: + get_layer_state_dict(self.module, prefix="") + + return shared_params_by_full_name + + def _copy_recovery_script(self, save_path): + base_dir = os.path.dirname(os.path.dirname(__file__)) + script = "zero_to_fp32.py" + src = os.path.join(base_dir, "utils", script) + dst = os.path.join(save_path, script) + #logger.info(f"creating recovery script {dst}") + copyfile(src, dst) + self._change_recovery_script_permissions(dst) + + def _change_recovery_script_permissions(self, dst): + # make executable (safeguard for file shares - Azure as example) + try: + os.chmod(dst, os.stat(dst).st_mode | stat.S_IEXEC) + except (FileNotFoundError, PermissionError) as e: + #this message is used in unit test TestZeRONonDistributed + logger.info( + f'Warning: Could not change permissions for {dst} due to error: {e}. Continuing without changing permissions.' + ) + + def _save_zero_checkpoint(self, save_path, tag): + zero_checkpoint_name = self._get_zero_ckpt_name(save_path, tag) + zero_sd = dict(optimizer_state_dict=self.optimizer.state_dict(), ds_config=self.config, ds_version=version) + self.checkpoint_engine.save(zero_sd, zero_checkpoint_name) + + if self.global_rank == 0: + self._copy_recovery_script(save_path) + ckpt_type = 'zero' if self.zero_optimization() else 'bf16_zero' + logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}') + + def _replace_module_consolidated_state_dict(self): + """ + Get a full non-partitioned state_dict with fp16 weights on cpu. + Important: this function must be called on all ranks and not just rank 0. + This is similar to nn.Module.state_dict (modelled after _save_to_state_dict) + This method is used for tensor parallel training. + + Returns: + OrderedDict: The consolidated state dictionary if the current process rank is 0, otherwise None. + """ + #TODO: If we use both Zero3 and tensor parallel simultaneously + # we need to consolidate the gather mechanisms of both. + state_dict = OrderedDict() if dist.get_rank() == 0 else None + + def get_layer_state_dict(module, prefix=""): + with GatherReplacedLayerParams(list(module.parameters(recurse=False)), module, enabled=True): + for name, param in module.named_parameters(recurse=False): + if param is None: + continue + key = prefix + name + if (dist.get_rank() == 0): + state_dict[key] = param.detach().cpu() + # print(key,module, param.detach().cpu().shape) + + for name, child in module.named_children(): + if child is not None: + get_layer_state_dict(child, prefix + name + ".") + + get_layer_state_dict(self.module, prefix="") + + # ensure that all GPU communication tasks are completed before the process exits + get_accelerator().synchronize() + return state_dict + + def _consolidated_16bit_state_dict(self, exclude_frozen_parameters=False): + """ + Consolidate the 16-bit state dictionary. + """ + if self.zero_optimization_stage() == ZeroStageEnum.weights: + return self._zero3_consolidated_16bit_state_dict(exclude_frozen_parameters) + elif self.autotp_size() > 1: + return self._replace_module_consolidated_state_dict() + + raise ValueError("consolidated_16bit_state_dict is only applicable to cases where weights are partitioned, " + "including Zero Stage 3 and tensor parallelism.") + + def _zero3_consolidated_16bit_state_dict(self, exclude_frozen_parameters=False): + """ + Get a full non-partitioned state_dict with fp16 weights on cpu. + Important: this function must be called on all ranks and not just rank 0. + This is similar to nn.Module.state_dict (modelled after _save_to_state_dict), but: + 1. consolidates the weights from different partitions on gpu0 + 2. works on one layer at a time to require as little gpu0 memory as possible, by + moving the already consolidated weights to cpu + 3. takes care to keep the shared params shared when gradually copying the params to cpu + Returns: + a consolidated fp16 ``state_dict`` on cpu on rank 0, ``None`` on other ranks + """ + if not self.zero_optimization_partition_weights(): + raise ValueError("this function requires ZeRO-3 mode") + + state_dict = OrderedDict() if dist.get_rank() == 0 else None + shared_params = {} + + def get_layer_state_dict(module, prefix=""): + # gather one layer at a time to be memory-efficient + # must use modifier_rank=0 to release GPU memory after each layer gathered + #see_memory_usage("before GatheredParameters", force=True) + with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0): + if dist.get_rank() == 0: + # handle params + for name, param in module.named_parameters(recurse=False): + if param is None or (exclude_frozen_parameters and not param.requires_grad): + continue + key = prefix + name + # can't rely on param.data_ptr() as it will be reused as weights gets + # gathered and reduced, but param.ds_id is unique across all zero weights + # (and shared params will have the same param.ds_id) + if param.ds_id in shared_params: + # shared weights + #print(f"`{key}` is shared with `{shared_params[param.ds_id]}`") + state_dict[key] = state_dict[shared_params[param.ds_id]] + else: + state_dict[key] = param.detach().cpu() + shared_params[param.ds_id] = key + #print(f"param {param.ds_id} {param.shape} {key} ") + + # now buffers - not sure if need to take care of potentially shared weights here + for name, buf in module.named_buffers(recurse=False): + if (buf is not None and name not in module._non_persistent_buffers_set): + state_dict[prefix + name] = buf.detach().cpu() + #see_memory_usage("after GatheredParameters", force=True) + + for name, child in module.named_children(): + if child is not None: + get_layer_state_dict(child, prefix + name + ".") + + # Prepare for checkpoint save by ensuring all parameters are partitioned + if self._optimizer_has_ckpt_event_prologue(): + self.optimizer.checkpoint_event_prologue() + + see_memory_usage("before get_layer_state_dict", force=False) + get_layer_state_dict(self.module, prefix="") + see_memory_usage("after get_layer_state_dict", force=False) + + if self._optimizer_has_ckpt_event_epilogue(): + self.optimizer.checkpoint_event_epilogue() + + return state_dict + + def save_fp16_model(self, save_dir, save_filename="pytorch_model.bin"): + """has been renamed to save_16bit_model, keeping this around for backwards + compatibility""" + return self.save_16bit_model(save_dir, save_filename) + + def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin", exclude_frozen_parameters=False): + """ + Save 16bit model weights + + This method saves the 16bit model weights at the desired destination. + + Arguments: + save_dir: Required. Directory for saving the model + save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin`` + exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. + + Returns: + ``True`` when a model has been saved, ``False`` otherwise. It will not be saved if + stage3_gather_16bit_weights_on_model_save is ``False``. + + Important: all processes must call this method and not just the process with rank 0. It is + because the processes need to work in sync to gather the weights. This method will hang + waiting to synchronize with other processes if it's called just for the process with rank 0. + + """ + + path = os.path.join(save_dir, save_filename) + + if self.zero_optimization_partition_weights(): + if self.zero_gather_16bit_weights_on_model_save(): + # consolidation is expensive in time and memory and therefore isn't a default + state_dict = self._zero3_consolidated_16bit_state_dict( + exclude_frozen_parameters=exclude_frozen_parameters) + else: + # the model will be bogus if not consolidated so don't confuse the user by saving it + logger.info( + f"Did not save the model {path} because stage3_gather_16bit_weights_on_model_save is False") + return False + else: + state_dict = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) + + tag = f"global_step{self.global_steps}" + tag = str(tag) + self.checkpoint_engine.create(tag) + + if dist.get_rank() == 0: + self.checkpoint_engine.makedirs(save_dir, exist_ok=True) + logger.info(f"Saving model weights to {path}, tag: {tag}") + self.checkpoint_engine.save(state_dict, path) + + self.checkpoint_engine.commit(tag) + + return True + + def empty_partition_cache(self): + """ + Release GPU memory consumed by offloaded model parameters. + """ + if hasattr(self.optimizer, 'empty_partition_cache'): + self.optimizer.empty_partition_cache() + gc.collect() + get_accelerator().empty_cache() + + def compile(self, backend=get_accelerator().get_compile_backend(), compile_kwargs={}, schedule=None) -> None: + """Compile the module using the specified backend and kwargs. + If a compiler_fn is set, it will be used instead of torch.compile(). + """ + # Avoid graph breaks + deepspeed.utils.nvtx.enable_nvtx = False + + if not is_compile_supported(): + raise RuntimeError("compile is not supported in your version of PyTorch.") + + if self.is_compiled: + return + + if 'backend' in compile_kwargs: + logger.warning("The `backend` in `compile_kwargs` will be overridden. Use the `backend` argument instead.") + + print(f"Compiling deepcompile={self.is_deepcompile_enabled()} backend={backend}") + + if self.is_deepcompile_enabled(): + assert self.zero_optimization_stage() == ZeroStageEnum.optimizer_states \ + or self.zero_optimization_stage() == ZeroStageEnum.weights \ + , "Currently DeepCompile supports stage 1 or 3 only." + + assert not isinstance(self.optimizer, + DeepSpeedZeRoOffload), "Currently DeepCompile is not supported without an optimizer." + + if schedule is not None: + + def passes_name_to_fn(passes): + for p in passes: + assert callable(p) or p in opt_passes, f"Unknown pass {p}" + return [p if callable(p) else opt_passes[p] for p in passes] + + schedule = [(step, passes_name_to_fn(passes)) for step, passes in schedule] + + assert backend in ['inductor', 'eager'], f"Backend {backend} is not supported for DeepCompile." + + compile_config = self._config.compile_config + if (("zero_optimization" in self.config and "offload_optimizer" in self.config["zero_optimization"] + and "offload_param" in self.config["zero_optimization"]) + and self._config.zero_config.offload_param.device == "cpu" + and self._config.zero_config.offload_optimizer.device == "cpu"): + compile_config.offload_parameters = True + if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states: + backend = init_z1(self, backend, compile_config, compile_kwargs, schedule) + elif self.zero_optimization_stage() == ZeroStageEnum.weights: + backend = init_z3(self, backend, compile_config, compile_kwargs, schedule) + + # create new dict to avoid modifying original dict + self.module.compile(**{**compile_kwargs, 'backend': backend}) + + self._is_compiled = True + + def get_compile_time(self): + from deepspeed.compile.backend import opt_pass_times + return opt_pass_times + + def register_compile_pass(self, pass_name: str, pass_fn: Callable) -> None: + register_compile_pass(pass_name, pass_fn) + + def is_deepcompile_enabled(self): + return self._config.compile_config.deepcompile + + @property + def is_compiled(self) -> bool: + return self._is_compiled + + def offload_states(self, + include: Container[OffloadStateTypeEnum] = None, + device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, + pin_memory: bool = True, + non_blocking: bool = False) -> None: + """Offload the engine's states to the specified device. + + Arguments: + include: Optional. The set of states to offload. If not provided, all states are offloaded. + device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported. + pin_memory: Optional. Whether to pin the memory of the offloaded states. + non_blocking: Optional. Whether to offload the states asynchronously. + """ + assert self.zero_optimization_stage( + ) == ZeroStageEnum.weights, "Moving buffers across devices is supported only for ZeRO stage 3." + + opt_offload_config = self.zero_offload_optimizer() + assert opt_offload_config is None or opt_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded optimizer states." + param_offload_config = self.zero_offload_param() + assert param_offload_config is None or param_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded parameters." + + assert not isinstance( + self.optimizer, + DeepSpeedZeRoOffload), "Moving states across devices is not supported without an optimizer." + + if device == OffloadDeviceEnum.none: + logger.warning("No device specified for offloading states.") + return + + if device == OffloadDeviceEnum.nvme: + raise ValueError("NVMe offload is not supported for offloading states.") + + self.optimizer.offload_states(include=include, device=device, pin_memory=pin_memory, non_blocking=non_blocking) + + def reload_states(self, non_blocking: bool = False) -> None: + """Reload the engine states to the original device. + + Arguments: + non_blocking: Optional. Whether to offload the states asynchronously. + """ + assert self.zero_optimization_stage( + ) == ZeroStageEnum.weights, "Moving buffers back is supported only for ZeRO stage 3." + + assert not isinstance( + self.optimizer, + DeepSpeedZeRoOffload), "Moving states across devices is not supported without an optimizer." + + self.optimizer.reload_states(non_blocking=non_blocking) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cf58eb1801af7cc02f3cc2057a5c843cd078b575 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .module import PipelineModule, LayerSpec, TiedLayerSpec +from .topology import ProcessTopology diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ff1b21f15e378d7d9c6fd3d72b3f2793d8e37fd Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..809da1c13956c9064c86a00507e163b4be78c421 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b3a04ecbf61774fabb1a46381ffd0455579f788 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5526e77e02f2160ad070af821598587c5a492b95 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..967f5e91b94b6d854ec14c531cdc1b666b62d98f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29ff0beae3091877bca5dcd20a1217ae0435b4e5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/p2p.py b/lib/python3.12/site-packages/deepspeed/runtime/pipe/p2p.py new file mode 100644 index 0000000000000000000000000000000000000000..ed6d80b8d4fbf5131d773d72bec8574f3105bac9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/pipe/p2p.py @@ -0,0 +1,182 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import msgpack +import typing + +import torch +from deepspeed import comm as dist + +from deepspeed.utils.torch import required_torch_version +from deepspeed.accelerator import get_accelerator + +_groups = None +_grid = None + +_async = [] + + +def can_send_recv() -> bool: + return required_torch_version(min_version=1.8) + + +#initializes adjacent process groups +#run this only after deepspeed.init_distributed() has been called +def init_process_groups(grid): + global _groups, _grid + _grid = grid + + assert _grid.pipe_parallel_size > 1, "There is no pipeline parallelism" + + if not can_send_recv(): + _groups = [dist.new_group(ranks=group) for group in _grid.p2p_groups] + + +def _is_valid_send_recv(src_stage, dest_stage): + first_stage = 0 + last_stage = _grid.pipe_parallel_size - 1 + assert abs(src_stage-dest_stage) == 1 or \ + (src_stage == first_stage and dest_stage == last_stage) or \ + (src_stage == last_stage and dest_stage == first_stage), \ + "Functionality currently limited to send and receive between adjacent ranks only" + + +def send(tensor, dest_stage, async_op=False): + global _groups + assert async_op == False, "Doesn't support async_op true" + src_stage = _grid.get_stage_id() + _is_valid_send_recv(src_stage, dest_stage) + + dest_rank = _grid.stage_to_global(stage_id=dest_stage) + if async_op: + global _async + op = dist.isend(tensor, dest_rank) + _async.append(op) + else: + + if can_send_recv(): + return dist.send(tensor, dest_rank) + else: + group = _get_send_recv_group(src_stage, dest_stage) + src_rank = _grid.stage_to_global(stage_id=src_stage) + return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) + + +def recv(tensor, src_stage, async_op=False): + global _groups + assert async_op == False, "Doesn't support async_op true" + dest_stage = _grid.get_stage_id() + _is_valid_send_recv(src_stage, dest_stage) + + src_rank = _grid.stage_to_global(stage_id=src_stage) + + if async_op: + global _async + op = dist.irecv(tensor, src_rank) + _async.append(op) + else: + if can_send_recv(): + return dist.recv(tensor, src_rank) + else: + group = _get_send_recv_group(src_stage, dest_stage) + return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) + + +def wait(): + global _async + for op in _async: + op.wait() + _async = [] + + get_accelerator().synchronize() + + +def send_obj(msg: typing.Any, dest: int): + """Send an arbitrary python object to ``dest``. + + Note: ``msg`` must be serializable by msgpack. + + WARN: This incurs a CPU -> GPU transfer and should be used sparingly + for performance reasons. + + Args: + msg (typing.Any): The object to send. + dest (int): Destination rank. + """ + # serialize the message + msg = msgpack.packb(msg) + # construct a tensor to send + msg = torch.ByteTensor(torch.ByteStorage.from_buffer(msg)).to(get_accelerator().device_name()) + + # Send meta and message + length_tensor = torch.tensor([len(msg)], dtype=torch.long).to(get_accelerator().device_name()) + dist.send(length_tensor, dst=dest) + dist.send(msg, dst=dest) + + +def recv_obj(sender: int) -> typing.Any: + """Receive an arbitrary python object from ``sender``. + + WARN: This incur a CPU <-> GPU transfers and should be used sparingly + for performance reasons. + + Args: + sender (int): The rank sending the message. + """ + # Get message meta + length = torch.tensor([0], dtype=torch.long).to(get_accelerator().device_name()) + dist.recv(length, src=sender) + + # Receive and deserialize + msg = torch.empty(length.item(), dtype=torch.uint8).to(get_accelerator().device_name()) + dist.recv(msg, src=sender) + + msg = msgpack.unpackb(msg.cpu().numpy().tobytes()) + + def _to(x): + """Recursively move to the current device.""" + if torch.is_tensor(x): + return x.to(get_accelerator().device_name()) + if isinstance(x, (tuple, list)): + ret = [_to(x_) for x_ in x] + if isinstance(x, tuple): + ret = tuple(ret) + return ret + # handle kwargs + if isinstance(x, dict): + ret = dict() + for key, val in x.items(): + ret[_to(key)] = _to(val) + return ret + + # Anything else is a no-op + return x + + msg = _to(msg) + return msg + + +def _get_send_recv_group(src_stage, dest_stage): + '''the group id is always the smaller rank unless its a wrap around''' + + stage_id = None + + first_stage = 0 + last_stage = _grid.pipe_parallel_size - 1 + + if (src_stage == first_stage and dest_stage == last_stage + or dest_stage == first_stage and src_stage == last_stage): + stage_id = last_stage + elif src_stage > dest_stage: + stage_id = dest_stage + else: + stage_id = src_stage + '''group_id corresponds to group of [group_id, group_id+1] + unless group_id is the rank of the last stage + in which case group_id corresponds to group[group_id-num_stages+1, group_id] + ''' + group_id = _grid.stage_to_global(stage_id=stage_id) + + return _groups[group_id] diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/topology.py b/lib/python3.12/site-packages/deepspeed/runtime/pipe/topology.py new file mode 100644 index 0000000000000000000000000000000000000000..328c199071009669f41a6c2f531d1e741bb690f4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/pipe/topology.py @@ -0,0 +1,456 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed import comm as dist + +from collections import namedtuple +from itertools import product as cartesian_product + + +class ProcessTopology: + """ Manages the mapping of n-dimensional Cartesian coordinates to linear + indices. This mapping is used to map the rank of processes to the grid + for various forms of parallelism. + + Each axis of the tensor is accessed by its name. The provided ordering + of the axes defines the layout of the topology. ProcessTopology uses a "row-major" + layout of the tensor axes, and so axes=['x', 'y'] would map coordinates (x,y) and + (x,y+1) to adjacent linear indices. If instead axes=['y', 'x'] was used, coordinates + (x,y) and (x+1,y) would be adjacent. + + Some methods return ProcessCoord namedtuples. + """ + + def __init__(self, axes, dims): + """Create a mapping of n-dimensional tensor coordinates to linear indices. + + Arguments: + axes (list): the names of the tensor axes + dims (list): the dimension (length) of each axis of the topology tensor + """ + + self.axes = axes # names of each topology axis + self.dims = dims # length of each topology axis + + # This is actually a class that lets us hash {'row':3, 'col':2} mappings + self.ProcessCoord = namedtuple('ProcessCoord', axes) + + self.mapping = {} + ranges = [range(d) for d in dims] + # example: 1, (0,0,1) + for global_rank, coord in enumerate(cartesian_product(*ranges)): + key = {axis: coord[self.axes.index(axis)] for axis in self.axes} + key = self.ProcessCoord(**key) + # for example, {ProcessCoord(row=0, col=1) : 1} + self.mapping[key] = global_rank + + def get_rank(self, **coord_kwargs): + """Return the global rank of a process via its coordinates. + + Coordinates are specified as kwargs. For example: + + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> X.get_rank(x=0, y=1) + 1 + """ + if len(coord_kwargs) != len(self.axes): + raise ValueError('get_rank() does not support slices. Use filter_match())') + + key = self.ProcessCoord(**coord_kwargs) + assert key in self.mapping, f'key {coord_kwargs} invalid' + return self.mapping[key] + + def get_axis_names(self): + """Return a list of the axis names in the ordering of the topology. """ + return self.axes + + def get_rank_repr(self, rank, omit_axes=['data', 'pipe'], inner_sep='_', outer_sep='-'): + """Return a string representation of a rank. + + This method is primarily used for checkpointing model data. + + For example: + >>> topo = Topo(axes=['a', 'b'], dims=[2, 2]) + >>> topo.get_rank_repr(rank=3) + 'a_01-b_01' + >>> topo.get_rank_repr(rank=3, omit_axes=['a']) + 'b_01' + + Args: + rank (int): A rank in the topology. + omit_axes (list, optional): Axes that should not be in the representation. Defaults to ['data', 'pipe']. + inner_sep (str, optional): [description]. Defaults to '_'. + outer_sep (str, optional): [description]. Defaults to '-'. + + Returns: + str: A string representation of the coordinate owned by ``rank``. + """ + omit_axes = frozenset(omit_axes) + axes = [a for a in self.get_axis_names() if a not in omit_axes] + names = [] + for ax in axes: + ax_rank = getattr(self.get_coord(rank=rank), ax) + names.append(f'{ax}{inner_sep}{ax_rank:02d}') + return outer_sep.join(names) + + def get_dim(self, axis): + """Return the number of processes along the given axis. + + For example: + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> X.get_dim('y') + 3 + """ + if axis not in self.axes: + return 0 + return self.dims[self.axes.index(axis)] + + def get_coord(self, rank): + """Return the coordinate owned by a process rank. + + The axes of the returned namedtuple can be directly accessed as members. For + example: + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> coord = X.get_coord(rank=1) + >>> coord.x + 0 + >>> coord.y + 1 + """ + for coord, idx in self.mapping.items(): + if idx == rank: + return coord + raise ValueError(f'rank {rank} not found in topology.') + + def get_axis_comm_lists(self, axis): + """ Construct lists suitable for a communicator group along axis ``axis``. + + Example: + >>> topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2]) + >>> topo.get_axis_comm_lists('pipe') + [ + [0, 4], # data=0, model=0 + [1, 5], # data=0, model=1 + [2, 6], # data=1, model=0 + [3, 7], # data=1, model=1 + ] + + Returns: + A list of lists whose coordinates match in all axes *except* ``axis``. + """ + + # We don't want to RuntimeError because it allows us to write more generalized + # code for hybrid parallelisms. + if axis not in self.axes: + return [] + + # Grab all axes but `axis` + other_axes = [a for a in self.axes if a != axis] + + lists = [] + + # Construct all combinations of coords with other_axes + ranges = [range(self.get_dim(a)) for a in other_axes] + for coord in cartesian_product(*ranges): + other_keys = {a: coord[other_axes.index(a)] for a in other_axes} + # now go over all ranks in `axis`. + sub_list = [] + for axis_key in range(self.get_dim(axis)): + key = self.ProcessCoord(**other_keys, **{axis: axis_key}) + sub_list.append(self.mapping[key]) + lists.append(sub_list) + + return lists + + def filter_match(self, **filter_kwargs): + """Return the list of ranks whose coordinates match the provided criteria. + + Example: + >>> X = ProcessTopology(axes=['pipe', 'data', 'model'], dims=[2, 2, 2]) + >>> X.filter_match(pipe=0, data=1) + [2, 3] + >>> [X.get_coord(rank) for rank in X.filter_match(pipe=0, data=1)] + [ProcessCoord(pipe=0, data=1, model=0), ProcessCoord(pipe=0, data=1, model=1)] + + Arguments: + **filter_kwargs (dict): criteria used to select coordinates. + + Returns: + The list of ranks whose coordinates match filter_kwargs. + """ + + def _filter_helper(x): + for key, val in filter_kwargs.items(): + if getattr(x, key) != val: + return False + return True + + coords = filter(_filter_helper, self.mapping.keys()) + return [self.mapping[coord] for coord in coords] + + def get_axis_list(self, axis, idx): + """Returns the list of global ranks whose coordinate in an axis is idx. + + For example: + >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) + >>> X.get_axis_list(axis='x', idx=0) + [0, 1, 2] + >>> X.get_axis_list(axis='y', idx=0) + [0, 3] + """ + + # This could be faster by generating the desired keys directly instead of + # filtering. + axis_num = self.axes.index(axis) + ranks = [self.mapping[k] for k in self.mapping.keys() if k[axis_num] == idx] + return ranks + + def world_size(self): + return len(self.mapping) + + def __str__(self): + return str(self.mapping) + + +def _prime_factors(N): + """ Returns the prime factorization of positive integer N. """ + if N <= 0: + raise ValueError("Values must be strictly positive.") + + primes = [] + while N != 1: + for candidate in range(2, N + 1): + if N % candidate == 0: + primes.append(candidate) + N //= candidate + break + return primes + + +class PipeDataParallelTopology(ProcessTopology): + """ A topology specialization for hybrid data and pipeline parallelism. + + Uses data parallelism on the last dimension to encourage gradient + reductions to use high-bandwidth intra-node links and lower-volume + pipeline communications to use low-bandwidth inter-node links. + """ + + def __init__(self, num_pp, num_dp): + super().__init__(axes=['pipe', 'data'], dims=[num_pp, num_dp]) + + +class PipeModelDataParallelTopology(ProcessTopology): + """ A topology for hybrid pipeline, model, and data parallelism. """ + + def __init__(self, num_pp, num_mp, num_dp): + super().__init__(axes=['pipe', 'data', 'model'], dims=[num_pp, num_dp, num_mp]) + + +class PipelineParallelGrid: + """Implements a grid object that stores the data parallel ranks + corresponding to each of the model parallel stages + + The grid object organizes the processes in a distributed pytorch job + into a 2D grid, of stage_id and data_parallel_id. + + self.stage_id and self.data_parallel_id stores the stage id + and the data parallel id of current process. + + self.dp_group groups the processes by stage_id. + self.dp_group[i], is a list containing all process ranks whose + stage_id is i. + + self.p2p_groups stores a list of tuple, where each tuple + stores process ranks of adjacent stages for a given data_parallel_id. + For example if num_stage is 5 then a tuple [7,8] represents stages [3, 4], + with data_parallel id = 1. A stage wrap around will appear as non-adjacent ranks, + for example tuple [4,0] with representing wrap-around stage 4 and 0, for + data_parallel_id = 0, or similarly [9,5] represents wrapped around stages [4,0] + for data_parallel_id = 1. + """ + + def __init__(self, topology=None, process_group=None): + # TODO use process_group if provided + self.global_rank = dist.get_rank() + self.world_size = dist.get_world_size() + if topology is not None: + if self.global_rank == 0: + print('Using topology:', topology) + self._topo = topology + else: + num_pp = 1 + num_dp = 1 + for idx, prime in enumerate(_prime_factors(self.world_size)): + if idx % 2 == 0: + num_pp *= prime + else: + num_dp *= prime + self._topo = PipeDataParallelTopology(num_dp=num_dp, num_pp=num_pp) + self.data_parallel_size = max(self._topo.get_dim('data'), 1) + self.pipe_parallel_size = max(self._topo.get_dim('pipe'), 1) + self.model_parallel_size = max(self._topo.get_dim('model'), 1) + self.slice_parallel_size = self.model_parallel_size + assert self._is_grid_valid(), "Invalid Grid" + + self.stage_id = self.get_stage_id() + self.data_parallel_id = self.get_data_parallel_id() + + # Create new ProcessGroups for all model parallelism. DeepSpeedLight uses these + # to detect overflow, etc. + self.ds_model_proc_group = None + self.ds_model_rank = -1 + for dp in range(self.data_parallel_size): + ranks = sorted(self._topo.get_axis_list(axis='data', idx=dp)) + if self.global_rank == 0: + #print(f'RANK={self.global_rank} building DeepSpeed model group: {ranks}') + pass + proc_group = dist.new_group(ranks=ranks) + if self.global_rank in ranks: + self.ds_model_proc_group = proc_group + self.ds_model_world_size = len(ranks) + self.ds_model_rank = ranks.index(self.global_rank) + assert self.ds_model_rank > -1 + assert self.ds_model_proc_group is not None + + # Create new ProcessGroup for gradient all-reduces - these are the data parallel groups + self.dp_group = [] + self.dp_groups = self._topo.get_axis_comm_lists('data') + for g in self.dp_groups: + proc_group = dist.new_group(ranks=g) + if self.global_rank in g: + self.dp_group = g + self.dp_proc_group = proc_group + + self.is_first_stage = (self.stage_id == 0) + self.is_last_stage = (self.stage_id == (self.pipe_parallel_size - 1)) + + self.p2p_groups = self._build_p2p_groups() + + # Create new ProcessGroup for pipeline collectives - these are pipe parallel groups + self.pp_group = [] + self.pp_proc_group = None + self.pipe_groups = self._topo.get_axis_comm_lists('pipe') + for ranks in self.pipe_groups: + if self.global_rank == 0: + #print(f'RANK={self.global_rank} building pipeline group: {ranks}') + pass + proc_group = dist.new_group(ranks=ranks) + if self.global_rank in ranks: + self.pp_group = ranks + self.pp_proc_group = proc_group + assert self.pp_proc_group is not None + + # Create new ProcessGroup for model (tensor-slicing) collectives + + # Short circuit case without model parallelism. + # TODO: it would be nice if topology had bcast semantics to avoid this branching + # case? + if self.model_parallel_size == 1: + for group_rank in range(self.world_size): + group_rank = [group_rank] + group = dist.new_group(ranks=group_rank) + if group_rank[0] == self.global_rank: + self.slice_group = group_rank + self.slice_proc_group = group + return + else: + self.mp_group = [] + self.model_groups = self._topo.get_axis_comm_lists('model') + for g in self.model_groups: + proc_group = dist.new_group(ranks=g) + if self.global_rank in g: + self.slice_group = g + self.slice_proc_group = proc_group + + def get_stage_id(self): + return self._topo.get_coord(rank=self.global_rank).pipe + + def get_data_parallel_id(self): + return self._topo.get_coord(rank=self.global_rank).data + + def _build_p2p_groups(self): + """Groups for sending and receiving activations and gradients across model + parallel stages. + """ + comm_lists = self._topo.get_axis_comm_lists('pipe') + p2p_lists = [] + for rank in range(self.world_size): + for l in comm_lists: + assert len(l) == self.pipe_parallel_size + if rank in l: + idx = l.index(rank) + buddy_rank = l[(idx + 1) % self.pipe_parallel_size] + p2p_lists.append([rank, buddy_rank]) + break # next global rank + assert len(p2p_lists) == self.world_size + return p2p_lists + + def _is_grid_valid(self): + ranks = 1 + for ax in self._topo.get_axis_names(): + ranks *= self._topo.get_dim(ax) + return ranks == dist.get_world_size() + + #returns the global rank of the process with the provided stage id + #which has the same data_parallel_id as caller process + def stage_to_global(self, stage_id, **kwargs): + me = self._topo.get_coord(self.global_rank) + transform = me._replace(pipe=stage_id, **kwargs)._asdict() + return self._topo.get_rank(**transform) + + def topology(self): + return self._topo + + # MPU functions for DeepSpeed integration + def get_global_rank(self): + return self.global_rank + + def get_pipe_parallel_rank(self): + """ The stage of the pipeline this rank resides in. """ + return self.get_stage_id() + + def get_pipe_parallel_world_size(self): + """ The number of stages in the pipeline. """ + return self.pipe_parallel_size + + def get_pipe_parallel_group(self): + """ The group of ranks within the same pipeline. """ + return self.pp_proc_group + + def get_data_parallel_rank(self): + """ Which pipeline this rank resides in. """ + return self.data_parallel_id + + def get_data_parallel_world_size(self): + """ The number of pipelines. """ + return self.data_parallel_size + + def get_data_parallel_group(self): + """ The group of ranks within the same stage of all pipelines. """ + return self.dp_proc_group + + # These are model parallel groups across all types of model parallelism. + # Deepspeed uses them to detect overflow, etc. + def get_model_parallel_rank(self): + return self.ds_model_rank + + def get_model_parallel_world_size(self): + return self.ds_model_world_size + + def get_model_parallel_group(self): + return self.ds_model_proc_group + + # For Megatron-style tensor slicing + def get_slice_parallel_rank(self): + if 'model' in self._topo.get_axis_names(): + return self._topo.get_coord(rank=self.global_rank).model + else: + return 0 + + def get_slice_parallel_world_size(self): + return self.slice_parallel_size + + def get_slice_parallel_group(self): + return self.slice_proc_group diff --git a/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3882393453512ac915a9786bbfcad349c249fb4a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .config import AUTOTP_MODE, get_tensor_parallel_config +from .tp_manager import TpTrainingManager diff --git a/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9be84138053ec4089a785f1de20cbcf59702efdc Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..532d657af8128952962052b158bf278967a38db6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..959e0e1dbe575bcdc672f189456924e307a52184 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/config.py b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/config.py new file mode 100644 index 0000000000000000000000000000000000000000..957984e9f8b339c195e94c685a9b1aa10e26cfbf --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/config.py @@ -0,0 +1,84 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import Enum +from deepspeed.runtime.config_utils import DeepSpeedConfigModel +import torch +from pydantic import Field +from typing import Optional + + +class AUTOTP_MODE(Enum): + TRAINING = "TRAINING" + INFERENCE = "INFERENCE" + + +class TPConfig(DeepSpeedConfigModel): + """ Configure tensor parallelism settings """ + + tp_size: int = 1 + """ Number of devices to split the model across using tensor parallelism. """ + + tp_grain_size: int = 1 + "The variable required by the autoTP parser has not been activated in training yet" + "as it depends on the gather logic that supports uneven partitioning. " + "Desired MLP/lm_head tp size granularity. DNN library favors tensor size in granularity of power of 2, we pick 64 as a default size." + + mpu: object = None + """ + A model parallelism unit object that implements + ``get_{model,data}_parallel_{rank,group,world_size}()``. + """ + + tp_group: object = None + + +class TPTrainingConfig(DeepSpeedConfigModel): + + dtype: torch.dtype = torch.float16 + """ + Desired model data type, will convert model to this type. + """ + + autotp_size: int = 0 + """ + In automatic tensor-parallelism training, 'tensor_parallel_size' + When set to 0, indicates that it is disabled. + """ + tp_overlap_comm: bool = False + """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """ + + tensor_parallel: TPConfig = Field({}, alias="tp") + """ + Configuration for tensor parallelism used to split the model across several + GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`. + """ + + injection_policy_tuple: Optional[tuple] = None + #The following parameters are required by autoTP parser. + ######################################## + keep_module_on_host: bool = False + """ + When loading checkpoints to model parameters, they are moved to the device. In very large models + this might fill the device and cause OOM. Setting this flag to true, will keep checkpoints on + host and not move them directly to the device (giving an option to quantize checkpoint data before + moving it to the device for example). + """ + + replace_with_kernel_inject: bool = Field(False, alias="kernel_inject") + """ + Set to true to inject inference kernels for models such as, Bert, GPT2, + GPT-Neo and GPT-J. Otherwise, the injection_dict provides the names of two + linear layers as a tuple: + `(attention_output projection, transformer output projection)` + """ + ######################################## + + +def get_tensor_parallel_config(ds_config): + + if 'tensor_parallel' in ds_config: + return TPTrainingConfig(**ds_config['tensor_parallel']) + return TPTrainingConfig() diff --git a/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..cf0b5a75c92aa3d426747917735835d1d11afe5f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from .config import TPTrainingConfig, TPConfig +from deepspeed.utils import groups +import deepspeed.comm as dist + + +class TpTrainingManager(): + + def __init__(self, model, tp_size, dtype): + self.module = model + self.config = self._initialize_config(dtype) + + from deepspeed.module_inject.auto_tp import AutoTP + from deepspeed import get_accelerator + + # Parse model configuration + parser_dict = AutoTP.tp_parser(model) + print("AutoTP: ", parser_dict) + + # Initialize TP configuration and model + self._initialize_tp_config(tp_size) + self._get_model_config_generate() + + # Synchronize random number generator state across devices + _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name()) + dist.broadcast(_rng_state, groups.get_tensor_model_parallel_src_rank(), self.tp_config.tp_group) + get_accelerator().set_rng_state(_rng_state.cpu()) + + # Apply injection policies + self._apply_policies(parser_dict) + + def _initialize_config(self, dtype): + """Initialize and return the DeepSpeed TP training configuration.""" + config = TPTrainingConfig() + config.dtype = dtype + return config + + def _apply_policies(self, parser_dict): + """Apply injection policies to the parsed modules.""" + for client_module, injection_policy in parser_dict: + self.config.injection_policy_tuple = injection_policy + self._apply_injection_policy(self.config, client_module) + + def _apply_injection_policy(self, config, client_module=None): + from deepspeed.module_inject import replace_transformer_layer + """Apply the given injection policy to a client module.""" + if isinstance(self.module, torch.nn.Module): + replace_transformer_layer(client_module, self.module, None, self.config, self.model_config) + + def _initialize_tp_config(self, tp_size): + """Perform TP configuration initialization.""" + self.tp_config = TPConfig() + self.tp_config.tp_size = tp_size + + groups._init_tp_mesh_device(tp_size) + self.tp_config.tp_group = groups.get_tensor_model_parallel_group() + self.config.tensor_parallel = self.tp_config + + def _get_model_config_generate(self): + """Generate and apply HF model configuration.""" + self.model_config = getattr(self.module, 'config', None) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..23fcf9ec13fbb5192af8c74fb8f3166fa7d622f7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .partition_parameters import ZeroParamType +from .partition_parameters import ZeroParamStatus +from .partition_parameters import Init +from .partition_parameters import GatheredParameters +from .partition_parameters import register_external_parameter + +from .tiling import TiledLinear +from .tiling import TiledLinearReturnBias + +from .mics import MiCS_Init + +from .stage3 import unwrap_model_for_generation diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f00944802fc3bf83cbd5d0d5932e061049d35d9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b14a13cf839cf8780657400724f50a1ce4e5d982 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92b08a9ecac8d07226b8c2f9a477a208933a67c7 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2668d080f5b36699033922df32cab23bc7b181ca Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..847a48953a4c266f0165c849e28ec0e22da7db08 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9796454f7d48d2e2bdf89f4b953066cc8c269bc5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60fff4d6f52e056326de02b7ceb20afefa6a845f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..757162a177c3eb6370421507beb654bc5c5d8137 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7dcec632487a20d25c1a110dfa963d5b5b9636e2 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d241577133a4b0f2f5e4bd5d01abb006358c051c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5386220fab0ffae7f47a7c2493d759c3c5fbff4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56966587888121b46887c8b69079dc10cd3df3bf Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..958c25e51bbf7601750f3ee0834e89c3eacbe43d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ced3fa5d53c929d2d8b8e7fb7c5ca65c17a8aaf Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/config.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/config.py new file mode 100644 index 0000000000000000000000000000000000000000..19ee9b51702e786773c3fd1781e3342031e805cc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/config.py @@ -0,0 +1,361 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +from typing import Optional, Dict, Any +from enum import Enum +from pydantic import Field, model_validator +from deepspeed.runtime.config_utils import get_scalar_param, pp_int, DeepSpeedConfigModel +from deepspeed.utils import logger +from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum + +# ZeRO optimization. By default, this optimization is not enabled. +# Users have to configure the desired optimization (0 means disabled) in params.json as below example: +ZERO_FORMAT = """ +ZeRO optimization should be enabled as: +"session_params": { + "zero_optimization": { + "stage": [0|1|2], + "stage3_max_live_parameters" : 1000000000, + "stage3_max_reuse_distance" : 1000000000, + "stage3_use_all_reduce_for_fetch_params": [true|false], + "stage3_module_granularity_threshold": 0, + "allgather_partitions": [true|false], + "use_multi_rank_bucket_allreduce": [true|false], + "allgather_bucket_size": 500000000, + "reduce_scatter": [true|false], + "contiguous_gradients" : [true|false] + "overlap_comm": [true|false], + "reduce_bucket_size": 500000000, + "load_from_fp32_weights": [true|false], + "cpu_offload": [true|false] (deprecated), + "cpu_offload_param" : [true|false] (deprecated), + "cpu_offload_use_pin_memory": [true|false] (deprecated), + "sub_group_size" : 1000000000000, + "offload_param": {...}, + "offload_optimizer": {...}, + "ignore_unused_parameters": [true|false], + "round_robin_gradients": [true|false], + "zero_hpz_partition_size": 1, + "zero_quantized_weights": [true|false], + "zero_quantized_nontrainable_weights": [true|false], + "zero_quantized_gradients": [true|false], + "memory_efficient_linear": [true|false], + "override_module_apply": [true|false], + "zeropp_loco_param": {...}, + "log_trace_cache_warnings" : [true|false], + } +} +""" + +ZERO_OPTIMIZATION = "zero_optimization" + + +def read_zero_config_deprecated(param_dict): + zero_config_dict = {} + zero_config_dict["stage"] = 1 if param_dict[ZERO_OPTIMIZATION] else 0 + if zero_config_dict["stage"] > 0: + zero_config_dict["allgather_bucket_size"] = get_scalar_param(param_dict, "allgather_size", 5e8) + logger.warning( + "DeepSpeedConfig: this format of ZeRO optimization setup is deprecated. Please use the following format: {}". + format(ZERO_FORMAT)) + return zero_config_dict + + +def get_zero_config(param_dict): + if ZERO_OPTIMIZATION in param_dict: + zero_config_dict = param_dict[ZERO_OPTIMIZATION] + if isinstance(zero_config_dict, bool): + zero_config_dict = read_zero_config_deprecated(param_dict) + else: + zero_config_dict = {} + return DeepSpeedZeroConfig(**zero_config_dict) + + +class ZeroStageEnum(int, Enum): + """ Enum class for possible zero stages """ + disabled = 0 + optimizer_states = 1 + gradients = 2 + weights = 3 + max_stage = 3 + + +class DeepSpeedZeroConfig(DeepSpeedConfigModel): + """ + Sets parameters for ZeRO optimizations. + """ + + stage: ZeroStageEnum = 0 + """ + Chooses different stages of ZeRO Optimizer. Stage 0, 1, 2, and 3 refer + to disabled, optimizer state partitioning, and optimizer+gradient state + partitioning, and optimizer+gradient+parameter partitioning, respectively. + """ + + contiguous_gradients: bool = True + """ + Copies the gradients to a contiguous buffer as they are produced. Avoids + memory fragmentation during backward pass. + """ + + reduce_scatter: bool = True + """ + Uses reduce or reduce scatter instead of allreduce to average gradients + """ + + reduce_bucket_size: int = Field(pp_int(5e8), ge=0) + """ + Number of elements reduced/allreduced at a time. Limits the memory required + for the allgather for large model sizes + """ + + use_multi_rank_bucket_allreduce: bool = True + """ + Combine the reduce buckets of the different ranks and do an All-Reduce instead of multiple Reduce ops. + This feature is useful when the model is small and we want to scale it on too many GPUs which therefore + reduces the message sizes of each packet. + """ + + allgather_partitions: bool = True + """ + Chooses between allgather collective or a series of broadcast collectives + to gather updated parameters from all the GPUs at the end of each step + """ + + allgather_bucket_size: int = Field(pp_int(5e8), ge=0) + """ + Number of elements allgathered at a time. Limits the memory required for + the allgather for large model sizes + """ + + overlap_comm: Optional[bool] = None # None for dynamic default value (see validator `overlap_comm_valid` below) + """ + Attempts to overlap the reduction of the gradients with backward computation + """ + + load_from_fp32_weights: bool = True + """ + Boolean indicating whether to initialize fp32 master weights from fp32 + copies in checkpoint (no precision loss) or from model's fp16 copies (with + precision loss). This can be used to initialize optimizer state even when + checkpoint is missing optimizer state. + """ + + elastic_checkpoint: bool = False + """ + Enable loading checkpoint that was saved by job with different GPU count. + No longer supported. + """ + + offload_param: Optional[DeepSpeedZeroOffloadParamConfig] = None + """ + Enable offloading of model parameters to CPU or NVMe. This frees up GPU + memory for larger models or batch sizes. Valid only with stage 3. Expects a + dictionary containing values for :any:`DeepSpeedZeroOffloadParamConfig`. + """ + + offload_optimizer: Optional[DeepSpeedZeroOffloadOptimizerConfig] = None + """ + Enable offloading of optimizer state to CPU or NVMe, and optimizer + computation to CPU. This frees up GPU memory for larger models or batch + sizes. Valid for ZeRO stage 1, 2, 3. Expects a dictionary containing values + for :any:`DeepSpeedZeroOffloadOptimizerConfig`. + """ + + sub_group_size: int = Field(pp_int(1e9), ge=0) + """ + Tile size for parameter processing to fit massive models (with trillions of + parameters). Used by ZeRO3-Offload and ZeRO-Infinity + """ + + cpu_offload_param: Optional[bool] = Field( + None, + json_schema_extra={ + "deprecated": True, + "new_param": "offload_param", + "new_param_fn": (lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu) + if val else None) + }, + ) + """ Deprecated, please use ``offload_param`` """ + + cpu_offload_use_pin_memory: Optional[bool] = Field( + None, + json_schema_extra={ + "deprecated": True, + "new_param": "offload_param or offload_optimizer", + "set_new_param": False + }, + ) + """ Deprecated, please use ``offload_param`` or ``offload_optimizer`` """ + + cpu_offload: Optional[bool] = Field( + None, + json_schema_extra={ + "deprecated": + True, + "new_param": + "offload_optimizer", + "new_param_fn": (lambda val: DeepSpeedZeroOffloadOptimizerConfig(device=OffloadDeviceEnum.cpu) + if val else None) + }, + ) + """ Deprecated, please use ``offload_optimizer`` """ + + prefetch_bucket_size: int = Field(pp_int(5e7), ge=0, alias="stage3_prefetch_bucket_size") + """ + Maximum number of parameter elements to fetch ahead of use. Used by ZeRO3, + ZeRO3-Offload, ZeRO-Infinity, and ZeRO-Inference. + """ + + param_persistence_threshold: int = Field(pp_int(1e5), ge=0, alias="stage3_param_persistence_threshold") + """ + Do not partition parameters smaller than this threshold. Smaller values use + less memory, but can greatly increase communication (especially + latency-bound messages). + """ + + model_persistence_threshold: int = Field(pp_int(sys.maxsize, "sys.maxsize"), + ge=0, + alias="stage3_model_persistence_threshold") + """ + Maximum number of parameter elements that can be persisted in GPU and not + partitioned. This imposes an upper bound on the number of unpartitioned + parameters resulting from param_persistence_threshold setting. Used by + ZeRO3-Offload, ZeRO-Infinity and ZeRO-Inference. + """ + + max_live_parameters: int = Field(pp_int(1e9), ge=0, alias="stage3_max_live_parameters") + """ + The maximum number of parameters resident per GPU before releasing. Smaller + values use less memory, but perform more communication. + """ + + max_reuse_distance: int = Field(pp_int(1e9), ge=0, alias="stage3_max_reuse_distance") + """ + Do not release a parameter if it will be reused within this threshold of + parameters. Smaller values use less memory, but perform more communication. + """ + + gather_16bit_weights_on_model_save: bool = Field(False, alias="stage3_gather_16bit_weights_on_model_save") + """ + Consolidate the weights before saving the model by ``save_16bit_model()``. + Since the weights are partitioned across GPUs, they aren’t part of + ``state_dict``, so this function automatically gathers the weights when + this option is enabled and then saves the fp16 model weights. + """ + + module_granularity_threshold: int = Field(pp_int(0), alias="stage3_module_granularity_threshold") + """ + The granularity of a module is determined by the ratio of "parameter_count / (1 + descendant count)". + ZeRO3 classifies modules with a granularity below the threshold as fine-grained, + which are treated as integral units during parameter fetching. This reduces host overhead + and the separate allgather overhead introduced by hooks for fine-grained layers when fetching parameters. + """ + + use_all_reduce_for_fetch_params: bool = Field(False, alias="stage3_use_all_reduce_for_fetch_params") + """ + Use all_reduce op when fetching module parameters at stage3. This improves performance by reducing + the overhead of concatenation and slicing on the host. + """ + + stage3_gather_fp16_weights_on_model_save: bool = Field(False, + json_schema_extra={ + "deprecated": True, + "new_param": "gather_16bit_weights_on_model_save" + }) + """ Deprecated, please use ``gather_16bit_weights_on_model_save`` """ + + ignore_unused_parameters: bool = True + """ + Unused parameters in modules may be unexpected in static networks, but + could be normal in dynamic networks. This controls whether or not training + should terminate with an error message when unused parameters are detected. + This is set to ``True`` by default, which means unused parameters are + ignored and training continues. Now is just used in stage 2. + """ + + legacy_stage1: bool = False + """ + For backward-compatibility enable old ZeRO stage 1 implementation. Use at + your own risk, will be deprecated soon. + """ + + round_robin_gradients: bool = False + """ + Stage 1 and 2 optimization for CPU offloading that parallelizes gradient + copying to CPU memory among ranks by fine-grained gradient partitioning. + Performance benefit grows with gradient accumulation steps (more copying + between optimizer steps) or GPU count (increased parallelism). + """ + zero_hpz_partition_size: int = Field(1, ge=0) + """ + Number of ranks in zero parameters partitioning secondary group + """ + zero_quantized_weights: bool = False + """ + Boolean indicating whether to quantize zero parameters (weights) + for efficient all_gather comm + """ + zero_quantized_nontrainable_weights: bool = False + """ + Boolean indicating whether to quantize non-trainable zero parameters (weights) + for efficient memory usage and communication. Different from zero_quantized_weights + that stores the weights in original precision and only perform quantization during communication, + this flag will store the weights in quantized precision. This is useful for LoRA training. + """ + zero_quantized_gradients: bool = False + """ + Boolean indicating whether to use quantized zero gradients + for efficient all_2_all_reduce comm + """ + zeropp_loco_param: Optional[Dict[str, Any]] = None + """ + This dictionary contains parameters for using LoCo-Zero++, with two key parameters: + - `err_beta`: A coefficient for the moving average of quantization errors before and after gradient computation. + It ranges between 0 and 1, with a default value of 0.8. + - `reset_T`: The number of steps after which the moving-average error buffer is cleared. The default value is 1024. + These parameters can be adjusted based on performance needs. Example configuration in ds config: + "zeropp_loco_param": { "err_beta": 0.8, "reset_T": 1024 }. + See LoCo paper for more details: (https://arxiv.org/abs/2407.04480). + """ + + mics_shard_size: int = Field(-1, json_schema_extra={"new_param": "mics_shard_size"}) + + mics_hierarchical_params_gather: bool = False + + memory_efficient_linear: bool = True + """ + Use memory efficient linear implementation, for Stage 3. + """ + """ + Whether force load checkpoint in pipeline mode, current only for Stage 3. + """ + pipeline_loading_checkpoint: bool = False + + override_module_apply: bool = True + """ + Override nn.Module apply function, for Stage 3. + """ + + log_trace_cache_warnings: bool = False + """ + Whether to log warnings from trace cache, such as invalidation events. + """ + + # Validators + @model_validator(mode="after") + def overlap_comm_valid(self): + if self.overlap_comm is None: + self.overlap_comm = self.stage == ZeroStageEnum.weights + return self + + @model_validator(mode="after") + def offload_ratio_check(self): + offload_config = self.offload_optimizer + if offload_config and offload_config.ratio < 1.0: + assert self.stage == ZeroStageEnum.weights, "Partial offloading only supported for ZeRO Stage 3." + return self diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py new file mode 100644 index 0000000000000000000000000000000000000000..35b3d5c7dd5d333e621cc554a73bb3c90e2a450e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py @@ -0,0 +1,287 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from deepspeed import comm as dist + + +def print_rank_0(message): + if dist.get_rank() == 0: + print(message) + + +class ContiguousMemoryAllocator(object): + + def __init__(self, size, dtype, device): + self.buffer = torch.zeros(size, dtype=dtype, device=device) + + #address to contiguous size available + self.contiguous_sizes = {} + + self.contiguous_sizes[0] = size + + #tensor id to its address + self.tensor_addresses = {} + + #tensor address to its size + self.tensor_sizes = {} + + #tensor address to ids + self.tensor_ids = {} + + #id to tensors + self.tensor_map = {} + + #id to params. Maps each tensor buffer to list of parameters that uses it + self.id_to_params = {} + + self.total_size = size + self.total_free = size + self.largest_contiguous = size + self.max_allocated = 0 + + self.count = 0 + + #create a tensor of size from the pre-allocated buffer + #if not enough free space will fail + #if not enough contiguous space, will defragment and allocate + def allocate_tensor(self, size): + free_before = self.total_free + + assert size <= self.total_free, "Not enough memory in buffer. Allocation failed" + if self.largest_contiguous < size: + print_rank_0("Needs defragmentation to allocate. Before Defragmentation:") + self.print_allocation(resolution=100) + self._defragment_memory() + #set the param data to the new tensor buffer locations + self._reset_param_data() + print_rank_0("After defragmentation:") + self.print_allocation(resolution=100) + + self.total_free = self.total_free - size + + allocated = self.total_size - self.total_free + if allocated > self.max_allocated: + self.max_allocated = allocated + + tensor_address = self._get_new_tensor_address(size) + + ret_tensor = self._get_new_tensor(tensor_address, size) + print_rank_0( + f"Free before allocation {free_before}. Allocating {size}. Free after allocation {self.total_free}. Max allocated {self.max_allocated}" + ) + assert self.total_free + size == free_before, "Allocation bookkeeping error" + + return ret_tensor + + #assigns the tensor data to the param data and keeps track of the assignment + #any change the underlying buffer from defragmentation will cause a + #reassignment of the param data + def assign_to_param(self, tensor, param, numel, shape): + tensor_id = id(tensor) + + assert tensor_id in self.tensor_map.keys(), "No such tensor allocated by the allocator." + assert tensor.numel() >= numel, "Assert tensor buffer does is not large enough" + assert not tensor_id in self.id_to_params.keys(), "This tensor has already been assigned to a param" + + self.id_to_params[tensor_id] = [param] + + replicated_tensor = tensor.narrow(0, 0, numel).view(shape) + param.data = replicated_tensor.data + param.contiguous_tensor_id = tensor_id + + #deletes the tensor and frees up the underlying buffer + def release_tensor(self, tensor): + free_before = self.total_free + tensor_id = id(tensor) + tensor_size = tensor.numel() + self._release_tensor(tensor_id) + self._unassign_params(tensor_id) + self.total_free += tensor_size + print_rank_0( + f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.") + assert self.total_free - tensor_size == free_before, "Release bookkeeping error" + + def release_tensor_with_id(self, tensor_id): + free_before = self.total_free + assert tensor_id in self.tensor_map.keys(), "Invalid tensor id" + tensor = self.tensor_map[tensor_id] + tensor_size = tensor.numel() + self._release_tensor(tensor_id) + self._unassign_params(tensor_id) + self.total_free += tensor_size + print_rank_0( + f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.") + assert self.total_free - tensor_size == free_before, "Release bookkeeping error" + + #shows the current memory allocation at specified resolution + def print_allocation(self, resolution=200): + total_size = self.buffer.numel() * 1.0 + empty = [] + for addr, size in self.contiguous_sizes.items(): + start = int(addr * resolution / total_size) + end = int((addr + size) * resolution / total_size) + empty.extend(range(start, end)) + s = '' + for i in range(resolution): + s += '.' if i in empty else '|' + print_rank_0(s) + + def max_allocated(self): + return self.max_allocated + + #to be called after defragmentation that moves the tensor buffers + #this call reassigns the data of all the parameters using the tensor buffers + def _reset_param_data(self): + for id, tensor in self.tensor_map.items(): + for param in self.id_to_params[id]: + param.data = tensor.narrow(0, 0, param.numel()).view(param.data.shape).data + + def _unassign_params(self, tensor_id): + if tensor_id in self.id_to_params.keys(): + del self.id_to_params[tensor_id] + + def _release_tensor(self, tensor_id): + assert tensor_id in self.tensor_addresses, f"Tensor id {tensor_id} not found" + + address = self.tensor_addresses[tensor_id] + contiguous_size = self.tensor_map[tensor_id].numel() + + del self.tensor_addresses[tensor_id] + del self.tensor_ids[address] + del self.tensor_map[tensor_id] + del self.tensor_sizes[address] + + self._consolidate_address(address, contiguous_size) + self.largest_contiguous = self._largest_contiguous() + + def _consolidate_address(self, address, contiguous_size): + + #consolidate next buffer + end_address = address + contiguous_size + if end_address in self.contiguous_sizes: + contiguous_size += self.contiguous_sizes[end_address] + del self.contiguous_sizes[end_address] + + #consolidate previous buffer + for addr, size in self.contiguous_sizes.items(): + if addr + size == address: + del self.contiguous_sizes[addr] + contiguous_size += size + address = addr + break + + self.contiguous_sizes[address] = contiguous_size + + def _defragment_memory(self): + empty_addresses = sorted(self.contiguous_sizes.keys()) + tensor_addresses = sorted(self.tensor_addresses.values()) + + tensor_index = 0 + + while tensor_index < len(tensor_addresses): + + empty_addr = empty_addresses[0] + empty_size = self.contiguous_sizes[empty_addr] + + tensor_addr = tensor_addresses[tensor_index] + tensor_size = self.tensor_sizes[tensor_addr] + tensor_id = self.tensor_ids[tensor_addr] + tensor = self.tensor_map[self.tensor_ids[tensor_addr]] + + assert tensor_size == tensor.numel(), \ + f"Size mismatch. {tensor_size} is allocated at addr {tensor_addr} but tensor size is {tensor.numel()} " + + assert empty_addr != tensor_addr, \ + f"Cannot have same empty address {empty_addr} and tensor address {tensor_addr}" + + if empty_addr < tensor_addr: + + if empty_size >= tensor_size: + dest_buffer = self.buffer.narrow(0, empty_addr, tensor_size) + src_buffer = self.buffer.narrow(0, tensor_addr, tensor_size) + dest_buffer.data.copy_(src_buffer.data) + else: + + #print_rank_0(f'empty addr : {empty_addr}, empty size {empty_size} tensor addr {tensor_addr} tensor size {tensor_size}') + src_addr = tensor_addr + dest_addr = empty_addr + while src_addr < (tensor_addr + tensor_size): + copy_size = min(empty_size, tensor_addr + tensor_size - src_addr) + + dest_buffer = self.buffer.narrow(0, dest_addr, copy_size) + src_buffer = self.buffer.narrow(0, src_addr, copy_size) + + dest_buffer.data.copy_(src_buffer.data) + + src_addr += copy_size + dest_addr += copy_size + + self._replace_old_address_with_new(tensor_id, empty_addr) + + tensor_index += 1 + + else: + tensor_index += 1 + + empty_addresses = sorted(self.contiguous_sizes.keys()) + + def _replace_old_address_with_new(self, tensor_id, new_address): + + tensor = self.tensor_map[tensor_id] + tensor_size = tensor.numel() + tensor.data = self.buffer.narrow(0, new_address, tensor_size).data + + self._release_tensor(tensor_id) + self._mark_as_occupied(new_address, tensor_size) + + self.tensor_ids[new_address] = tensor_id + self.tensor_map[tensor_id] = tensor + self.tensor_addresses[tensor_id] = new_address + self.tensor_sizes[new_address] = tensor_size + + def _get_new_tensor_address(self, size): + tensor_address = None + for address, contiguous_size in self.contiguous_sizes.items(): + if contiguous_size >= size and \ + (tensor_address is None or \ + contiguous_size < self.contiguous_sizes[tensor_address]): + tensor_address = address + assert tensor_address is not None, "address cannot be None" + return tensor_address + + def _get_new_tensor(self, address, size): + available_contiguous_size = self.contiguous_sizes[address] + + assert size <= available_contiguous_size, \ + f"Tensor numel {size} is large than available contiguous size {available_contiguous_size}" + self.count += 1 + new_tensor = self.buffer.narrow(0, address, size) + tensor_id = id(new_tensor) + self.tensor_addresses[tensor_id] = address + self.tensor_sizes[address] = size + + self.tensor_ids[address] = tensor_id + self.tensor_map[tensor_id] = new_tensor + + self._mark_as_occupied(address, size) + + return new_tensor + + def _largest_contiguous(self): + if len(self.contiguous_sizes) > 0: + return max([size for _, size in self.contiguous_sizes.items()]) + else: + return 0 + + def _mark_as_occupied(self, address, size): + available_contiguous_size = self.contiguous_sizes[address] + del self.contiguous_sizes[address] + + if available_contiguous_size != size: + self.contiguous_sizes[address + size] = available_contiguous_size - size + + self.largest_contiguous = self._largest_contiguous() diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/linear.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..8c8db60768ebca9f7f972c13e0d3d20176158c96 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/linear.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +#Linear Module to use with ZeRO Stage 3 to allow for parameter memory release +#after the module execution during forward +#Instead of saving variables using save_for_backward, we save variable ids +#Allowing us to retrieve the variable without creating pointer to it +#Which allows for underlying tensor to be garbage collected +#When partitioned as needed by the Zero Stage 3 optimizer +#TODO instead of patching Linear module, we could patch the ctx.save_for_backward +#ctx.saved_tensors so that this approach works for all nn modules that are built upon +#torch.nn.function. However the issue is that many modules uses C++ implementations +#which does not have pytorch implementation. Eg torch.addmm which acts as a functional +#when implemented outside of torch.autograd.Function + +import math +import functools + +import torch +from torch import Tensor +from torch.nn.parameter import Parameter +from torch.nn import init +from torch.nn.modules.module import Module +from deepspeed.runtime.utils import noop_decorator +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator + + +def print_rank_0(message, debug=False, force=False): + if dist.get_rank() == 0 and (debug or force): + print(message) + + +try: + # Fix `torch.[device].amp.custom_fwd/bwd` FutureWarning in torch 2.4 + if hasattr(torch, 'amp') and hasattr(torch.amp, 'custom_fwd') and hasattr(torch.amp, 'custom_bwd'): + autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=get_accelerator().device_name()) + autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=get_accelerator().device_name()) + else: + # original implementation + autocast_custom_fwd = get_accelerator().amp().custom_fwd + autocast_custom_bwd = get_accelerator().amp().custom_bwd +except (ImportError, AttributeError) as exp: + autocast_custom_fwd = noop_decorator + autocast_custom_bwd = noop_decorator + + +class LinearFunctionForZeroStage3(torch.autograd.Function): + + # Note that both forward and backward are @staticmethods + @staticmethod + @autocast_custom_fwd + # bias is an optional argument + def forward(ctx, input, weight, bias=None): + + ctx.save_for_backward(input, weight, bias) + + if input.dim() == 2 and bias is not None: + # fused op is marginally faster + ret = torch.addmm(bias, input, weight.t()) + else: + output = input.matmul(weight.t()) + if bias is not None: + output += bias + ret = output + + return ret + + # This function has only a single output, so it gets only one gradient + @staticmethod + @autocast_custom_bwd + def backward(ctx, grad_output): + # This is a pattern that is very convenient - at the top of backward + # unpack saved_tensors and initialize all gradients w.r.t. inputs to + # None. Thanks to the fact that additional trailing Nones are + # ignored, the return statement is simple even when the function has + # optional inputs. + input, weight, bias = ctx.saved_tensors + + grad_input = grad_weight = grad_bias = None + + #print(f"backward shaped grad_output {grad_output.shape}, input {input.shape}, weight {weight.shape} and bias {bias.shape if bias is not None else None}") + # These needs_input_grad checks are optional and there only to + # improve efficiency. If you want to make your code simpler, you can + # skip them. Returning gradients for inputs that don't require it is + # not an error. + if ctx.needs_input_grad[0]: + #print(f"Computing grad input weight {weight.shape} grad_output {grad_output.shape}") + grad_input = grad_output.matmul(weight) + #print(f"Computed grad input {grad_input.shape}") + if ctx.needs_input_grad[1]: + #print("Computing grad weight") + dim = grad_output.dim() + if dim > 2: + grad_weight = grad_output.reshape(-1, + grad_output.shape[-1]).t().matmul(input.reshape(-1, input.shape[-1])) + else: + grad_weight = grad_output.t().matmul(input) + #print(f"Computed grad weight grad_weight {grad_weight.shape}") + if bias is not None and ctx.needs_input_grad[2]: + #print("Computing grad bias") + if dim > 2: + grad_bias = grad_output.sum([i for i in range(dim - 1)]) + else: + grad_bias = grad_output.sum(0) + #print("Done computing grad bias") + #print("needs bias") + #print(f"backward shaped grad_input {grad_input.shape}, grad_weight {grad_weight.shape}, grad_bias {grad_bias.shape if grad_bias is not None else None}") + return grad_input, grad_weight, grad_bias + + +def zero3_linear_wrap(input, weight, bias=None): + if bias is None: + return LinearFunctionForZeroStage3.apply(input, weight) + else: + return LinearFunctionForZeroStage3.apply(input, weight, bias) + + +class LinearModuleForZeroStage3(Module): + r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`. + The weights are pre-transposed and stored as A^T instead of transposing during each + forward. Memory savings proportional to the parameter size. + + Args: + in_features: size of each input sample + out_features: size of each output sample + bias: If set to ``False``, the layer will not learn an additive bias. + Default: ``True`` + + Shape: + - Input: :math:`(N, *, H_{in})` where :math:`*` means any number of + additional dimensions and :math:`H_{in} = \text{in\_features}` + - Output: :math:`(N, *, H_{out})` where all but the last dimension + are the same shape as the input and :math:`H_{out} = \text{out\_features}`. + + Attributes: + weight: the learnable weights of the module of shape + :math:`(\text{out\_features}, \text{in\_features})`. The values are + initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where + :math:`k = \frac{1}{\text{in\_features}}` + bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. + If :attr:`bias` is ``True``, the values are initialized from + :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where + :math:`k = \frac{1}{\text{in\_features}}` + + Examples:: + + >>> m = nn.Linear(20, 30) + >>> input = torch.randn(128, 20) + >>> output = m(input) + >>> print(output.size()) + torch.Size([128, 30]) + """ + __constants__ = ['in_features', 'out_features'] + in_features: int + out_features: int + weight: Tensor + + def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: + super(LinearModuleForZeroStage3, self).__init__() + print("Building ZeRO module") + self.in_features = in_features + self.out_features = out_features + self.weight = Parameter(torch.Tensor(out_features, in_features)) + if bias: + self.bias = Parameter(torch.Tensor(out_features)) + else: + self.register_parameter('bias', None) + self.reset_parameters() + + def reset_parameters(self) -> None: + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) + bound = 1 / math.sqrt(fan_in) + init.uniform_(self.bias, -bound, bound) + + def forward(self, input: Tensor) -> Tensor: + return LinearFunctionForZeroStage3.apply(input, self.weight, self.bias) + + def extra_repr(self) -> str: + return 'in_features={}, out_features={}, bias={}'.format(self.in_features, self.out_features, self.bias + is not None) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/mics.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/mics.py new file mode 100644 index 0000000000000000000000000000000000000000..628bf86a61da23182b2f33cc7b6edab2c1aee41f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/mics.py @@ -0,0 +1,472 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sys +from typing import List + +import deepspeed +import torch +from deepspeed import comm as dist +from deepspeed.runtime.zero.utils import is_zero_param +from deepspeed.runtime.zero.mics_utils import (MiCS_CommGroups, create_mics_comm_groups, scale_tensors) +from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload +from deepspeed.runtime.zero.partition_parameters import Init, AllGatherCoalescedHandle, ZeroParamStatus +from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 +from deepspeed.utils import instrument_w_nvtx, log_dist, logger +from deepspeed.accelerator import get_accelerator +from torch import Tensor +from torch.nn import Parameter + + +def has_hierarchical_all_gather_groups(comm_groups: MiCS_CommGroups): + result = False + if comm_groups.param_intra_node_group is not None and comm_groups.param_inter_node_shard_group is not None: + result = True + return result + + +class MiCS_AllGatherCoalescedHandle(AllGatherCoalescedHandle): + """ This handle assumes that no need to + copy data out from a contiguous tensor + """ + + def __init__(self, allgather_handle, params: List[Parameter], partitions: List[Tensor], world_size: int) -> None: + super().__init__(allgather_handle, params, partitions, world_size) + + def wait(self, **kwargs) -> None: + """ + """ + # let the current stream to op + try: + # print("HANDLE", self.allgather_handle) + instrument_w_nvtx(self.allgather_handle.wait)() + except (ValueError, RuntimeError) as e: + log_dist( + f"WARNING: Runtime Error while waiting the collective all-gather, possibly due to the _IllegalWork", + ranks=[0]) + log_dist(f"Error message: {e}", ranks=[0]) + + if self.complete: + return + + for _, param in enumerate(self.params): + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + param.ds_status = ZeroParamStatus.AVAILABLE + + self.complete = True + + +class MiCS_Init(Init): + + def __init__(self, + module=None, + data_parallel_group=None, + sequence_data_parallel_group=None, + mem_efficient_linear=True, + remote_device=None, + pin_memory=False, + config_dict_or_path=None, + config=None, + enabled=True, + dtype=None, + mpu=None): + """A context manager to partition the model parameters during the model + construction with MiCS partition strategy. Model states are partitioned + to the number of devices specified via ``mics_shard_size`` field in the + deepspeed config json file. The context manager also introduces + hierarchical communication method to reduce the cost of inter-node + communications, which can be enabled with + ``mics_hierarchical_params_gather`` field in deepspeed config. + + Args: + module (``torch.nn.Module``, optional): If provided, partition the model as + if it was constructed in the context. + data_parallel_group (``deepspeed.comm`` process group, optional): + The group of processes to partition among. Defaults to all processes. + Synonymous with sequence data parallel group for param partitioning + across both sequence and data parallel groups. + mem_efficient_linear (bool, optional): Replace + torch.nn.functional.linear with an implementation that allows + DeepSpeed to partition parameters. Defaults to ``True``. + remote_device (string, optional): The initial device to store model + weights e.g., ``cpu``, ``nvme``. Passing ``"cpu"`` will create the model in CPU + memory. The model may still be moved to GPU based on the + offload settings for training. Defaults to param offload device if a config is + defined, otherwise GPU. + pin_memory (bool, optional): Potentially increase performance by + using pinned memory for model weights. ``remote_device`` must be + ``"cpu"``. Defaults to pin_memory value in config, otherwise ``False``. + config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration + for swapping fp16 params to NVMe. + config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead. + enabled (bool, optional): If ``False``, this context has no + effect. Defaults to ``True``. + dtype (``dtype``, optional): Can be used to change the data type of the parameters. + Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None`` + mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,world_size}. + + This context follows the same logic as ``deepspeed.zero.Init()``, but + with the modification for partition size of each parameter. + + Examples + -------- + + #. Allocate a model and partition it among all processes: + + .. code-block:: python + # the config_dict_or_path is required to let the context manager know + # how partition the parameters. + # The configuration has to include the field ``mics_shard_size`` + with deepspeed.zero.MiCS_Init(config_dict_or_path=ds_config): + model = MyLargeModel() + + + #. Allocate a model in pinned CPU memory and partition it among a subgroup of processes: + + .. code-block:: python + + with deepspeed.zero.MiCS_Init(data_parallel_group=mpu.get_data_parallel_group(), + remote_device="cpu", + pin_memory=True + config_dict_or_path=ds_config): + model = MyLargeModel() + + + #. Partition an already-allocated model in CPU memory: + + .. code-block:: python + + model = deepspeed.zero.MiCS_Init(module=model, + config_dict_or_path=ds_config) + """ + + assert config_dict_or_path is not None, "Must provide configuration for MiCS Initialization" + _ds_config = deepspeed.runtime.config.DeepSpeedConfig(config_dict_or_path, mpu) + if not dist.is_initialized(): + dist.init_distributed() + assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm" + + if data_parallel_group is None: + ds_process_group = dist.get_world_group() + else: + ds_process_group = data_parallel_group + + if sequence_data_parallel_group is not None: + logger.warning( + f"sequence_data_parallel_group' is deprecated and will be removed. Use 'data_parallel_group' instead.") + if data_parallel_group is not None: + raise ValueError( + "Both 'data_parallel_group' and 'sequence_data_parallel_group' were specified. Please provide only one of these arguments." + ) + self.ds_process_group = sequence_data_parallel_group + + self.mics_comm_groups = create_mics_comm_groups( + _ds_config.mics_shard_size, + ds_process_group, + hierarchical_allgather=_ds_config.mics_hierarchial_params_gather, + mpu=mpu) + + super().__init__(module, data_parallel_group, mem_efficient_linear, remote_device, pin_memory, + config_dict_or_path, config, enabled, dtype, mpu) + + def _convert_to_deepspeed_param(self, param): + super()._convert_to_deepspeed_param(param) + # attach communication groups to every param + param.comm = self.mics_comm_groups + + # record existing all_gather_coalesced implementation + # so that we can fallback later + old_all_gather_coalesced = param.all_gather_coalesced + + def _param_all_gather_coalesced(params, param_buffers=None, **kwargs): + """""" + mics_comm_groups: MiCS_CommGroups = params[0].comm + hierarchical_all_gather = has_hierarchical_all_gather_groups(mics_comm_groups) + if dist.has_coalescing_manager() and hierarchical_all_gather: + return self._hierarchical_all_gather_params(params, param_buffers) + elif dist.has_coalescing_manager(): + return self._flat_all_gather_with_coalescing_manager(params, param_buffers) + else: + return old_all_gather_coalesced(params, **kwargs) + + # change the all_gather_coalesced method + param.all_gather_coalesced = _param_all_gather_coalesced + + def _pre_all_gather(self, params, params_buffers=None): + # fetches from nvme if the partition is not available and in nvme + self._ensure_availability_of_partitioned_params(params) + + for param in params: + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(param.ds_summary()) + param.ds_status = ZeroParamStatus.INFLIGHT + + # ensure that each rank has params in same order. the allgather + # is done by flattening the parameter list into a single tensor that + # can be allgathered in a single call - this means that if each rank + # gives a list of the same parameters in a different order we will + # silently get incorrect parameter values, and have very difficult + # to debug correctness issues. + params = sorted(params, key=lambda p: p.ds_id) + return params, params_buffers + + def _flat_all_gather_with_coalescing_manager(self, params, params_buffers=None): + """""" + # must have to change the status of the param + # and ensure they are on the device + params, params_buffers = self._pre_all_gather(params, params_buffers) + + mics_comm_groups: MiCS_CommGroups = params[0].comm + param_shard_size = mics_comm_groups.param_shard_size + + output_tensors = [] + input_tensors = [] + for i, p in enumerate(params): + t_size = p.ds_tensor.ds_numel * param_shard_size + if params_buffers is not None and params_buffers[i] is not None: + assert params_buffers[i].numel( + ) == t_size, f'params_to_gather_buffers[{i}] size {params_buffers[i].numel()} does not match with t_size {t_size}' + flat_out = params_buffers[i] + else: + flat_out = torch.empty(t_size, dtype=p.dtype, device=self.local_device, requires_grad=False).view(-1) + output_tensors.append(flat_out) + _flat_input = p.ds_tensor.data.view(-1) + input_tensors.append(_flat_input) + + all_gather_handle = dist.all_gather_coalesced(output_tensors, + input_tensors, + group=mics_comm_groups.param_shard_group, + async_op=True) + + for idx, param in enumerate(params): + param.data = output_tensors[idx].narrow(0, 0, param.ds_numel).view(param.ds_shape).data + + return MiCS_AllGatherCoalescedHandle(allgather_handle=all_gather_handle, + params=params, + partitions=[], + world_size=param_shard_size) + + def _hierarchical_all_gather_params(self, params, params_buffers=None): + """""" + params, params_buffers = self._pre_all_gather(params, params_buffers) + + mics_comm_groups: MiCS_CommGroups = params[0].comm + local_rank = dist.get_rank(group=mics_comm_groups.param_intra_node_group) + inter_node_comm_group = mics_comm_groups.param_inter_node_shard_group + intra_node_comm_group = mics_comm_groups.param_intra_node_group + param_shard_size = mics_comm_groups.param_shard_size + + inter_node_size = dist.get_world_size(group=inter_node_comm_group) + intra_node_size = dist.get_world_size(group=intra_node_comm_group) + param_tensors = [] + for i, p in enumerate(params): + param_size = p.ds_tensor.ds_numel * param_shard_size + if params_buffers is not None and params_buffers[i] is not None: + assert params_buffers[i].numel( + ) == param_size, f'param_buffers[{i}] size {params_buffers[i].numel()} does not match with param_size {param_size}' + param_tensor = params_buffers[i] + else: + param_tensor = torch.empty(param_size, dtype=p.dtype, device=self.local_device, + requires_grad=False).view(-1) + param_tensors.append(param_tensor) + + # inter node all-gather + inter_outputs = [] + inter_inputs = [] + for i, p in enumerate(params): + inter_size = p.ds_tensor.ds_numel * inter_node_size + _out = param_tensors[i].narrow(0, local_rank * inter_size, inter_size) + inter_outputs.append(_out) + inter_inputs.append(p.ds_tensor.data.view(-1).to(self.local_device)) + # sync enqueue + dist.all_gather_coalesced(inter_outputs, inter_inputs, group=inter_node_comm_group, async_op=False) + + # intra node all-gather + intra_outputs = [] + intra_inputs = [] + for i, p in enumerate(params): + # partition param into multiple chunks for allgather + # because inter-node all-gather outputs are in a continues memory + # while in param memory, those inter-node data are placed in different + # location. + # each chunk is an intra-node output + param_chunk = param_tensors[i].view( + (inter_node_size, intra_node_size, p.ds_tensor.ds_numel)).narrow(1, local_rank, 1) + param_chunk.copy_(inter_outputs[i].detach().clone().view(param_chunk.size())) + output_chunks = torch.chunk(param_tensors[i], inter_node_size) + for j, _out in enumerate(output_chunks): + intra_chunk_size = intra_node_size * p.ds_tensor.ds_numel + local_offset = local_rank * p.ds_tensor.ds_numel + _in = param_tensors[i].narrow(0, j * intra_chunk_size + local_offset, p.ds_tensor.ds_numel) + intra_outputs.append(_out) + intra_inputs.append(_in) + + all_gather_handle = dist.all_gather_coalesced(intra_outputs, + intra_inputs, + group=intra_node_comm_group, + async_op=True) + for i, param in enumerate(params): + param.data = param_tensors[i].narrow(0, 0, param.ds_numel).view(param.ds_shape).data + + return MiCS_AllGatherCoalescedHandle( + allgather_handle=all_gather_handle, + params=params, + partitions=[], + world_size=param_shard_size, + ) + + def get_partition_dp_group(self, param): + return param.comm.param_shard_group + + def get_partition_rank(self): + return self.mics_comm_groups.param_shard_rank + + @property + def num_partitions(self): + return self.mics_comm_groups.param_shard_size + + +class MiCS_Offload(DeepSpeedZeRoOffload): + """ Wrapper to change the behavior for parameter sharding + """ + + def _convert_to_zero_parameters(self, ds_config, module, mpu): + """ overload the parent class function for convert the parameters + + """ + log_dist(f'Convert to zero parameters from MiCS Offload manager', ranks=[0]) + non_zero_params = [p for p in module.parameters() if not is_zero_param(p)] + if non_zero_params: + zero_params = [p for p in module.parameters() if is_zero_param(p)] + if zero_params: + zero_params[0].convert_to_zero_parameters(param_list=non_zero_params) + else: + group = None + if mpu: + group = mpu.get_data_parallel_group() + + MiCS_Init(module=module, + data_parallel_group=group, + dtype=self.dtype, + config_dict_or_path=ds_config, + remote_device=self.offload_device, + pin_memory=self.offload_param_pin_memory, + mpu=mpu) + + +class MiCS_Optimizer(DeepSpeedZeroOptimizer_Stage3): + """ + MiCS Optimizer + """ + + def __init__(self, + module, + init_optimizer, + timers, + ds_config, + static_loss_scale=1, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + contiguous_gradients=True, + reduce_bucket_size=500000000, + prefetch_bucket_size=50000000, + max_reuse_distance=1000000000, + max_live_parameters=1000000000, + param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, + dp_process_group=None, + reduce_scatter=True, + overlap_comm=False, + offload_optimizer_config=None, + offload_param_config=None, + sub_group_size=1000000000000, + offload_ratio=0.0, + mpu=None, + clip_grad=0, + gradient_accumulation_dtype=torch.float16, + communication_data_type=torch.float16, + postscale_gradients=True, + gradient_predivide_factor=1, + gradient_accumulation_steps=1, + elastic_checkpoint=False, + aio_config=None): + + log_dist("Init MiCS optimizer", ranks=[0]) + super().__init__(module, init_optimizer, timers, ds_config, static_loss_scale, dynamic_loss_scale, + dynamic_loss_args, verbose, contiguous_gradients, reduce_bucket_size, prefetch_bucket_size, + max_reuse_distance, max_live_parameters, param_persistence_threshold, + model_persistence_threshold, dp_process_group, reduce_scatter, overlap_comm, + offload_optimizer_config, offload_param_config, sub_group_size, offload_ratio, mpu, clip_grad, + gradient_accumulation_dtype, communication_data_type, postscale_gradients, + gradient_predivide_factor, gradient_accumulation_steps, elastic_checkpoint, aio_config) + first_param = next(module.parameters()) + # overload the dp_process_group and partition_count + assert hasattr(first_param, "comm"), " ".join([ + "Sharded parameters don't have the MiCS_CommGroups attached.", + "Might due to the use of deepspeed.zero.Init context for initializing the weights.", + "To use MiCS sharding, please use deepspeed.zero.MiCS_Init instead for initializing parameter." + ]) + self.dp_process_group = first_param.comm.param_shard_group + self.partition_count = first_param.comm.param_shard_size + + def initialize_ds_offload( + self, + *args, + **kwargs, + ): + return MiCS_Offload(*args, **kwargs) + + def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: + grad_buffers = super().partition_grads(params_to_release, grad_partitions) + # perform all-reduce among replication groups + # the function will perform accumulation boundary check + self.allreduce_mics_shard_grads(params_to_release, grad_buffers) + + @instrument_w_nvtx + def allreduce_mics_shard_grads(self, params, partitioned_grads_buffers: List[Tensor]): + """ + """ + # TODO: improve the condition check + if not self.is_gradient_accumulation_boundary or \ + len(partitioned_grads_buffers) == 0: + return + + mics_comm_groups: MiCS_CommGroups = params[0].comm + param_repli_group = mics_comm_groups.param_repli_group + param_repli_size = mics_comm_groups.param_repli_size + + if param_repli_size is None or param_repli_size <= 1: + return + if not get_accelerator().on_accelerator(partitioned_grads_buffers[0]): + raise RuntimeError("Local sharding has no support for CPU offloading") + + if dist.has_all_reduce_coalesced(): + scale_tensors(partitioned_grads_buffers, param_repli_size) + dist.all_reduce_coalesced(tensors=partitioned_grads_buffers, group=param_repli_group) + else: + # manually coalescing all-reduce + aggregated_buffer: Tensor = torch.cat(partitioned_grads_buffers) + aggregated_buffer.div_(param_repli_size) + dist.all_reduce(aggregated_buffer, group=param_repli_group) + offset = 0 + for grad_buff in partitioned_grads_buffers: + grad_buff.view(-1).copy_(aggregated_buffer.narrow(0, offset, grad_buff.numel())) + offset += grad_buff.numel() + + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None): + r""" Loading the ZeRO-3/MiCS partitioned checkpoints + Because the self.dp_process_group is replaced with the communicator for + partition group we can call the load_state_dict logic from ZeRO-3. + """ + super().load_state_dict(state_dict_list, load_optimizer_states, load_from_fp32_weights, checkpoint_folder) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/mics_utils.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/mics_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..06b83160bd6c4cd93f5250a12c67f5abda01d920 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/mics_utils.py @@ -0,0 +1,203 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from dataclasses import dataclass +from typing import List + +import numpy as np +import torch +from torch import Tensor + +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.utils import logger + + +def _log_rank0(msg): + if dist.get_rank() == 0: + logger.info(msg) + + +@torch.jit.script +def scale_tensors(tensors: List[Tensor], scale: int): + for t in tensors: + t.div_(scale) + + +@dataclass +class MiCS_CommGroups: + """""" + param_shard_group = None + param_shard_size = -1 + param_shard_rank = -1 + + param_repli_group = None + param_repli_size = -1 + param_repli_rank = -1 + + param_intra_node_group = None + param_inter_node_shard_group = None + + +def create_mics_comm_groups( + shard_size, + dp_group, + hierarchical_allgather=False, + mpu=None, +): + """ + create shard-group, replicate-group from config_file + TODO: consider broadcast the config from rank0 + + Returns: + MiCS_CommGroups + """ + # env var for debugging purpose + ndevices_per_node = int(os.environ.get("NDEV_PER_NODE", get_accelerator().device_count())) + _log_rank0(f'creating MiCS communication groups with per node device size {ndevices_per_node}') + groups = MiCS_CommGroups() + + if mpu is not None: + assert dp_group == mpu.get_data_parallel_group() + + # full size of the world + world_size = dist.get_world_size() + # global rank + global_rank = dist.get_rank() + + config = _generate_mics_config(world_size, ndevices_per_node, shard_size, 1) + ranks_of_shard_group = config['shard_groups'] + ranks_of_repli_group = config['replicate_groups'] + if len(ranks_of_repli_group) == 0: + assert len(ranks_of_shard_group) == 1, "replicate groups are empty only for single shard group" + for r in ranks_of_shard_group[0]: + ranks_of_repli_group.append([r]) + + # for simplicity + assert _sizes_all_same(ranks_of_repli_group), "replicate groups must have the same size" + assert _sizes_all_same(ranks_of_shard_group), "shard groups must have the same size" + + assert sum([len(g) for g in ranks_of_shard_group]) == dist.get_world_size(), "all sharded ranks " + if len(ranks_of_shard_group) > 1: # if only shard on one group then no need for replicate groups + assert len(ranks_of_shard_group) == len( + ranks_of_repli_group[0]), "number of shard groups must equal to the size of each replicate group" + + global_rank = dist.get_rank() + # create shard groups + for shard_ranks in ranks_of_shard_group: + _group = dist.new_group(shard_ranks) + if global_rank in shard_ranks: + groups.param_shard_group = _group + groups.param_shard_size = len(shard_ranks) + groups.param_shard_rank = dist.get_rank(_group) + logger.info(f'rank {global_rank}, shard group' + f' {groups.param_shard_rank}/{dist.get_world_size(group=_group)}') + + # create replicate groups + for repli_ranks in ranks_of_repli_group: + if len(repli_ranks) > 1: + _group = dist.new_group(repli_ranks) + if global_rank in repli_ranks: + groups.param_repli_group = _group + groups.param_repli_size = len(repli_ranks) + groups.param_repli_rank = dist.get_rank(group=_group) + logger.info(f'rank {global_rank} ' + f'replicate group {groups.param_repli_rank}/{dist.get_world_size(group=_group)}') + else: + groups.param_repli_group = None + groups.param_repli_size = 1 + groups.param_repli_rank = 0 + logger.info(f'rank {global_rank} replicate group 0/1') + + # assign shard group size as world size + assert groups.param_shard_size == len(ranks_of_shard_group[0]) + + if hierarchical_allgather: + # create hierarchy inter-node, intra-node groups + # n_span_nodes = config['shard_span'] + n_span_nodes = config['span_nodes'] + assert n_span_nodes > 1, "sharding spans on single node, no need for hierarchy allgather" + assert len(ranks_of_shard_group[0]) % n_span_nodes == 0 + + n_gpu_per_node = len(ranks_of_shard_group[0]) // n_span_nodes + intra_node_ranks_group = [] + inter_node_ranks_group = [] + for shard_group in ranks_of_shard_group: + _intra_node_ranks = [] + for i in range(0, len(shard_group), n_gpu_per_node): + _intra_node_ranks.append(shard_group[i:i + n_gpu_per_node]) + _inter_node_ranks = [] + for i in range(n_gpu_per_node): + _ranks = [_g[i] for _g in _intra_node_ranks] + _inter_node_ranks.append(_ranks) + + intra_node_ranks_group.append(_intra_node_ranks) + inter_node_ranks_group.append(_inter_node_ranks) + + _log_rank0(f"create for hierarchy all-gather groups: intra nodes {intra_node_ranks_group}") + _log_rank0(f"create for hierarchy all-gather groups: inter nodes {inter_node_ranks_group}") + + # create communicators + for shard_group in intra_node_ranks_group: + for intra_node_ranks in shard_group: + _group = dist.new_group(intra_node_ranks) + if global_rank in intra_node_ranks: + groups.param_intra_node_group = _group + _log_rank0(f'create group for intra node ranks {intra_node_ranks}') + + for shard_group in inter_node_ranks_group: + for inter_node_ranks in shard_group: + _group = dist.new_group(inter_node_ranks) + if global_rank in inter_node_ranks: + groups.param_inter_node_shard_group = _group + _log_rank0(f'create group for inter node ranks {inter_node_ranks}') + return groups + + +def _generate_mics_config(world_size, ndev_per_node, shard_size, pp_size=1): + """Generating the configuration for sharding This shard config generation assume + that the pipeline stages are partitioned in order, i.e., first ranks + hold the stage0, etc. + + Args: + + shard_size (int): zero3 data-parallel shard size, FIXME: + change the name later + + pp_size (int): pipeline parallel size, currently, only work with + pipeline parallelism + zero + + """ + assert world_size % pp_size == 0 + assert (world_size // pp_size) % shard_size == 0, \ + f"dp group size is not dividable by dp_shard_size, "\ + f" (world_size {world_size}, pp_size {pp_size}, dp_shard_size {shard_size})" + + config = {} + shard_groups = np.arange(world_size).reshape(-1, shard_size) + replicate_groups = [] + for i in range(shard_size): + same_shard_ranks = shard_groups[:, i].tolist() + n_ranks = len(same_shard_ranks) + replicate_size = n_ranks // pp_size + replicate_groups.extend([same_shard_ranks[j:j + replicate_size] for j in range(0, n_ranks, replicate_size)]) + + config['replicate_groups'] = replicate_groups + config['shard_groups'] = shard_groups.tolist() + config["span_nodes"] = len(shard_groups[0]) // ndev_per_node + return config + + +def _sizes_all_same(groups): + """all groups have same length""" + all_same = True + for g in groups: + if len(g) != len(groups[0]): + return False + return all_same diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/offload_config.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/offload_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ca35d7a7d1691c8ec4e6a7016cdb22b31ef459ac --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/offload_config.py @@ -0,0 +1,109 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import Enum +from pathlib import Path +from pydantic import Field, model_validator +from typing import Optional + +from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int + + +class OffloadDeviceEnum(str, Enum): + """ Enum for valid offload devices """ + none = "none" + cpu = "cpu" + nvme = "nvme" + + +class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel): + """ Set options for parameter offload. Valid only with stage 3. """ + + device: OffloadDeviceEnum = "none" + """ + Device memory to offload model parameters. Supported options are `cpu` and + `nvme`. + """ + + nvme_path: Optional[Path] = None + """ Filesystem path for NVMe device for parameter offloading. """ + + buffer_count: int = Field(5, ge=0) + """ Number of buffers in buffer pool for parameter offloading to NVMe. """ + + buffer_size: int = Field(pp_int(1e8), ge=0) + """ Size of buffers in buffer pool for parameter offloading to NVMe. """ + + max_in_cpu: int = Field(pp_int(1e9), ge=0) + """ + Number of parameter elements to maintain in CPU memory when offloading to + NVMe is enabled. + """ + + pin_memory: bool = False + """ + Offload to page-locked CPU memory. This could boost throughput at the cost + of extra memory overhead. + """ + + +class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel): + """ Set options for optimizer offload. Valid with stage 1, 2, and 3. """ + + device: OffloadDeviceEnum = "none" + """ + Device memory to offload optimizer state. Supported options are `cpu` and + `nvme`. Optimizer computation is offload to CPU regardless of device option. + """ + + nvme_path: Optional[Path] = None + """ Filesystem path for NVMe device for optimizer state offloading. """ + + buffer_count: int = Field(4, ge=0) + """ + Number of buffers in buffer pool for optimizer state offloading to NVMe. + This should be at least the number of states maintained per parameter by + the optimizer. For example, Adam optimizer has 4 states (parameter, + gradient, momentum, and variance). + """ + + pin_memory: bool = False + """ + Offload to page-locked CPU memory. This could boost throughput at the cost + of extra memory overhead. + """ + + pipeline_read: bool = False + """ + For tile-based optimizer step processing, overlap read of next tile with + computation of current tile. Used in ZeRO-Infinity. + """ + + pipeline_write: bool = False + """ + For tile-based optimizer step processing, overlap write of previous tile + with computation of current tile. + """ + + fast_init: bool = False + """ Enable fast optimizer initialization when offloading to NVMe. """ + + ratio: float = Field(1.0, ge=0.0, le=1.0) + """ Percentage of offloaded optimizer states to CPU Adam. Only valid with ZeRO Stage 3.""" + + @model_validator(mode="after") + def set_pipeline(self): + pipeline = self.pipeline_read or self.pipeline_write + self.__dict__["pipeline"] = pipeline + return self + + +class OffloadStateTypeEnum(str, Enum): + """ Enum for internal buffer types """ + optim_states = "optim_states" + hp_params = "hp_params" + lp_params = "lp_params" + lp_grads = "lp_grads" + contiguous_grad_buffer = "contiguous_grad_buffer" diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/offload_states.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/offload_states.py new file mode 100644 index 0000000000000000000000000000000000000000..6a8e50c0f9d989bd803543dd1e939dc435d15410 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/offload_states.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Set +import torch + +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.offload_config import OffloadStateTypeEnum + + +def _make_offload_state_key(key): + return f"{key}_offload_buffer" + + +def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_key(state, key): + offload_buf_key = _make_offload_state_key(key) + if offload_buf_key not in state: + state[offload_buf_key] = torch.empty_like(state[key], device=device) + if pin_memory: + state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) + state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) + state[key].data = state[offload_buf_key] + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_key(state, "exp_avg_sq") + + +def reload_adam_states(optimizer, device, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_back_key(state, key): + state[key].data = state[_make_offload_state_key(key)].to(device, non_blocking=non_blocking) + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_back_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_back_key(state, "exp_avg_sq") + + +def get_state_devices(model, state: OffloadStateTypeEnum) -> Set[torch.device]: + """Retrieve the devices of the specified state of the model. + + Args: + model (DeepSpeedEngine): The model whose device allocations are to be checked. + state (OffloadStateTypeEnum): The specific state for which the devices should be retrieved. + + Returns: + Set[torch.device]: A set of devices of the specified state. + + """ + if state == OffloadStateTypeEnum.hp_params: + return set(model.optimizer.get_hp_param_device(p) for p in model.parameters()) + elif state == OffloadStateTypeEnum.lp_params: + return set(p.ds_tensor.device for p in model.parameters()) + elif state == OffloadStateTypeEnum.lp_grads: + return {model.optimizer.grad_partitions_flat_buffer.device} + elif state == OffloadStateTypeEnum.optim_states: + return set(model.optimizer.get_hp_param_device(p, "exp_avg") for p in model.parameters()) | \ + set(model.optimizer.get_hp_param_device(p, "exp_avg_sq") for p in model.parameters()) + elif state == OffloadStateTypeEnum.contiguous_grad_buffer: + if model.optimizer._DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer == None: + return {} + return {model.optimizer._DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer.device} diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/parameter_offload.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/parameter_offload.py new file mode 100644 index 0000000000000000000000000000000000000000..fe53deb47ab82498fd7224768264f5251a32b3e0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/parameter_offload.py @@ -0,0 +1,589 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import torch +from collections import OrderedDict +from deepspeed.utils import z3_leaf_module, set_z3_leaf_module +from deepspeed.runtime.utils import see_memory_usage +from deepspeed.runtime.zero.utils import apply_to_tensors_only, is_zero_param +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.runtime.zero.partition_parameters import _init_external_params +from deepspeed.runtime.zero.partition_parameters import * +from deepspeed.runtime.zero.partitioned_param_coordinator import PartitionedParameterCoordinator, InflightParamRegistry, iter_params +from deepspeed.accelerator import get_accelerator +from deepspeed import utils + +FWD_MODULE_STACK = list() + + +#for each tensor in outputs run the forward_function and register backward_function as hook +def _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, outputs): + if type(outputs) is tuple: + touched_outputs = [] + for output in outputs: + touched_output = _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, + output) + touched_outputs.append(touched_output) + return tuple(touched_outputs) + elif type(outputs) is torch.Tensor: + forward_function(outputs) + if outputs.requires_grad: + outputs.register_hook(backward_function) + return outputs + else: + return outputs + + +class ZeROOrderedDict(OrderedDict): + + def __init__(self, parent_module, *args, **kwargs): + """A replacement for ``collections.OrderedDict`` to detect external ZeRO params. + + Args: + parent_module (``collections.OrderedDict``): the collection to replace + """ + + super().__init__(*args, **kwargs) + self._parent_module = parent_module + self._in_forward = False + + def __reduce__(self): + r0, _, *r2 = super().__reduce__() + return (r0, (self._parent_module, )) + tuple(r2) + + def __getitem__(self, key): + param = super().__getitem__(key) + + # Params can be registered as None (e.g., bias) + if param is None: + return param + + # TODO: only weaken this check during compilation + if hasattr(param, "ds_status") and param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + if self._parent_module._parameters._in_forward: + register_external_parameter(FWD_MODULE_STACK[-1], param) + param.all_gather() + print_rank_0(f'Registering external parameter from getter {key} ds_id = {param.ds_id}', force=False) + + return param + + +def _inject_parameters(module, cls): + for module in module.modules(): + module._original_parameters = module._parameters + + if cls == ZeROOrderedDict: + new_param = cls(parent_module=module) + else: + new_param = cls() + + for key, param in module._parameters.items(): + new_param[key] = param + + module._parameters = new_param + + +class DeepSpeedZeRoOffload(object): + + def __init__( + self, + module, + timers, + ds_config, + overlap_comm=True, + prefetch_bucket_size=50000000, + max_reuse_distance=1000000000, + max_live_parameters=1000000000, + param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, + dp_process_group=None, + offload_param_config=None, + mpu=None, + zero_param_parallel_group=None, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + zero_module_granularity_threshold=0, + log_trace_cache_warnings=False, + ): + + see_memory_usage("DeepSpeedZeRoOffload initialize [begin]", force=True) + + print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False) + + self.module = module + self.timers = timers + self.dtype = list(module.parameters())[0].dtype + self.dp_process_group = dp_process_group + self.offload_device = None + self.offload_param_pin_memory = False + self.zero_param_parallel_group = zero_param_parallel_group + self.zero_quantized_weights = zero_quantized_weights + self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights + self.log_trace_cache_warnings = log_trace_cache_warnings + + if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none: + self.offload_device = offload_param_config.device + self.offload_param_pin_memory = offload_param_config.pin_memory + + self._convert_to_zero_parameters(ds_config, module, mpu) + + for m in module.modules(): + _init_external_params(m) + + _inject_parameters(module, ZeROOrderedDict) + + self.param_numel_persistence_threshold = int(param_persistence_threshold) + self.model_persistence_threshold = int(model_persistence_threshold) + self.persistent_parameters = self.mark_persistent_parameters(self.param_numel_persistence_threshold, + self.model_persistence_threshold) + + self._prefetch_bucket_sz = int(prefetch_bucket_size) + self._max_reuse_distance_in_numel = int(max_reuse_distance) + self._max_available_parameters_in_numel = int(max_live_parameters) + self.__allgather_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream( + ) if overlap_comm else get_accelerator().default_stream() + + if not hasattr(module, "ds_inflight_param_registry"): + module.ds_inflight_param_registry = InflightParamRegistry() + self.__inflight_param_registry = module.ds_inflight_param_registry + + self.fast_sharding_for_leaf_module = False + + if zero_module_granularity_threshold > 0: + self.min_granularity_value = sys.maxsize + self.min_granularity_layer = None + self.granularity_info = set() + self.z3_leaf_layers = [] + self._set_z3_leaf_modules_by_threshold(module, zero_module_granularity_threshold) + self.fast_sharding_for_leaf_module = True + + self.param_coordinator = PartitionedParameterCoordinator( + prefetch_bucket_sz=self._prefetch_bucket_sz, + max_reuse_distance_in_numel=self._max_reuse_distance_in_numel, + max_available_parameters_in_numel=self._max_available_parameters_in_numel, + allgather_stream=self.__allgather_stream, + inflight_param_registry=self.__inflight_param_registry, + prefetch_nvme=self.offload_device == OffloadDeviceEnum.nvme, + timers=self.timers, + zero_quantized_weights=self.zero_quantized_weights, + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights, + fast_sharding_for_leaf_module=self.fast_sharding_for_leaf_module, + log_trace_cache_warnings=self.log_trace_cache_warnings, + ) + + self.forward_hooks = [] + self.backward_hooks = [] + + self.setup_zero_stage3_hooks() + print_rank_0( + f'Created module hooks: forward = {len(self.forward_hooks)}, backward = {len(self.backward_hooks)}', + force=False) + + see_memory_usage("DeepSpeedZeRoOffload initialize [end]", force=True) + + @instrument_w_nvtx + def partition_all_parameters(self): + """Partitioning Parameters that were not partitioned usually if parameters + of modules whose input parameters do not require grad computation do not + trigger post call and will therefore will remain unpartitioned""" + self.get_param_coordinator().release_and_reset_all(self.module) + for param in iter_params(self.module, recurse=True): + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(f"{param.ds_summary()} expected to be released") + + def get_param_coordinator(self): + return self.param_coordinator + + def empty_partition_cache(self): + self.partition_all_parameters() + + def _convert_to_zero_parameters(self, ds_config, module, mpu): + non_zero_params = [p for p in module.parameters() if not is_zero_param(p)] + if non_zero_params: + zero_params = [p for p in module.parameters() if is_zero_param(p)] + if zero_params: + zero_params[0].convert_to_zero_parameters(param_list=non_zero_params) + else: + group = None + # parallel_state_sp doesn't have get_data_parallel_group + if mpu and hasattr(mpu, "get_data_parallel_group"): + group = mpu.get_data_parallel_group() + + Init(module=module, + data_parallel_group=group, + dtype=self.dtype, + config_dict_or_path=ds_config, + remote_device=self.offload_device, + pin_memory=self.offload_param_pin_memory, + mpu=mpu, + zero_param_parallel_group=self.zero_param_parallel_group, + zero_quantized_weights=self.zero_quantized_weights, + zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights) + + def destroy(self): + self._remove_module_hooks() + + def _remove_module_hooks(self): + num_forward_hooks = len(self.forward_hooks) + num_backward_hooks = len(self.backward_hooks) + + for hook in self.forward_hooks: + hook.remove() + + for hook in self.backward_hooks: + hook.remove() + + self.fwd_pre_hook.remove() + + print_rank_0(f'Deleted module hooks: forward = {num_forward_hooks}, backward = {num_backward_hooks}', + force=False) + + def setup_zero_stage3_hooks(self): + self.hierarchy = 0 + + #reset step if in inference mode + @instrument_w_nvtx + def _start_of_forward_hook(module, *args): + + self.get_param_coordinator().reset_step() + + self.fwd_pre_hook = self.module.register_forward_pre_hook(_start_of_forward_hook) + + #likely one of them should be enough but just to be safe + self._register_deepspeed_module(self.module) + + # Add top module to stack trace + global FWD_MODULE_STACK + FWD_MODULE_STACK.append(self.module) + + def mark_persistent_parameters(self, param_threshold, model_threshold): + persistent_params = [] + total_persistent_parameters = 0 + params_count = 0 + for name, param in self.module.named_parameters(recurse=True): + if param.ds_numel + total_persistent_parameters > model_threshold: + continue + + if param.ds_numel <= param_threshold: + params_count += 1 + param.ds_persist = True + persistent_params.append(param) + total_persistent_parameters += param.ds_numel + + print_rank_0( + f"Parameter Offload: Total persistent parameters: {total_persistent_parameters} in {params_count} params", + force=True) + + return persistent_params + + def _register_deepspeed_module(self, module, count=[0]): + my_count = count[0] + module.ds_id = my_count + + #print(f"{module.__class__} : {module.ds_id}") + + if z3_leaf_module(module): + for param in module.parameters(): + param.ds_z3_leaf_module = module + else: + for child in module.children(): + count[0] = count[0] + 1 + self._register_deepspeed_module(child, count=count) + + @torch.compiler.disable + def _pre_forward_module_hook(module, *args): + self.pre_sub_module_forward_function(module) + + @instrument_w_nvtx + def _post_forward_module_hook(module, input, output): + + global FWD_MODULE_STACK + FWD_MODULE_STACK.pop() + if output is None: + output = [] + elif not isinstance(output, (list, tuple)): + if torch.is_tensor(output): + output = [output] + else: + #print(f'got UNKNOWN type {type(output)}') + outputs = [] + output = output if isinstance(output, dict) else vars(output) + for name, val in output.items(): + if not name.startswith('__') and torch.is_tensor(val): + outputs.append(val) + output = outputs + + for item in filter(lambda item: is_zero_param(item) or hasattr(item, 'ds_param_alias'), output): + key = id(item) if hasattr(item, 'ds_id') else id(item.ds_param_alias) + actual_external_param = item if hasattr(item, 'ds_id') else item.ds_param_alias + + if not any(key in m._external_params for m in FWD_MODULE_STACK): + actual_external_param.is_external_param = True + module_to_register = FWD_MODULE_STACK[-1] + register_external_parameter(module_to_register, actual_external_param) + print_rank_0( + f'Registering dangling parameter for module {module_to_register.__class__.__name__}, ds_id = {actual_external_param.ds_id}.', + force=False) + + # It's possible that the parameter was already external to the completed module. If so, remove it the + # registration as it will be covered by the outer module instead. + if key in module._external_params: + print_rank_0( + f' Unregistering nested dangling parameter from module {module.__class__.__name__}, ds_id = {actual_external_param.ds_id}', + force=False) + unregister_external_parameter(module, actual_external_param) + + actual_external_param.all_gather() + + self.post_sub_module_forward_function(module) + + def _bwd_hook_unexpected_inputs_msg(value): + return f"A module has unknown inputs or outputs type ({type(value)}) and the tensors embedded in it cannot be detected. " \ + "The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and " \ + "output tensors and therefore may not get triggered properly." + + def _pre_backward_module_hook(module, inputs, output): + + return apply_to_tensors_only(module.pre_bwd_fn.apply, + output, + warning_msg_fn=_bwd_hook_unexpected_inputs_msg) + + #This is an alternate to doing _post_backward_module_hook + #it uses tensor.register_hook instead of using torch.autograd.Function + def _alternate_post_backward_module_hook(module, inputs): + module.ds_grads_remaining = 0 + + #print(f"Before Forward {module.__class__.__name__}") + + def _run_after_backward_hook(*unused): + module.ds_grads_remaining = module.ds_grads_remaining - 1 + if module.ds_grads_remaining == 0: + #print(f"After backward {module.__class__.__name__}") + self.post_sub_module_backward_function(module) + + def _run_before_forward_function(input): + if input.requires_grad: + module.ds_grads_remaining += 1 + + return _apply_forward_and_backward_to_tensors_only(module, _run_before_forward_function, + _run_after_backward_hook, inputs) + + @torch.compiler.disable + def _post_backward_module_hook(module, inputs): + module.ds_grads_remaining = 0 + + return apply_to_tensors_only(module.post_bwd_fn.apply, + inputs, + warning_msg_fn=_bwd_hook_unexpected_inputs_msg) + + # Pre forward hook + self.forward_hooks.append(module.register_forward_pre_hook(_pre_forward_module_hook)) + + # Post forward hook + self.forward_hooks.append(module.register_forward_hook(_post_forward_module_hook)) + + # Pre backward hook + if not hasattr(module, "pre_bwd_fn"): + + @instrument_w_nvtx + def _run_before_backward_function(sub_module): + # some models (e.g. Albert) may run multiple forwards on the same layer in a loop + # before doing backwards, so each backward will need a pre-fetch - using reference + # counting to support this scenario + #print(f"COUNTER before: {sub_module.applied_pre_backward_ref_cnt}") + if sub_module.applied_pre_backward_ref_cnt > 0: + self.pre_sub_module_backward_function(sub_module) + sub_module.applied_pre_backward_ref_cnt -= 1 + #print(f"COUNTER after: {sub_module.applied_pre_backward_ref_cnt}") + + class PreBackwardFunctionForModule(torch.autograd.Function): + + @staticmethod + def forward(ctx, outputs): + # Capture `module` and _run_before_backward_function + ctx.module = module + ctx.pre_backward_function = _run_before_backward_function + if not hasattr(ctx.module, "applied_pre_backward_ref_cnt"): + ctx.module.applied_pre_backward_ref_cnt = 0 + ctx.module.applied_pre_backward_ref_cnt += 1 + outputs = outputs.detach() + return outputs + + @staticmethod + def backward(ctx, *args): + ctx.pre_backward_function(ctx.module) + return args + + module.pre_bwd_fn = PreBackwardFunctionForModule + + self.backward_hooks.append(module.register_forward_hook(_pre_backward_module_hook)) + + # post backward hook + if not hasattr(module, "post_bwd_fn"): + + @instrument_w_nvtx + def _run_after_backward_function(sub_module): + if sub_module.ds_grads_remaining == 0: + self.post_sub_module_backward_function(sub_module) + + class PostBackwardFunctionModule(torch.autograd.Function): + + @staticmethod + def forward(ctx, output): + ctx.module = module + if output.requires_grad: + #TODO SOME TIMES post backward does not seem to be triggered debug in detail + #Should only cause increase in memory not correctness issue + #if output.grad_fn.__class__.__name__ == 'ViewBackward': + # ctx.view=True + # print(f"Warning view tensor for input to module : {module.__class__.__name__}. Backward hooks may not trigger properly") + #assert len(module.parameters(recurse=False)), "The input tensor to the module is a view, and autograd Function or register_hook is not triggered with view tensors." + #if module.ds_grads_remaining == 0: + # print(f"Before Forward: {ctx.module.__class__.__name__}") + module.ds_grads_remaining += 1 + ctx.post_backward_function = _run_after_backward_function + output = output.detach() + return output + + @staticmethod + def backward(ctx, *args): + ctx.module.ds_grads_remaining = ctx.module.ds_grads_remaining - 1 + if ctx.module.ds_grads_remaining == 0: + ctx.post_backward_function(ctx.module) + return args + + module.post_bwd_fn = PostBackwardFunctionModule + + self.backward_hooks.append(module.register_forward_pre_hook(_post_backward_module_hook)) + + @torch.no_grad() + def pre_sub_module_forward_function(self, sub_module): + see_memory_usage(f"Before sub module function {sub_module.__class__.__name__}", force=False) + + global FWD_MODULE_STACK + FWD_MODULE_STACK.append(sub_module) + + param_coordinator = self.get_param_coordinator() + param_coordinator.trace_prologue(sub_module) + if param_coordinator.is_record_trace(): + param_coordinator.record_module(sub_module) + param_coordinator.fetch_sub_module(sub_module, forward=True) + + see_memory_usage(f"Before sub module function {sub_module.__class__.__name__} after fetch", force=False) + + @torch.no_grad() + def post_sub_module_forward_function(self, sub_module): + see_memory_usage( + f"After sub module function {sub_module.__class__.__name__} {sub_module.ds_id} before release", + force=False) + + param_coordinator = self.get_param_coordinator() + param_coordinator.release_sub_module(sub_module, forward=True) + + see_memory_usage( + f"After sub module function {sub_module.__class__.__name__} {sub_module.ds_id} after release", + force=False) + + @torch.no_grad() + def pre_sub_module_backward_function(self, sub_module): + # assert sub_module.training, "backward pass is invalid for module in evaluation mode" + param_coordinator = self.get_param_coordinator() + param_coordinator.trace_prologue(sub_module) + if param_coordinator.is_record_trace(): + param_coordinator.record_module(sub_module) + param_coordinator.fetch_sub_module(sub_module, forward=False) + + @torch.no_grad() + def post_sub_module_backward_function(self, sub_module): + # assert sub_module.training, "backward pass is invalid for module in evaluation mode" + see_memory_usage( + f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} before release", + force=False) + + self.get_param_coordinator().release_sub_module(sub_module, forward=False) + + see_memory_usage( + f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} after release", + force=False) + + def _set_z3_leaf_modules_by_threshold(self, module, zero_module_granularity_threshold): + + self._get_granularity_recursively(module) + print_rank_0(f"{'MODULE NAME'.ljust(30)}|{'GRANULARITY VALUE'.rjust(20)}", force=True) + for granularity in self.granularity_info: + print_rank_0(granularity, force=True) + + if self.min_granularity_value <= zero_module_granularity_threshold: + self._set_leaf_by_threshold_preorder(module, zero_module_granularity_threshold) + utils.logger.info( + f"z3_leaf_module was set by stage3_module_granularity_threshold:{zero_module_granularity_threshold}") + for layer in self.z3_leaf_layers: + print_rank_0(f"{layer.__class__.__name__}:{layer.ds_model_granularity}", force=True) + else: + utils.logger.warning( + f"The smallest module granularity is [{self.min_granularity_layer}:{self.min_granularity_value}]. "\ + f"To make stage3_module_granularity_threshold effective, you need to set stage3_module_granularity_threshold >= {self.min_granularity_value}. "\ + f"Current Value:{zero_module_granularity_threshold}" + ) + + def _get_granularity_recursively(self, module): + """This function is used to recursively obtain the granularity of each module.""" + + # avoid setting as leaf for particularly large models, even if the granularity is very small + # an oversized leaf module increases the number of live parameters, introducing memory overhead + Z3_MAX_LEAF_SIZE = 1e9 + + if not list(module.parameters()): + # skip Modules without parameters, such as GELU, etc. + module.ds_model_granularity = sys.maxsize + return 0, 0 + + num_layers = 0 + num_params = 0 + num_params += sum(p.ds_numel for p in module.parameters(recurse=False)) + if not any(module.children()): + # torch leaf module + module.ds_model_granularity = sys.maxsize + return 1, num_params + + for child in module.children(): + layers_in_child, params_in_child = self._get_granularity_recursively(child) + num_layers += layers_in_child + num_params += params_in_child + + if module.__class__.__name__ in torch.nn.modules.container.__all__: + # Do not set container modules like ModuleList as leaf modules + # as this will prevent hooks from being set on their children + # and they may do not invoke the forward method + module.ds_model_granularity = sys.maxsize + return num_layers, num_params + + num_layers += 1 + ds_model_granularity = (num_params // num_layers) if num_params <= Z3_MAX_LEAF_SIZE else sys.maxsize + module.ds_model_granularity = ds_model_granularity + # module.ds_model_num_layers = num_layers + # module.ds_model_num_params = num_params + if self.min_granularity_value > ds_model_granularity: + self.min_granularity_value = ds_model_granularity + self.min_granularity_layer = module.__class__.__name__ + self.granularity_info.add(f"{module.__class__.__name__.ljust(30)}|{str(ds_model_granularity).rjust(20)}") + + return num_layers, num_params + + def _set_leaf_by_threshold_preorder(self, module, granularity_treshhold): + '''Set modules as leaf modules based on the threshold, prioritizing parent nodes.''' + + num_params = sum(p.ds_numel for p in module.parameters()) + if num_params == 0: + # skip Modules without parameters, such as GELU, etc. + return + if module.ds_model_granularity <= granularity_treshhold: + set_z3_leaf_module(module, True) + self.z3_leaf_layers.append(module) + return + + for sub_module in module.children(): + self._set_leaf_by_threshold_preorder(sub_module, granularity_treshhold) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py new file mode 100644 index 0000000000000000000000000000000000000000..e9f6bc17405bb3535f2c4ab14a111947bcfe33d9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py @@ -0,0 +1,2258 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +import os +import types +from typing import Callable, Iterable +from enum import Enum +import functools +import itertools +from typing import List +from collections import defaultdict +import logging +import torch +from torch import Tensor +from deepspeed import comm as dist +from torch.nn import Module +from torch.nn import Parameter + +from .linear import zero3_linear_wrap + +from deepspeed.utils import groups +import deepspeed +from ..utils import see_memory_usage, get_only_unique_item +from deepspeed.runtime.zero.config import DeepSpeedZeroConfig +from deepspeed.runtime.zero.utils import assert_ints_same_as_other_ranks, is_zero_param +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.runtime.config_utils import get_config_default +from deepspeed.utils import instrument_w_nvtx, logger +from deepspeed.comm.comm import init_distributed +from deepspeed.utils.debug import (debug_param2name_id_shape, debug_param2name_id_shape_device, debug_module2name, + debug_param2name_id, debug_param2name_id_shape_status) +from deepspeed.accelerator import get_accelerator +from ..swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper, PartitionedParamStatus +from deepspeed.inference.quantization.utils import _quantize_param, WEIGHT_QUANTIZATION_LAYERS, wrap_quantized_functional, wrap_load_from_state_dict + +partitioned_param_data_shape = [0] +zero_init_context = 0 +top_level_context = None + + +class NoGatherHandle: + + def __init__(self, param: Parameter) -> None: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to be available") + + if hasattr(param.ds_tensor, "ds_quant_scale"): + param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to( + device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape) + else: + param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(), + non_blocking=True).view(param.ds_shape) + self.__param = param + + def wait(self, **kwargs) -> None: + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().synchronize() + self.__param.ds_status = ZeroParamStatus.AVAILABLE + + +class NoGatherCoalescedHandle: + + def __init__(self, params: List[Parameter]) -> None: + self.__params = params + self.__complete = False + + for param in self.__params: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to not be available") + if hasattr(param.ds_tensor, "ds_quant_scale"): + param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to( + device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape) + else: + param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(), + non_blocking=True).view(param.ds_shape) + + @instrument_w_nvtx + def wait(self, **kwargs) -> None: + if self.__complete: + return + + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().synchronize() + for param in self.__params: + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + param.ds_status = ZeroParamStatus.AVAILABLE + + self.__complete = True + + +def _dist_allgather_fn(input_tensor: Tensor, output_tensor: Tensor, group=None): + return instrument_w_nvtx(dist.allgather_fn)(output_tensor, input_tensor, group=group, async_op=True) + + +def print_rank_0(message, debug=False, force=False): + rank = dist.get_rank() + if rank == 0 and (debug or force): + print(message) + # other variations + # - print for all ranks w/o interleaving + # printflock(f"[{rank}] {message}") + # - print to log file per rank + # log_rank_file(rank, message) + + +def debug_rank0(msg: str) -> None: + if dist.get_rank() == 0: + logger.debug(msg) + + +def _init_external_params(module): + if not hasattr(module, '_external_params'): + module._external_params = {} + + def external_parameters(self): + return self._external_params.items() + + def all_parameters(self): + return itertools.chain(self.named_parameters(self, recurse=False), external_parameters(self)) + + module.ds_external_parameters = types.MethodType(external_parameters, module) + module.all_parameters = types.MethodType(all_parameters, module) + + +def register_external_parameter(module, parameter): + """Instruct DeepSpeed to coordinate ``parameter``'s collection and partitioning in + the forward and backward passes of ``module``. + + This is used when a parameter is accessed outside of its owning module's + ``forward()``. DeepSpeed must know to collect it from its partitioned + state and when to release the memory. + + .. note:: + This is only applicable to training with ZeRO stage 3. + + Args: + module (``torch.nn.Module``): The module that requires ``parameter`` in its forward pass. + parameter (``torch.nn.Parameter``): The parameter to register. + + Raises: + RuntimeError: If ``parameter`` is not of type ``torch.nn.Parameter``. + + + Examples + ======== + + #. Register a weight that is used in another module's forward pass (line 6). + Parameter ``layer1.weight`` is used by ``layer2`` (line 11). + + .. code-block:: python + :linenos: + :emphasize-lines: 6,11 + + class ModuleZ3(torch.nn.Module): + def __init__(self, *args): + super().__init__(self, *args) + self.layer1 = SomeLayer() + self.layer2 = OtherLayer() + deepspeed.zero.register_external_parameter(self, self.layer1.weight) + + def forward(self, input): + x = self.layer1(input) + # self.layer1.weight is required by self.layer2.forward + y = self.layer2(x, self.layer1.weight) + return y + """ + if not isinstance(parameter, torch.nn.Parameter): + raise RuntimeError('Parameter is not a torch.nn.Parameter') + + if not hasattr(module, '_external_params'): + _init_external_params(module) + + key = id(parameter) + module._external_params[key] = parameter + + +def unregister_external_parameter(module, parameter): + """Reverses the effects of :meth:`register_external_parameter`. + + Args: + module (``torch.nn.Module``): The module to affect. + parameter (``torch.nn.Parameter``): The parameter to unregister. + + Raises: + RuntimeError: If ``parameter`` is not of type ``torch.nn.Parameter``. + RuntimeError: If ``parameter`` is not a registered external parameter of ``module``. + """ + if not isinstance(parameter, torch.nn.Parameter): + raise RuntimeError('Parameter is not a torch.nn.Parameter') + + if not hasattr(module, '_external_params') or id(parameter) not in module._external_params: + raise RuntimeError('Parameter is not a registered external parameter of module.') + + key = id(parameter) + del module._external_params[key] + + +class ZeroParamType(Enum): + + # same as regular pytorch parameters + NORMAL = 1 + + # parameters are partitioned across data parallel process + PARTITIONED = 2 + + # the parameter is held with a unique process rank + # and is not available on all other process + REMOTE = 3 + + +class ZeroParamStatus(Enum): + # parameters are fully present and ready for use on all processes + AVAILABLE = 1 + + # parameters are either partitioned or remote in some or all process + NOT_AVAILABLE = 2 + + # parameters are being gathered. + INFLIGHT = 3 + + +_orig_torch_tensor = torch.tensor +_orig_torch_empty = torch.empty +_orig_torch_zeros = torch.zeros +_orig_torch_ones = torch.ones +_orig_torch_full = torch.full +_orig_torch_arange = torch.arange +_orig_torch_eye = torch.eye +_orig_torch_randn = torch.randn + + +def zero_wrapper_for_fp_tensor_constructor(fn: Callable, target_fp_dtype: torch.dtype) -> Callable: + + def wrapped_fn(*args, **kwargs) -> Tensor: + if kwargs.get("device", None) is None: + kwargs['device'] = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])) + tensor: Tensor = fn(*args, **kwargs) + if tensor.is_floating_point(): + tensor.data = tensor.data.to(target_fp_dtype) + + return tensor + + return wrapped_fn + + +def get_new_tensor_fn_for_dtype(dtype: torch.dtype) -> Callable: + + def new_tensor(cls, *args, **kwargs) -> Tensor: + device = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])) + if not args: + args = (0, ) + tensor = _orig_torch_empty(0, device=device).new_empty(*args, **kwargs) + if tensor.is_floating_point(): + tensor = tensor.to(dtype) + + return tensor + + return new_tensor + + +# https://stackoverflow.com/a/63851681/9201239 +def get_all_subclasses(cls, include_root=True): + subclass_list = [] + + def recurse(cl): + for subclass in cl.__subclasses__(): + subclass_list.append(subclass) + recurse(subclass) + + recurse(cls) + + ret = set(subclass_list) + if include_root: + ret.add(cls) + return ret + + +@instrument_w_nvtx +def free_param(param: Parameter) -> None: + """Free underlying storage of a parameter.""" + assert not param.ds_active_sub_modules, param.ds_summary() + if get_accelerator().on_accelerator(param.data): + # need to make sure that we don't free the parameter while it is still + # being used for computation + if not get_accelerator().is_synchronized_device(): + param.data.record_stream(get_accelerator().current_stream()) + # param.data doesn't store anything meaningful in partitioned state + param.data = torch.empty(0, dtype=param.dtype, device=param.device) + param.ds_status = ZeroParamStatus.NOT_AVAILABLE + + +reuse_buffers = False +temp_contiguous_tensor = None +empty_buffers = {} + + +# Inserts _post_init_method at the end of init method +# for all sub classes of torch.nn.Module +class InsertPostInitMethodToModuleSubClasses(object): + num_module_parameters = 0 + num_module_elements = 0 + + def __init__(self, enabled=True, mem_efficient_linear=True, ds_config=None, dtype=None): + self.mem_efficient_linear = mem_efficient_linear + self.enabled = enabled + self._set_dtype(ds_config, dtype) + assert self.dtype in [ + torch.half, torch.bfloat16, torch.float + ], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.bfloat16, torch.float]" + self.wrapped_cls = set() + self.skip_init_depth = 0 + + self.quantized_initialization = None + if ds_config is not None and ds_config.weight_quantization_config and ds_config.weight_quantization_config.quantized_initialization: + self.quantized_initialization = ds_config.weight_quantization_config.quantized_initialization + + def __enter__(self): + if not self.enabled: + return + + global zero_init_context + if zero_init_context == 0: + self.patch_init_and_builtins() + global top_level_context + top_level_context = self + + zero_init_context += 1 + + def __exit__(self, exc_type, exc_value, traceback): + if not self.enabled: + return + + global zero_init_context + zero_init_context -= 1 + + # Exiting the top level context + if zero_init_context == 0: + self.unpatch_init_and_builtins() + global top_level_context + top_level_context = None + + if dist.get_rank() == 0: + billion_elems = InsertPostInitMethodToModuleSubClasses.num_module_elements / 1e9 + num_params = InsertPostInitMethodToModuleSubClasses.num_module_parameters + logger.info( + f"finished initializing model - num_params = {num_params}, num_elems = {billion_elems:.2f}B") + + # Now that we cleaned up the metaclass injection, raise the exception. + if exc_type is not None: + return False + + # To be implemented by inheriting classes + def _post_init_method(self, module): + pass + + def _set_dtype(self, ds_config, dtype): + if ds_config is not None and dtype is None: + if ds_config.bfloat16_enabled and ds_config.fp16_enabled: + raise RuntimeError("bfloat16 and fp16 cannot be enabled at once") + + if ds_config.bfloat16_enabled: + self.dtype = torch.bfloat16 + elif ds_config.fp16_enabled: + self.dtype = torch.half + else: + self.dtype = torch.float + else: + self.dtype = dtype or torch.float16 if get_accelerator().is_fp16_supported( + ) else torch.bfloat16 if get_accelerator().is_bf16_supported else torch.float32 + + def patch_init_and_builtins(self): + + def apply_with_gather(orig_module_apply_fn: Callable) -> Callable: + """many models make use of child modules like Linear or Embedding which + perform their own weight initialization in their __init__ methods, + but will then have more weight initialization in a parent module's __init__ + method that modifies weights of child modules, which is typically done + using the Module.apply method. + + since the Init context manager partitions child modules immediately after + they are initialized, without modifying apply we would entirely skip + any initialization done by parent modules. + + to get around this issue, we wrap the function passed to Module.apply + so that the applied function is applied to child modules correctly. + """ + + def get_wrapped_fn_to_apply(fn_to_apply: Callable) -> Callable: + if hasattr(fn_to_apply, "wrapped"): + return fn_to_apply + + @functools.wraps(fn_to_apply) + def wrapped_fn_to_apply(module_to_apply_fn_to: Module) -> None: + """gathers parameters before calling apply function. afterwards + parameters are broadcasted to ensure consistency across all ranks + then re-partitioned. + + takes the following steps: + 1. allgathers parameters for the current module being worked on + 2. calls the original function + 3. broadcasts root rank's parameters to the other ranks + 4. re-partitions the parameters + """ + + # TODO Delay error checking for dangling partitioned parameters to post module init + # raise RuntimeError(f"not all parameters for {module_to_apply_fn_to.__class__.__name__}, " + # f"were zero params, is it possible that the parameters were " + # f"overwritten after they were initialized? " + # f"params: {[p for p in module_to_apply_fn_to.parameters(recurse=False)]} ") + + params_to_apply_fn_to: Iterable[Parameter] = list( + sorted([p for p in module_to_apply_fn_to.parameters(recurse=False) if is_zero_param(p)], + key=lambda p: p.ds_id)) + + for param in params_to_apply_fn_to: + param.all_gather() + + fn_to_apply(module_to_apply_fn_to) + + for param in params_to_apply_fn_to: + dist.broadcast(param.data, 0, group=param.ds_process_group) + + for param in params_to_apply_fn_to: + param.partition(has_been_updated=True) + + wrapped_fn_to_apply.wrapped = True + + return wrapped_fn_to_apply + + @functools.wraps(orig_module_apply_fn) + def wrapped_apply(module: Module, fn_to_apply: Callable) -> None: + orig_module_apply_fn(module, get_wrapped_fn_to_apply(fn_to_apply)) + + return wrapped_apply + + def hook_for_skip_init(module): + # this function is intended for handling the logic of torch.nn.utils.skip_init + # skip_init:module_cls(*args, **kwargs).to_empty(device=final_device), where kwargs['device']='meta' + # the function call occurs between module_cls(*args, **kwargs) and to_empty(device=final_device). + def partition_after_empty_init(f): + + @functools.wraps(f) + def wrapper(module, *args, **kwargs): + _module = f(module, *args, **kwargs) + # here is the post-hook for module.apply(empty_like...) + # after module.apply(empty_like...), the module has completed its empty init on real device + # since skip_init won't involve any computations or weight adjustments, we can directly utilize post_init + self._post_init_method(_module) + return _module + + return wrapper + + def post_wrapper_to_empty(f): + # append some wrapper restoration after to_empty() call + @functools.wraps(f) + def wrapper(*args, **kwargs): + res = f(*args, **kwargs) + # restore _apply hook + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _disable_class_apply(subclass) + # self restore + module.to_empty = f + return res + + return wrapper + + def _enable_class_apply(cls): + if '_apply' in cls.__dict__: + cls._old_apply_of_skip_init_hook = cls._apply + cls._apply = partition_after_empty_init(cls._apply) + + def _disable_class_apply(cls): + if hasattr(cls, '_old_apply_of_skip_init_hook'): + cls._apply = cls._old_apply_of_skip_init_hook + + # add hooks for to_empty: apply_(empty_like) + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _enable_class_apply(subclass) + + # add a restore hook when exiting skip_init + module.to_empty = post_wrapper_to_empty(module.to_empty) + + def partition_after(f): + + @functools.wraps(f) + def wrapper(module, *args, **kwargs): + + # important logic: We want to run post_init only after child's __init__ is + # completed, and do nothing after __init__ of any of its parents and grandparents in + # the inheritance ancestry. This way the partitioning will need to happen only once + # when the whole object is ready to be partitioned and not before. This is because + # often the child module will need to tweak the weights - for example running a + # custom weights init function. So if a parent created the weights param, the child + # won't need to gather it in order to tweak it + + print_rank_0(f'Before initializing {module.__class__.__name__}', force=False) + + is_child_module = False + if not hasattr(module, "_ds_child_entered"): + # child's __init__ was called, since parents all see the same object they can now skip post_init + is_child_module = True + setattr(module, "_ds_child_entered", True) + + init_on_meta = 'device' in kwargs and kwargs['device'] == 'meta' + if init_on_meta: + self.skip_init_depth += 1 + + f(module, *args, **kwargs) + if init_on_meta and self.skip_init_depth == 1: + # check and handle the logic of empty_init + hook_for_skip_init(module) + if is_child_module: + # child's __init__ is done, now we can run a single post_init on the child object + delattr(module, "_ds_child_entered") + + print_rank_0(f'Running post_init for {module.__class__.__name__}', force=False) + if self.skip_init_depth == 0: + self._post_init_method(module) + + print_rank_0(f'After initializing followed by post init for {module.__class__.__name__}', force=False) + if init_on_meta: + self.skip_init_depth -= 1 + + return wrapper + + def _enable_class(cls): + if '__init__' in cls.__dict__: + cls._old_init = cls.__init__ + cls.__init__ = partition_after(cls.__init__) + + def _init_subclass(cls, **kwargs): + if '__init__' in cls.__dict__: + cls._old_init = cls.__init__ + cls.__init__ = partition_after(cls.__init__) + + # Replace .__init__() for all existing subclasses of torch.nn.Module recursively + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _enable_class(subclass) + + # holding onto some methods so we can put them back the way they were in __exit__ + torch.nn.modules.module.Module._old_init_subclass = torch.nn.modules.module.Module.__init_subclass__ + torch.nn.modules.module.Module._old_apply = torch.nn.modules.module.Module.apply + torch.Tensor.__old_new__ = torch.Tensor.__new__ + + # Replace .__init__() for future subclasses of torch.nn.Module + torch.nn.modules.module.Module.__init_subclass__ = classmethod(_init_subclass) + if Init.override_module_apply: + torch.nn.modules.module.Module.apply = apply_with_gather(torch.nn.modules.module.Module._old_apply) + + self._add_tensor_creation_wrappers() + + if self.mem_efficient_linear: + print_rank_0( + "nn.functional.linear has been overridden with a more memory efficient version. This will persist unless manually reset.", + force=False) + if not hasattr(InsertPostInitMethodToModuleSubClasses, "linear_bk"): + InsertPostInitMethodToModuleSubClasses.linear_bk = torch.nn.functional.linear + torch.nn.functional.linear = zero3_linear_wrap + + if self.quantized_initialization: + print_rank_0("nn.functional.linear has been overridden with quantized linear version.", force=False) + torch.nn.functional.linear = wrap_quantized_functional(torch.nn.functional.linear) + torch.nn.functional.embedding = wrap_quantized_functional(torch.nn.functional.embedding) + for cls in WEIGHT_QUANTIZATION_LAYERS: + cls._load_from_state_dict = wrap_load_from_state_dict(cls._load_from_state_dict) + + logger.info("Enable Zero3 engine with INT4 quantization.") + + self.patched = True + + def unpatch_init_and_builtins(self): + if self.patched: + + def _disable_class(cls): + if hasattr(cls, '_old_init'): + cls.__init__ = cls._old_init + + for subclass in get_all_subclasses(torch.nn.modules.module.Module): + _disable_class(subclass) + + # putting methods back the way we found them + torch.nn.modules.module.Module.__init_subclass__ = torch.nn.modules.module.Module._old_init_subclass + if Init.override_module_apply: + torch.nn.modules.module.Module.apply = torch.nn.modules.module.Module._old_apply + + self._remove_tensor_creation_wrappers() + + self.patched = False + + def _add_tensor_creation_wrappers(self): + torch.Tensor.__new__ = get_new_tensor_fn_for_dtype(self.dtype) + torch.tensor = zero_wrapper_for_fp_tensor_constructor(_orig_torch_tensor, self.dtype) + torch.empty = zero_wrapper_for_fp_tensor_constructor(_orig_torch_empty, self.dtype) + torch.zeros = zero_wrapper_for_fp_tensor_constructor(_orig_torch_zeros, self.dtype) + torch.ones = zero_wrapper_for_fp_tensor_constructor(_orig_torch_ones, self.dtype) + torch.full = zero_wrapper_for_fp_tensor_constructor(_orig_torch_full, self.dtype) + torch.arange = zero_wrapper_for_fp_tensor_constructor(_orig_torch_arange, self.dtype) + torch.eye = zero_wrapper_for_fp_tensor_constructor(_orig_torch_eye, self.dtype) + torch.randn = zero_wrapper_for_fp_tensor_constructor(_orig_torch_randn, self.dtype) + + def _remove_tensor_creation_wrappers(self): + torch.Tensor.__new__ = torch.Tensor.__old_new__ + torch.tensor = _orig_torch_tensor + torch.empty = _orig_torch_empty + torch.zeros = _orig_torch_zeros + torch.ones = _orig_torch_ones + torch.full = _orig_torch_full + torch.arange = _orig_torch_arange + torch.eye = _orig_torch_eye + torch.randn = _orig_torch_randn + + +def shutdown_init_context(): + """ + This function is used to initialize deepspeed engine inside the context of Init. + We need to remove the wrappers but keep the context. + """ + if top_level_context: + top_level_context.unpatch_init_and_builtins() + + +def restore_init_context(): + """ + This function is used to restore the wrappers after deepspeed engine is initialized. + """ + if top_level_context: + top_level_context.patch_init_and_builtins() + + +class AllGatherHandle: + + def __init__(self, handle, param: Parameter, quantization=None) -> None: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to be available") + + self.__handle = handle + self.__param = param + self.__quantization = quantization + + def wait(self, handle_dependency=True) -> None: + instrument_w_nvtx(self.__handle.wait)() + if self.__quantization: + instrument_w_nvtx(self.__quantization.quant_handle.wait)() + self.__param.data = self.__quantization.backend.dequantize( + self.__quantization.quantized_param, self.__quantization.scale_buffer).to(self.__param.device) + self.__param.ds_status = ZeroParamStatus.AVAILABLE + + +class AllGatherCoalescedHandle: + + data_buffer = [] + + def __init__( + self, + allgather_handle, + params: List[Parameter], + partitions: List[Tensor], + world_size: int, + use_secondary_tensor=False, + quantization=None, + ) -> None: + self.allgather_handle = allgather_handle + self.params = params + self.partitions = partitions + self.world_size = world_size + self.use_secondary_tensor = use_secondary_tensor + self.complete = False + self.quantization = quantization + + for param in self.params: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to not be available") + + @instrument_w_nvtx + def wait(self, handle_dependency=True) -> None: + if self.complete: + return + + instrument_w_nvtx(self.allgather_handle.wait)() + + if self.quantization: + instrument_w_nvtx(self.quantization.quant_handle.wait)() + flat_tensor = self.quantization.backend.dequantize( + self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device) + + self.partitions: List[Parameter] = [] + for i in range(self.world_size): + self.partitions.append( + flat_tensor.narrow(0, self.quantization.partition_sz * i, self.quantization.partition_sz)) + + # split the single tensor out into individual tensors + param_offset = 0 + for param in self.params: + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + partitions: List[Tensor] = [] + ds_tensor_numel = param.ds_tensor.ds_numel + if self.use_secondary_tensor: + ds_tensor_numel *= param.ds_secondary_tensor_num_of_groups + for rank in range(self.world_size): + param_start = rank * ds_tensor_numel + if param_start < param.ds_numel: + part_to_copy = self.partitions[rank].narrow(0, param_offset, + min(param.ds_numel - param_start, ds_tensor_numel)) + partitions.append(part_to_copy) + param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape) + param.ds_status = ZeroParamStatus.AVAILABLE + if not get_accelerator().is_synchronized_device() and handle_dependency: + for part_to_copy in partitions: + part_to_copy.record_stream(get_accelerator().current_stream()) + + param_offset += ds_tensor_numel + + self.complete = True + if not get_accelerator().is_synchronized_device() and not handle_dependency: + # if the device needs to handle dependencies and opts for explicit processing outside the function. + AllGatherCoalescedHandle.data_buffer.append(partitions) + + @staticmethod + def free_buffer(): + AllGatherCoalescedHandle.data_buffer = [] + + +class MultipleAllGatherHandles: + + def __init__(self, handles: List[AllGatherCoalescedHandle]): + self.handles = handles + + def wait(self, handle_dependency=True) -> None: + for handle in self.handles: + handle.wait(handle_dependency) + + +class AllReduceCoalescedHandle: + + def __init__(self, handle, params: List[Parameter]) -> None: + self.handle = handle + self.params = params + self.complete = False + + for param in self.params: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to not be available") + + @instrument_w_nvtx + def wait(self) -> None: + if self.complete: + return + + instrument_w_nvtx(self.handle.wait)() + + for param in self.params: + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + param.ds_status = ZeroParamStatus.AVAILABLE + + self.complete = True + + +class QuantizationInfo: + # a placeholder object to store all quant related vars used in handles + def __init__(self) -> None: + self.quantized_param = None + self.backend = None + self.quant_handle = None + self.scale_buffer = None + + +class CUDAQuantizer: + async_flag = True + target_group_size = 8000 # the optimal size is 4k, so we set the target to be below 8k + group_size_cache = dict() + quantizer_cuda_module = None + + def __init__(self) -> None: + if CUDAQuantizer.quantizer_cuda_module is None: + CUDAQuantizer.quantizer_cuda_module = deepspeed.ops.op_builder.QuantizerBuilder().load() + + def quantize(self, param, groups=None): + if groups is None: + try: + groups = self.group_size_cache[param.numel()] + except KeyError: + groups = math.ceil(param.numel() / self.target_group_size) + while groups < param.numel(): + if param.numel() % (8 * groups) == 0: + break + groups += 1 + while True: + if param.numel() % (8 * groups * 2) == 0 and param.numel( + ) / groups > self.target_group_size: #hard limit of 16k group_size + groups *= 2 + else: + break + assert ( + param.numel() % (8 * groups) == 0 + ), f"Qantized weight requires the number of weights be a multiple of 8. Yet {param.numel()} cannot be divided by 8*{groups}" + assert (param.numel() / groups < 16000), f"{param.numel()} / {groups} is larger than 16k" + assert param.numel( + ) > groups, f"Adaptive grouping algorithm cannot find a group size for input tensor of size {param.numel()}" + self.group_size_cache[param.numel()] = groups + return self.quantizer_cuda_module.quantize(param.to(get_accelerator().device_name()), groups, 8, + self.quantizer_cuda_module.Symmetric) + + def dequantize(self, quantized_param, scale): + return self.quantizer_cuda_module.dequantize(quantized_param, scale, scale.numel(), 8, + self.quantizer_cuda_module.Symmetric) + + +def _no_gather_coalesced(params: Iterable[Parameter]) -> AllGatherCoalescedHandle: + for param in params: + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(f"expect param.ds_status == ZeroParamStatus.NOT_AVAILABLE, got{param.ds_summary()}") + param.ds_status = ZeroParamStatus.INFLIGHT + + params = sorted(params, key=lambda p: p.ds_id) + if len(params) == 1: + param, = params + return NoGatherHandle(param) + return NoGatherCoalescedHandle(params) + + +# Replaces all parameters in module with Scattered Parameters +class Init(InsertPostInitMethodToModuleSubClasses): + param_id = 0 + param_persistence_threshold = get_config_default(DeepSpeedZeroConfig, "param_persistence_threshold") + model_persistence_threshold = get_config_default(DeepSpeedZeroConfig, "model_persistence_threshold") + num_persisted_parameters = 0 + num_persisted_elements = 0 + apply_param_persistence = False + override_module_apply = get_config_default(DeepSpeedZeroConfig, "override_module_apply") + + def __init__(self, + module=None, + data_parallel_group=None, + mem_efficient_linear=True, + remote_device=None, + pin_memory=False, + config_dict_or_path=None, + config=None, + enabled=True, + dtype=None, + mpu=None, + zero_param_parallel_group=None, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + sequence_data_parallel_group=None, + param_swapper=None): + """A context to enable massive model construction for training with + ZeRO-3. Models are automatically partitioned (or, sharded) across the + system and converted to half precision. + + Args: + module (``torch.nn.Module``, optional): If provided, partition the model as + if it was constructed in the context. + data_parallel_group (``deepspeed.comm`` process group, optional): + The group of processes to partition among. Defaults to all processes. + Synonymous with sequence data parallel group for param partitioning + across both sequence and data parallel groups. + mem_efficient_linear (bool, optional): Replace + torch.nn.functional.linear with an implementation that allows + DeepSpeed to partition parameters. Defaults to ``True``. + remote_device (string, optional): The initial device to store model + weights e.g., ``cpu``, ``nvme``. Passing ``"cpu"`` will create the model in CPU + memory. The model may still be moved to GPU based on the + offload settings for training. Defaults to param offload device if a config is + defined, otherwise GPU. + pin_memory (bool, optional): Potentially increase performance by + using pinned memory for model weights. ``remote_device`` must be + ``"cpu"``. Defaults to pin_memory value in config, otherwise ``False``. + config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration + for swapping fp16 params to NVMe and other things like ``dtype``. + config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead. + enabled (bool, optional): If ``False``, this context has no + effect. Defaults to ``True``. + dtype (``dtype``, optional): Can be used to change the data type of the parameters. + Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None`` + mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,world_size}. + zero_param_parallel_group(``object``, optional): Parallel (comm) group for dual partitioning of ZeRO params. + zero_quantized_weights (bool, optional): If ``True``, turn on quantized weights in all gather weights. Default is ``False`` + zero_quantized_nontrainable_weights (bool, optional): If ``True``, nontrainable weights will be stored in quantized format. Default is ``False`` + param_swapper (``deepspeed.runtime.swap_tensor.partitioned_param_swapper.AsyncPartitionedParameterSwapper``, optional): [Experimental] Use existing parameter swapper. Defaults to ``None``. + This argument will be removed in the near future. + + This context accelerates model initialization and enables models that + are too large to allocate in their entirety in CPU memory. It has the + following effects: + + #. allocates tensors to either GPU or CPU memory or NVMe + #. converts floating point tensors to half precision + #. immediately partitions tensors among the group of data-parallel devices + #. (*optional*) replaces ``torch.nn.functional.linear`` with a more + memory-efficient implementation + + These modifications allow for models that exceed the size of local CPU/GPU + memory/NVMe, but fit within the total NVMe capacity (*i.e.*, aggregate CPU + or GPU memory or NVMe) across all nodes. Consider initializing a model with one + trillion parameters, whose weights occupy two terabytes (TB) in half + precision. The initial CPU allocation in full precision requires 4TB of + memory *per process*, and so a system with 8 GPUs per node would need 32TB of + CPU memory due to data-parallel redundancies. Instead, by immediately + partitioning tensors we remove the redundancies. The result is that + regardless of the number of GPUs, we still only require the original 4TB. This + allows for a linear increase in model size with the aggregate system memory. + For example, if a node has 1TB of memory and 8 GPUs, we could fit a trillion + parameter model with 4 nodes and 32 GPUs. + + Important: If the fp16 weights of the model can't fit onto a single GPU memory + this feature must be used. + + .. note:: + Initializes ``deepspeed.comm`` if it has not already been done so. + See :meth:`deepspeed.init_distributed` for more information. + + .. note:: + Only applicable to training with ZeRO-3. + + Examples + -------- + + #. Allocate a model and partition it among all processes: + + .. code-block:: python + + with deepspeed.zero.Init(): + model = MyLargeModel() + + + #. Allocate a model in pinned CPU memory and partition it among a subgroup of processes: + + .. code-block:: python + + with deepspeed.zero.Init(data_parallel_group=mpu.get_data_parallel_group(), + remote_device="cpu", + pin_memory=True): + model = MyLargeModel() + + + #. Partition an already-allocated model in CPU memory: + + .. code-block:: python + + model = deepspeed.zero.Init(module=model) + """ + if config is not None: + config_dict_or_path = config + logger.warning('zero.Init: the `config` argument is deprecated. Please use `config_dict_or_path` instead.') + _ds_config = deepspeed.runtime.config.DeepSpeedConfig(config_dict_or_path, + mpu) if config_dict_or_path is not None else None + if _ds_config is not None: + mem_efficient_linear = _ds_config.zero_config.memory_efficient_linear + + super().__init__(enabled=enabled, mem_efficient_linear=mem_efficient_linear, ds_config=_ds_config, dtype=dtype) + if not dist.is_initialized(): + init_distributed() + assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm" + + if data_parallel_group is None: + self.ds_process_group = dist.get_world_group() + else: + self.ds_process_group = data_parallel_group + + if sequence_data_parallel_group is not None: + logger.warning( + f"sequence_data_parallel_group' is deprecated and will be removed. Use 'data_parallel_group' instead.") + if data_parallel_group is not None: + raise ValueError( + "Both 'data_parallel_group' and 'sequence_data_parallel_group' were specified. Please provide only one of these arguments." + ) + self.ds_process_group = sequence_data_parallel_group + + self.rank = dist.get_rank(group=self.ds_process_group) + self.dp_world_size = dist.get_world_size(group=self.ds_process_group) + + self.zero_param_process_group = zero_param_parallel_group + if _ds_config is not None and _ds_config.zero_config.zero_hpz_partition_size > 1 and self.zero_param_process_group is None: + groups._create_zero_param_parallel_group(_ds_config.zero_config.zero_hpz_partition_size) + self.zero_param_process_group = groups._get_zero_param_intra_parallel_group() + + self.num_ranks_in_param_group = self.dp_world_size + self.rank_in_group = self.rank + self.num_param_groups = 1 + + if self.zero_param_process_group is not None: + self.num_ranks_in_param_group = groups._get_zero_param_intra_parallel_group_world_size() + self.num_param_groups = int(self.dp_world_size / self.num_ranks_in_param_group) + self.rank_in_group = groups._get_zero_param_intra_parallel_rank_in_mygroup() + print_rank_0(f"hpZeRO group size: {self.num_ranks_in_param_group}", force=True) + + logger.debug( + "hpZeRO partition parameter my rank in world {} my rank in group {} ranks in my param partition group: {} " + .format(self.rank, self.rank_in_group, groups._get_zero_param_intra_parallel_group_ranks())) + + # Local device is the device where the parameters are consumed, must be default device. + # It is the device where parameters are fully instantiated using allgather + self.local_device = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])) + get_accelerator().set_device(self.local_device) + + self.quantized_weights = zero_quantized_weights + if _ds_config is not None and _ds_config.zero_config.zero_quantized_weights and not self.quantized_weights: + self.quantized_weights = _ds_config.zero_config.zero_quantized_weights + self.quantized_nontrainable_weights = zero_quantized_nontrainable_weights + if _ds_config is not None and _ds_config.zero_config.zero_quantized_nontrainable_weights and not self.quantized_nontrainable_weights: + self.quantized_nontrainable_weights = _ds_config.zero_config.zero_quantized_nontrainable_weights + + self.module = module + if (self.quantized_weights or self.quantized_nontrainable_weights): + self.quantizer_module = CUDAQuantizer() + print_rank_0(f'Using quantizer for weights: {self.quantizer_module.__class__.__name__}', force=True) + + if _ds_config is not None: + Init.override_module_apply = _ds_config.zero_config.override_module_apply + + if _ds_config.zero_config.offload_param is not None: + remote_device = _ds_config.zero_config.offload_param.device + pin_memory = _ds_config.zero_config.offload_param.pin_memory + + self._validate_remote_device(remote_device, _ds_config) + + # Remote device is the device where parameter partitions are stored + # It can be same as local_device or it could be CPU or NVMe. + self.remote_device = self.local_device if remote_device in [None, OffloadDeviceEnum.none] else remote_device + self.pin_memory = pin_memory if (self.remote_device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme + ]) else False + + # Enable fp16 param swapping to NVMe + if self.remote_device == OffloadDeviceEnum.nvme: + self.param_swapper = param_swapper or AsyncPartitionedParameterSwapper(_ds_config, self.dtype) + else: + self.param_swapper = None + + # If we are provided an already-allocated module to prepare. + if module is not None: + assert isinstance(module, torch.nn.Module) + self._convert_to_zero_parameters(module.parameters(recurse=True)) + + self.use_all_gather_into_tensor = dist.has_all_gather_into_tensor() + if not self.use_all_gather_into_tensor: + logger.info(f"all_gather_into_tensor API is not available in torch {torch.__version__}") + + self.use_all_reduce_for_fetch_params = get_config_default(DeepSpeedZeroConfig, + "use_all_reduce_for_fetch_params") + if _ds_config is not None: + self.use_all_reduce_for_fetch_params = _ds_config.zero_config.use_all_reduce_for_fetch_params + + def _update_persist_config(self, ds_config): + Init.apply_param_persistence = True + Init.param_persistence_threshold = ds_config.zero_config.param_persistence_threshold + Init.model_persistence_threshold = ds_config.zero_config.model_persistence_threshold // self.num_partitions + + def _zero_init_param(self, param): + self._convert_to_deepspeed_param(param) + if dist.get_world_group() == self.get_dp_process_group(): + dist.broadcast(param.data, 0, self.get_dp_process_group()) + else: + dist.broadcast(param.data, dist.get_global_rank(self.get_dp_process_group(), 0), + self.get_dp_process_group()) + param.partition() + + def _convert_to_zero_parameters(self, param_list): + for param in param_list: + if is_zero_param(param): + continue + + param.data = param.data.to(self.local_device) + self._zero_init_param(param) + + def _validate_remote_device(self, remote_device, ds_config): + if ds_config is not None: + if remote_device in [None, OffloadDeviceEnum.cpu]: + if ds_config.zero_config.offload_param is not None: + offload_param_device = ds_config.zero_config.offload_param.device + assert offload_param_device != OffloadDeviceEnum.nvme, \ + f"'device' in DeepSpeed Config cannot be {offload_param_device} if remote device is {remote_device}." + + if remote_device == OffloadDeviceEnum.nvme: + assert ds_config.zero_config.offload_param is not None, \ + f'"offload_param" must be defined in DeepSpeed Config if remote device is {OffloadDeviceEnum.nvme}.' + + assert ds_config.zero_config.offload_param.nvme_path is not None, \ + f'"nvme_path" in DeepSpeed Config cannot be None if remote device is {OffloadDeviceEnum.nvme}' + + def _post_init_method(self, module): + #see_memory_usage(f"Before converting params in {module.__class__.__name__}", force=False) + print_rank_0(f'Converting Params in {module.__class__.__name__}', force=False) + see_memory_usage(f"Before converting and partitioning params in {module.__class__.__name__}", force=False) + + for name, param in module.named_parameters(recurse=False): + print_rank_0(f'Analyzing param {name} in {module.__class__.__name__}', force=False) + InsertPostInitMethodToModuleSubClasses.num_module_parameters += 1 + InsertPostInitMethodToModuleSubClasses.num_module_elements += param.numel() + if not is_zero_param(param): + if not get_accelerator().on_accelerator(param): + param.data = param.data.to(self.local_device) + + if name == 'weight' and self.quantized_initialization and type(module) in WEIGHT_QUANTIZATION_LAYERS: + _quantize_param(param, self.quantized_initialization) + + self._zero_init_param(param) + print_rank_0( + f"Partitioning param {debug_param2name_id_shape(param)} module={debug_module2name(module)}") + + see_memory_usage( + f"Param count {InsertPostInitMethodToModuleSubClasses.num_module_elements}. After converting and partitioning params in {module.__class__.__name__}", + force=False) + + def _convert_to_deepspeed_param(self, param): + + # Partitioned, Normal, Remote + param.ds_param_type = ZeroParamType.PARTITIONED + + # Replicated vs Partitioned vs Inflight + param.ds_status = ZeroParamStatus.AVAILABLE + + # Stores the shape of the original tensor + param.ds_shape = param.shape + + # Stores the number of elements in the original parameter without padding + param.ds_numel = param.numel() + + # Stores the partitioned copy of the tensor + param.ds_tensor = None + + # Keeps track of how many active sub-modules need this param at any given point in time + param.ds_active_sub_modules = set() + + # If this flag is true, then the parameters are replicated throughput training + # And only partitioned before the step + if Init.apply_param_persistence and param.ds_numel <= Init.param_persistence_threshold and Init.num_persisted_elements + param.ds_numel <= Init.model_persistence_threshold: + param.ds_persist = True + Init.num_persisted_parameters += 1 + Init.num_persisted_elements += param.ds_numel + else: + param.ds_persist = False + + param.is_external_param = False + + # The group that the parameter is scattered across. + param.ds_process_group = self.ds_process_group + + # Stores the secondary partitioned copy of the tensor + param.ds_secondary_tensor = None + + #Process group for secondary partition all (group) gather + param.ds_zero_param_process_group = self.zero_param_process_group + param.ds_secondary_tensor_group_size = self.num_ranks_in_param_group + param.ds_secondary_tensor_num_of_groups = self.num_param_groups + + # This is set to the Async Param swapper if remote device is nvme + # else this is set to None + param.nvme_swapper = self.param_swapper + + # DeepSpeed Param ID + param.ds_id = Init.param_id + Init.param_id += 1 + + def all_gather(param_list=None, async_op=False, hierarchy=0): + cls = param + if param_list is None: + param_list = [cls] + return self._all_gather(param_list, async_op=async_op, hierarchy=hierarchy) + + def _all_gather_dtype(dtype, params, world_size, rank_in_group, ds_process_group): + partition_sz = sum(p.ds_tensor.ds_numel for p in params) + + use_secondary_tensor = params[0].ds_secondary_tensor is not None + + if use_secondary_tensor: + partition_sz = sum(p.ds_tensor.ds_numel * p.ds_secondary_tensor_num_of_groups for p in params) + + flat_tensor = torch.empty(partition_sz * world_size, + dtype=dtype, + device=get_accelerator().current_device_name(), + requires_grad=False) + + partitions: List[Parameter] = [] + for i in range(world_size): + partitions.append(flat_tensor.narrow(0, partition_sz * i, partition_sz)) + + if use_secondary_tensor: + instrument_w_nvtx( + torch.cat)([p.ds_secondary_tensor.to(get_accelerator().current_device_name()) for p in params], + out=partitions[rank_in_group]) + else: + instrument_w_nvtx(torch.cat)([p.ds_tensor.to(get_accelerator().current_device_name()) for p in params], + out=partitions[rank_in_group]) + handle = _dist_allgather_fn(partitions[rank_in_group], flat_tensor, ds_process_group) + #Fix get_partition_dp_group(params[0])) + + return AllGatherCoalescedHandle( + allgather_handle=handle, + params=params, + partitions=partitions, + world_size=world_size, + use_secondary_tensor=use_secondary_tensor, + ) + + @instrument_w_nvtx + def all_gather_coalesced(params: Iterable[Parameter], + safe_mode: bool = False, + quantize: bool = False) -> AllGatherCoalescedHandle: + + # fetches from nvme if the partition is not available and in nvme + self._ensure_availability_of_partitioned_params(params) + + if self.num_partitions == 1: + return _no_gather_coalesced(params) + + for param in params: + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(param.ds_summary()) + param.ds_status = ZeroParamStatus.INFLIGHT + + #use appropriate all gather process group + ds_process_group = self.ds_process_group + rank_in_group = self.rank + world_size = self.dp_world_size + use_secondary_tensor = params[0].ds_secondary_tensor is not None + if self.zero_param_process_group and use_secondary_tensor: + ds_process_group = self.zero_param_process_group #intragroup + rank_in_group = self.rank_in_group + world_size = self.num_ranks_in_param_group + + #pprint(dir(ds_process_group)) + # ensure that each rank has params in same order. the allgather + # is done by flattening the parameter list into a single tensor that + # can be allgathered in a single call - this means that if each rank + # gives a list of the same parameters in a different order we will + # silently get incorrect parameter values, and have very difficult + # to debug correctness issues. + params = sorted(params, key=lambda p: p.ds_id) + + if logger.isEnabledFor(logging.DEBUG): + debug_rank0(f"-allgather_coalesced: {[p.ds_id for p in params]}") + + if safe_mode: + # ensure that same list (with same ordering) of parameters are + # being allgathered across all ranks, otherwise could mix + # data between tensors. + assert_ints_same_as_other_ranks([p.ds_id for p in params]) + # ensure that tensors from each rank agree on the same ds_numel + # otherwise could mix data between tensors. + assert_ints_same_as_other_ranks([p.ds_tensor.ds_numel for p in params]) + + if len(params) == 1: + # have an opportunity to avoid some intermediate memory allocations + param = params[0] + buffer_size = math.ceil(param.ds_numel / world_size) * world_size + if use_secondary_tensor: + buffer_size = param.ds_secondary_tensor.shape[0] * world_size #make sure out is appropriately sized + + param_ds_tensor = param.ds_secondary_tensor if use_secondary_tensor else param.ds_tensor + param_buffer = torch.empty( + buffer_size, + dtype=param_ds_tensor.dtype if not quantize else torch.int8, + device=get_accelerator().current_device_name(), + requires_grad=False, + ) + if not quantize: + handles = _dist_allgather_fn( + param_ds_tensor.to(get_accelerator().current_device_name()), + param_buffer, + ds_process_group, + ) + param.data = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(param.device) + #print_rank_0(f"{param.shape=}", force=True) + #print_rank_0(f"{param_buffer.shape=}", force=True) + + return AllGatherHandle(handles, param) + else: + if hasattr(param_ds_tensor, "ds_quant_scale"): + scales = param_ds_tensor.ds_quant_scale + quantized_param = param_ds_tensor.data + else: + quantized_param, scales = self.quantizer_module.quantize(param_ds_tensor) + handle = _dist_allgather_fn(quantized_param.to(get_accelerator().current_device_name()), + param_buffer, ds_process_group) + + quant_scale_buffer = torch.empty( + scales.numel() * world_size, + dtype=scales.dtype, + device=get_accelerator().current_device_name(), + requires_grad=False, + ) + quant_handle = _dist_allgather_fn(scales.to(get_accelerator().current_device_name()), + quant_scale_buffer, ds_process_group) + quant_info = QuantizationInfo() + quant_info.quantized_param = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to( + param.device) + quant_info.backend = self.quantizer_module + quant_info.quant_handle = quant_handle + quant_info.scale_buffer = quant_scale_buffer + return AllGatherHandle(handle, param, quantization=quant_info) + + else: + if self.use_all_reduce_for_fetch_params and not quantize and not use_secondary_tensor: + # Use all_reduce instead of all_gather to fetch the module params + flat_buffer_size = sum(p.ds_numel_aligned for p in params) + flat_tensor = torch.zeros(flat_buffer_size, + dtype=get_only_unique_item(p.ds_tensor.dtype for p in params), + device=get_accelerator().current_device_name(), + requires_grad=False) + start_param = 0 + for param in params: + param.data = flat_tensor.narrow(0, start_param, param.ds_numel).view(param.ds_shape) + start = start_param + param.ds_tensor.ds_numel * self.get_partition_rank() + flat_tensor.narrow(0, start, param.ds_tensor.ds_numel).copy_(param.ds_tensor) + + start_param += param.ds_numel + + handle = dist.all_reduce(flat_tensor, group=ds_process_group, async_op=True) + + return AllReduceCoalescedHandle(handle=handle, params=params) + else: + if not quantize: + dtype_params = defaultdict(list) + for p in params: + dtype_params[p.ds_tensor.dtype].append(p) + handles = [] + for dtype, params in dtype_params.items(): + handles.append( + _all_gather_dtype(dtype, params, world_size, rank_in_group, ds_process_group)) + + return MultipleAllGatherHandles(handles) + + else: + partition_sz = sum(p.ds_tensor.ds_numel for p in params) + + if use_secondary_tensor: + partition_sz = sum(p.ds_tensor.ds_numel * p.ds_secondary_tensor_num_of_groups + for p in params) + + flat_tensor = torch.empty(partition_sz * world_size, + dtype=torch.int8, + device=get_accelerator().current_device_name(), + requires_grad=False) + + if use_secondary_tensor: + if hasattr(params[0].ds_secondary_tensor, "ds_quant_scale"): + quantized_param = instrument_w_nvtx(torch.cat)([ + p.ds_secondary_tensor.data.to(get_accelerator().current_device_name()) + for p in params + ]) + scales = instrument_w_nvtx(torch.cat)([ + p.ds_secondary_tensor.ds_quant_scale.to(get_accelerator().current_device_name()) + for p in params + ]) + else: + quantized_param, scales = self.quantizer_module.quantize( + instrument_w_nvtx(torch.cat)([ + p.ds_secondary_tensor.to(get_accelerator().current_device_name()) + for p in params + ])) + else: + if hasattr(params[0].ds_tensor, "ds_quant_scale"): + quantized_param = instrument_w_nvtx(torch.cat)( + [p.ds_tensor.data.to(get_accelerator().current_device_name()) for p in params]) + scales = instrument_w_nvtx(torch.cat)([ + p.ds_tensor.ds_quant_scale.to(get_accelerator().current_device_name()) + for p in params + ]) + else: + quantized_param, scales = self.quantizer_module.quantize( + instrument_w_nvtx(torch.cat)( + [p.ds_tensor.to(get_accelerator().current_device_name()) for p in params])) + quant_scale_buffer = torch.empty( + scales.numel() * world_size, + dtype=torch.float32, + device=get_accelerator().current_device_name(), + requires_grad=False, + ) + handle = _dist_allgather_fn(quantized_param, flat_tensor, ds_process_group) + quant_handle = _dist_allgather_fn(scales, quant_scale_buffer, ds_process_group) + quant_info = QuantizationInfo() + quant_info.quantized_param = flat_tensor + quant_info.backend = self.quantizer_module + quant_info.quant_handle = quant_handle + quant_info.scale_buffer = quant_scale_buffer + quant_info.partition_sz = partition_sz + quant_info.world_size = world_size + return AllGatherCoalescedHandle( + allgather_handle=handle, + params=params, + partitions=None, + world_size=world_size, + use_secondary_tensor=use_secondary_tensor, + quantization=quant_info, + ) + + def partition(param_list=None, hierarchy=0, has_been_updated=False, free_data=True): + cls = param + print_rank_0(f"{'--'*hierarchy}----Partitioning param {debug_param2name_id_shape_device(cls)}", + force=False) + if param_list is None: + param_list = [cls] + self._partition(param_list, has_been_updated=has_been_updated, free_data=True) + + def reduce_gradients_at_owner(param_list=None, hierarchy=0): + cls = param + if param_list is None: + param_list = [cls] + print_rank_0( + f"{'--'*hierarchy}----Reducing Gradients for param with ids {[param.ds_id for param in param_list]} to owner" + ) + self._reduce_scatter_gradients(param_list) + + def partition_gradients(param_list=None, partition_buffers=None, hierarchy=0, accumulate=False): + cls = param + print_rank_0( + f"{'--'*hierarchy}----Partitioning param gradient with id {debug_param2name_id_shape_device(cls)}") + if param_list is None: + param_list = [cls] + if isinstance(partition_buffers, torch.Tensor): + partition_buffers = [partition_buffers] + + self._partition_gradients(param_list, partition_buffers=partition_buffers, accumulate=accumulate) + + def aligned_size(): + return self._aligned_size(param) + + def padding_size(): + return self._padding_size(param) + + def partition_numel(): + return self._partition_numel(param) + + def item_override(): + param.all_gather() + return param._orig_item() + + def ds_summary(slf: torch.Tensor, use_debug_name: bool = False) -> dict: + return { + "id": debug_param2name_id(slf) if use_debug_name else slf.ds_id, + "status": slf.ds_status.name, + "numel": slf.numel(), + "ds_numel": slf.ds_numel, + "shape": tuple(slf.shape), + "ds_shape": tuple(slf.ds_shape), + "requires_grad": slf.requires_grad, + "grad_shape": tuple(slf.grad.shape) if slf.grad is not None else None, + "persist": slf.ds_persist, + "active_sub_modules": slf.ds_active_sub_modules, + "ds_tensor.shape": slf.ds_tensor.shape if slf.ds_tensor is not None else None + } + + def convert_to_zero_parameters(param_list): + self._convert_to_zero_parameters(param_list) + + def allgather_before(func: Callable) -> Callable: + + def wrapped(*args, **kwargs): + param.all_gather() + return func(*args, **kwargs) + + return wrapped + + # Collectives for gathering and partitioning parameters + param.all_gather = all_gather + param.all_gather_coalesced = all_gather_coalesced + param.partition = partition + + # Collective for averaging gradients + param.reduce_gradients_at_owner = reduce_gradients_at_owner + param.partition_gradients = partition_gradients + + # Partitioning size utilities + param.aligned_size = aligned_size + param.padding_size = padding_size + param.partition_numel = partition_numel + param.ds_summary = types.MethodType(ds_summary, param) + + param.item = allgather_before(param.item) + + param.convert_to_zero_parameters = convert_to_zero_parameters + + def _aligned_size(self, param): + return param.ds_numel + self._padding_size(param) + + def _padding_size(self, param): + remainder = param.ds_numel % self.num_partitions + return (self.num_partitions - remainder) if remainder else 0 + + def _partition_numel(self, param): + return param.ds_tensor.ds_numel + + def _ensure_availability_of_partitioned_params(self, params): + swap_in_list = [] + swap_in_flight = [] + for param in params: + if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE: + assert param.ds_tensor.final_location == OffloadDeviceEnum.nvme and param.ds_status == ZeroParamStatus.NOT_AVAILABLE + swap_in_list.append(param) + if param.ds_tensor.status == PartitionedParamStatus.INFLIGHT: + assert param.ds_tensor.final_location == OffloadDeviceEnum.nvme and param.ds_status == ZeroParamStatus.NOT_AVAILABLE + swap_in_flight.append(param) + if len(swap_in_list) > 0: + swap_in_list[0].nvme_swapper.swap_in(swap_in_list, async_op=False) + elif len(swap_in_flight) > 0: + swap_in_flight[0].nvme_swapper.synchronize_reads() + + @instrument_w_nvtx + def _all_gather(self, param_list, async_op=False, hierarchy=None): + + # fetches from nvme if the partition is not available and in nvme + self._ensure_availability_of_partitioned_params(param_list) + + handles = [] + all_gather_list = [] + for param in param_list: + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + if async_op: + handle = self._allgather_param(param, async_op=async_op, hierarchy=hierarchy) + param.ds_status = ZeroParamStatus.INFLIGHT # if async_op else ZeroParamStatus.AVAILABLE + handles.append(handle) + else: + all_gather_list.append(param) + # note: param_list may contain params that are already in flight / aviailable. So we need to use all_gather_list + if not async_op: + if len(all_gather_list) == 1: + ret_value = self._allgather_params(all_gather_list, hierarchy=hierarchy) + else: + all_gather_quantize_list = [] + all_gather_nonquantize_list = [] + for param in all_gather_list: + if hasattr(param.ds_tensor, + "ds_quant_scale") or (hasattr(param, "ds_secondary_tensor") + and hasattr(param.ds_secondary_tensor, "ds_quant_scale")): + all_gather_quantize_list.append(param) + else: + all_gather_nonquantize_list.append(param) + # _allgather_params_coalesced always return None + self._allgather_params_coalesced(all_gather_nonquantize_list, hierarchy, quantize=False) + self._allgather_params_coalesced(all_gather_quantize_list, hierarchy, quantize=True) + for param in all_gather_list: + param.ds_status = ZeroParamStatus.AVAILABLE + return None + + return handles + + def _partition(self, param_list, force=False, has_been_updated=False, free_data=True): + for param in param_list: + print_rank_0(f"Before Partitioning Param {param.ds_id}", force=False) + if self.zero_param_process_group is not None: + self._partition_param_sec(param) + self._partition_param(param, has_been_updated=has_been_updated, free_data=True) + + param.ds_status = ZeroParamStatus.NOT_AVAILABLE + # if param.ds_tensor is not None: + # assert id(param.data) == id(param.ds_tensor.data), \ + # "After the parameters are initially partitioned, make sure we are not recreating the partition." + #print_rank_0(f"After Partitioning Param {param.ds_id} {param.ds_tensor.size()} {param.ds_tensor}",force=False) + @instrument_w_nvtx + def _partition_param(self, param, buffer=None, has_been_updated=False, free_data=True): + assert param.ds_status is not ZeroParamStatus.INFLIGHT, f" {param} Cannot partition a param in flight" + global reuse_buffers + print_rank_0(f"Param id {param.ds_id} status is {param.ds_status}", force=False) + if param.ds_status is ZeroParamStatus.AVAILABLE: + print_rank_0(f"Partitioning param id {param.ds_id} reuse buffers {reuse_buffers}", force=False) + # if reuse_buffers and False: + # numel = buffer.numel() + # buffer = param.data.view(-1) + # print_rank_0( + # "Returning buffer for param {param.ds_id} with numel {param.ds_numel} to empty buffers", + # force=False) + # if numel in empty_buffers: + # empty_buffers[numel].append(buffer) + + # if deepspeed.comm.get_rank(): + # print(f"Releasing {param.data.numel()}") + + if param.ds_tensor is not None and not has_been_updated: ##param already partitioned + + #print_rank_0(f"Param {param.ds_id} pri {param.ds_tensor.size()} loc? {param.ds_tensor.final_location}", force=True) + #param.data = param.ds_tensor.data + + see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=False) + # param.data does not store anything meaningful in partitioned state + if free_data: + free_param(param) + see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=False) + + if param.ds_tensor.final_location == OffloadDeviceEnum.nvme: + print_rank_0(f"Param {param.ds_id} partition released since it exists in nvme", force=False) + param.nvme_swapper.remove_partition_and_release_buffers([param]) + print_rank_0( + f"after swap Param {param.ds_id} {param.ds_tensor.shape} partition released since it exists in nvme", + force=False) + + return + + tensor_size = self._aligned_size(param) + partition_size = tensor_size // self.num_partitions + if param.ds_tensor is None: + final_location = None + if self.remote_device == OffloadDeviceEnum.nvme and self.param_swapper.swappable_tensor( + numel=partition_size): + final_location = OffloadDeviceEnum.nvme + buffer = self.param_swapper.get_buffer(param, partition_size) + partitioned_tensor = torch.empty(0, dtype=param.dtype, device=buffer.device) + partitioned_tensor.data = buffer.data + print_rank_0(f"ID {param.ds_id} Initializing partition for the first time for nvme offload.") + + else: + if param.ds_persist: + device = self.local_device + elif self.remote_device == OffloadDeviceEnum.nvme: + device = OffloadDeviceEnum.cpu + else: + device = self.remote_device + + partitioned_tensor = torch.empty(partition_size, dtype=param.dtype, device=device) + # quantize the tensor if it's not trainable + if not param.requires_grad and self.quantized_nontrainable_weights: + partitioned_tensor, partitioned_tensor.ds_quant_scale = self.quantizer_module.quantize( + partitioned_tensor) + + if device == OffloadDeviceEnum.cpu and self.pin_memory: + partitioned_tensor = get_accelerator().pin_memory(partitioned_tensor) + + partitioned_tensor.requires_grad = False + param.ds_tensor = partitioned_tensor + param.ds_tensor.ds_numel = partition_size + param.ds_tensor.status = PartitionedParamStatus.AVAILABLE + param.ds_tensor.final_location = final_location + param.ds_numel_aligned = tensor_size + + start = partition_size * self.get_partition_rank() + end = start + partition_size + + one_dim_param = param.contiguous().view(-1) + + if start < param.ds_numel and end <= param.ds_numel: + src_tensor = one_dim_param.narrow(0, start, partition_size) + + with torch.no_grad(): + # make sure param.ds_tensor requires_grad always be false, + # otherwise, torch tracer will complain. + param.ds_tensor.copy_(src_tensor) + + #partitioned_tensor = src_tensor.clone().detach().to(self.remote_device) + + else: + # partitioned_tensor = torch.zeros(partition_size, + # dtype=param.dtype, + # device=self.remote_device ) + + if start < param.ds_numel: + elems_to_copy = param.ds_numel - start + with torch.no_grad(): + # make sure param.ds_tensor requires_grad always be false, + # otherwise, torch tracer will complain. + param.ds_tensor.narrow(0, 0, + elems_to_copy).copy_(one_dim_param.narrow(0, start, elems_to_copy)) + + #print(f"Remote device {self.remote_device}") + + #param.ds_tensor = partitioned_tensor + + #param.data = param.ds_tensor.data + + # param.data does not store anything meaningful in partitioned state + + see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=False) + free_param(param) + see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=False) + + if param.ds_tensor.final_location == OffloadDeviceEnum.nvme: + self.param_swapper.swap_out_and_release([param]) + print_rank_0(f"ID {param.ds_id} Offloaded to nvme offload and buffers released.") + see_memory_usage(f"ID {param.ds_id} Offloaded to nvme offload and buffers released.", force=False) + + print_rank_0(f"ID {param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}") + + @instrument_w_nvtx + def _partition_param_sec(self, param, buffer=None, has_been_updated=False): + assert param.ds_status is not ZeroParamStatus.INFLIGHT, f" {param} Cannot partition a param in flight" + global reuse_buffers + ##support for NVME secondary param offload + #print_rank_0(f"SEC Param id {param.ds_id} status is {param.ds_status}", force=True) + if param.ds_status is ZeroParamStatus.AVAILABLE: + if param.ds_secondary_tensor is not None and not has_been_updated: ##param already partitioned + return + #check padding + tensor_size = self._aligned_size(param) + partition_size = tensor_size // self.dp_world_size + + secondary_partition_size = int(tensor_size // self.num_ranks_in_param_group) + if param.ds_secondary_tensor is None: + final_location = None + secondary_partitioned_tensor = torch.empty(secondary_partition_size, + dtype=param.dtype, + device=self.remote_device) + + if self.pin_memory: + secondary_partitioned_tensor = secondary_partitioned_tensor.pin_memory() + # quantize the tensor if it's not trainable + if not param.requires_grad and self.quantized_nontrainable_weights: + secondary_partitioned_tensor, secondary_partitioned_tensor.ds_quant_scale = self.quantizer_module.quantize( + secondary_partitioned_tensor) + secondary_partitioned_tensor.requires_grad = False + param.ds_secondary_tensor = secondary_partitioned_tensor + param.ds_secondary_tensor.ds_numel = secondary_partition_size + param.ds_secondary_tensor.status = PartitionedParamStatus.AVAILABLE + param.ds_secondary_tensor.final_location = final_location + + #use rank in group for secondary tensor + secondary_start = secondary_partition_size * self.rank_in_group + + secondary_end = secondary_start + secondary_partition_size + + one_dim_param = param.contiguous().view(-1) + + # ds_numel is unpadded, so the last chunk of the secondary tensor might not be secondary_partition_size + sec_numel = max(0, min(param.ds_numel - secondary_start, secondary_partition_size)) + + # copy from full tensor to secondary tensor + param.ds_secondary_tensor.narrow(0, 0, + sec_numel).copy_(one_dim_param.narrow(0, secondary_start, sec_numel)) + + # TODO: This is a temporary fix to avoid the issue that 2nd tensor all-gather happens before 2nd tensor partition is done + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().synchronize() + + print_rank_0(f"{param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}", + force=False) + + def _param_status(self, param): + if param.ds_tensor is not None: + print_rank_0( + f"Param id {param.ds_id}, param status: {param.ds_status}, param numel {param.ds_numel}, partitioned numel {param.ds_tensor.numel()}, data numel {param.data.numel()}" + ) + else: + print_rank_0( + f"Param id {param.ds_id}, param status: {param.ds_status}, param numel {param.ds_numel}, partitioned ds_tensor {param.ds_tensor}, data numel {param.data.numel()}" + ) + + def _allgather_param(self, param, async_op=False, hierarchy=0): + + partition_size = param.ds_tensor.ds_numel + + tensor_size = partition_size * self.num_partitions + aligned_param_size = self._aligned_size(param) + assert tensor_size == aligned_param_size, f'param id {param.ds_id} aligned size {aligned_param_size} does not match tensor size {tensor_size}' + + print_rank_0( + f"{'--'* hierarchy}---- Before allocating allgather param {debug_param2name_id_shape_status(param)} partition size={partition_size}" + ) + + see_memory_usage( + f'Before allocate allgather param {debug_param2name_id_shape_status(param)} partition_size={partition_size} ', + force=False) + flat_tensor = torch.zeros(aligned_param_size, dtype=param.dtype, device=param.device).view(-1) + see_memory_usage( + f'After allocate allgather param {debug_param2name_id_shape_status(param)} {aligned_param_size} {partition_size} ', + force=False) + + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + + print_rank_0( + f"{'--'* hierarchy}----allgather param with {debug_param2name_id_shape_status(param)} partition size={partition_size}" + ) + # if not flat_tensor.numel() > 100000: + # replicated_tensor = flat_tensor.narrow(0, + # 0, + # param.ds_numel).view(param.ds_shape) + # param.data = replicated_tensor.data + # return None + if self.use_all_gather_into_tensor: + handle = dist.all_gather_into_tensor(flat_tensor, + param.ds_tensor.to(get_accelerator().device_name()), + group=self.get_partition_dp_group(param), + async_op=async_op) + else: + partitions = [] + for i in range(self.num_partitions): + partitions.append(flat_tensor.narrow(0, partition_size * i, partition_size)) + + if i == dist.get_rank(group=self.get_partition_dp_group(param)): + partitions[i].data.copy_(param.ds_tensor.data, non_blocking=True) + + handle = dist.all_gather(partitions, + partitions[self.get_partition_rank()], + group=self.get_partition_dp_group(param), + async_op=async_op) + + replicated_tensor = flat_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape) + param.data = replicated_tensor.data + return handle + + def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False): + """ blocking call + avoid explicit memory copy in _allgather_params + """ + if len(param_list) == 0: + return + + if self.num_partitions == 1: + handle = _no_gather_coalesced(param_list) + handle.wait() + return None + + # collect local tensors and partition sizes + partition_sizes = [] + local_tensors = [] + if quantize: + quantize_scale_sizes = [] + quantize_scale_tensors = [] + for param in param_list: + partition_sizes.append(param.ds_tensor.ds_numel) + local_tensors.append(param.ds_tensor.to(get_accelerator().device_name())) + if quantize: + quantize_scale_sizes.append(param.ds_tensor.ds_quant_scale.numel()) + quantize_scale_tensors.append(param.ds_tensor.ds_quant_scale.to(get_accelerator().device_name())) + # allocate memory for allgather params + allgather_params = [] + if quantize: + allgather_quantize_scale = [] + for psize in partition_sizes: + tensor_size = psize * self.num_partitions + flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, + device=self.local_device).view(-1) + flat_tensor.requires_grad = False + allgather_params.append(flat_tensor) + if quantize: + for psize in quantize_scale_sizes: + tensor_size = psize * self.num_partitions + flat_tensor = torch.empty(tensor_size, + dtype=param_list[0].ds_tensor.ds_quant_scale.dtype, + device=self.local_device).view(-1) + flat_tensor.requires_grad = False + allgather_quantize_scale.append(flat_tensor) + + # launch + launch_handles = [] + launch_quantize_handles = [] + for param_idx, param in enumerate(param_list): + input_tensor = local_tensors[param_idx].view(-1) + + if self.use_all_gather_into_tensor: + # try the _all_gather_base from Pytorch master + h = dist.all_gather_into_tensor(allgather_params[param_idx], + input_tensor, + group=self.get_partition_dp_group(param), + async_op=True) + if quantize: + quantize_handle = dist.all_gather_into_tensor(allgather_quantize_scale[param_idx], + quantize_scale_tensors[param_idx], + group=self.get_partition_dp_group(param), + async_op=True) + launch_quantize_handles.append(quantize_handle) + else: + output_list = [] + for i in range(self.num_partitions): + psize = partition_sizes[param_idx] + partition = allgather_params[param_idx].narrow(0, i * psize, psize) + output_list.append(partition) + if not get_accelerator().on_accelerator(partition): + logger.warning( + f'param {param_idx}, partition {i} is not on CUDA, partition shape {partition.size()}') + + # back to old all_gather function + h = dist.all_gather(output_list, input_tensor, group=self.get_partition_dp_group(param), async_op=True) + if quantize: + output_scale_list = [] + for i in range(self.num_partitions): + psize = quantize_scale_sizes[param_idx] + partition = allgather_quantize_scale[param_idx].narrow(0, i * psize, psize) + output_scale_list.append(partition) + quant_handle = dist.all_gather(output_scale_list, + quantize_scale_tensors[param_idx], + group=self.get_partition_dp_group(param), + async_op=True) + launch_quantize_handles.append(quant_handle) + launch_handles.append(h) + + # Wait ensures the operation is enqueued, but not necessarily complete. + launch_handles[-1].wait() + if quantize: + for quant_handle in launch_quantize_handles: + quant_handle.wait() + + # assign to param.data (not copy) + for i, param in enumerate(param_list): + gathered_tensor = allgather_params[i] + if quantize: + gathered_tensor = self.quantizer_module.dequantize(gathered_tensor, allgather_quantize_scale[i]) + param.data = gathered_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape).data + + # guarantee the communication to be completed + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + + return None + + @torch.no_grad() + def _allgather_params(self, param_list, hierarchy=0): + if len(param_list) == 0: + return + + partition_size = sum([param.ds_tensor.ds_numel for param in param_list]) + + tensor_size = partition_size * self.num_partitions + flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, device=self.local_device) + partitions = [] + for i in range(self.num_partitions): + start = partition_size * i + + partitions.append(flat_tensor.narrow(0, start, partition_size)) + + if i == self.get_partition_rank(): + offset = 0 + for param in param_list: + param_numel = param.ds_tensor.ds_numel + + partitions[i].narrow(0, offset, param_numel).copy_(param.ds_tensor.data) + + offset += param_numel + + if hasattr(param_list[0], 'ds_quant_scale'): + scale_size = sum([param.ds_tensor.ds_quant_scale.numel() for param in param_list]) + scale_tensor_size = scale_size * self.world_size + flat_scale_tensor = torch.empty(scale_tensor_size, + dtype=param_list[0].ds_tensor.ds_quant_scale.dtype, + device=self.local_device) + scale_partitions = [] + for i in range(self.world_size): + start = scale_tensor_size * i + scale_partitions.append(flat_scale_tensor.narrow(0, start, scale_tensor_size)) + if i == self.rank: + offset = 0 + for param in param_list: + param_scale_numel = param.ds_tensor.ds_quant_scale.ds_numel + + scale_partitions[i].narrow(0, offset, + param_scale_numel).copy_(param.ds_tensor.ds_quant_scale.data) + + offset += param_scale_numel + + dist.all_gather_into_tensor(flat_tensor, + partitions[self.get_partition_rank()], + group=self.get_partition_dp_group(param), + async_op=False) + if hasattr(param_list[0], 'ds_quant_scale'): + dist.all_gather(flat_scale_tensor, + param_list[0].ds_quant_scale, + group=self.get_partition_dp_group(param), + async_op=False) + param_offset = 0 + + for param in param_list: + param_partition_size = param.ds_tensor.ds_numel + param_size = param.ds_numel + replicated_tensor = torch.empty(param.ds_shape, dtype=param.ds_tensor.dtype, device=self.local_device) + + for i in range(self.num_partitions): + + start = i * partition_size + + param_start = i * param_partition_size + + if param_start < param_size: + numel_to_copy = min(param_size - param_start, param_partition_size) + + part_to_copy = partitions[i].narrow(0, param_offset, numel_to_copy) + + replicated_tensor.view(-1).narrow(0, param_start, numel_to_copy).copy_(part_to_copy) + #param_offset += param.data.numel() + param_offset += param.ds_tensor.ds_numel + if hasattr(param_list[0], 'ds_quant_scale'): + replicated_tensor = self.quantizer_module.dequantize(replicated_tensor, flat_scale_tensor) + param.data = replicated_tensor.data + + return None + + def _reduce_scatter_gradients(self, param_list): + #print_rank_0([param.grad for param in param_list]) + #assert any([param.grad is None for param in param_list]), "None gradients cannot be reduce scattered" + + handles_and_reduced_partitions = [] + for param in param_list: + assert param.grad.numel( + ) == param.ds_numel, f"{param.grad.numel()} != {param.ds_numel} Cannot reduce scatter gradients whose size is not same as the params" + + handles_and_reduced_partitions.append(self._reduce_scatter_gradient(param)) + + for param, (handle, reduced_partition) in zip(param_list, handles_and_reduced_partitions): + if handle is not None: + handle.wait() + + # some ranks may have partitions that are padded to go beyond the grad size. + # For these ranks the output of reduce scatter is a separate buffer and needs + # to be copied in + partition_size = param.ds_tensor.ds_numel + start = self.get_partition_rank() * partition_size + end = start + partition_size + #print_rank_0("REduce scatter was executed for param {param.ds_id}") + if start < param.ds_numel < end: + elements = param.ds_numel - start + param.grad.view(-1).narrow(0, start, elements).copy_(reduced_partition.narrow(0, 0, elements)) + + def _reduce_scatter_gradient(self, param): + + partition_size = param.ds_tensor.ds_numel + #output = torch.empty(partition_size, dtype=param.dtype, device=param.device) + + total_size = partition_size * self.num_partitions + input_list = [] + + for i in range(self.num_partitions): + + start = i * partition_size + end = start + partition_size + + #print("before reduce scatter gradients") + if start < param.ds_numel and end <= param.ds_numel: + input = param.grad.view(-1).narrow(0, start, partition_size) + else: + input = torch.zeros(partition_size, dtype=param.dtype, device=param.device) + + if start < param.ds_numel: + elements = param.ds_numel - start + input.narrow(0, 0, elements).copy_(param.grad.view(-1).narrow(0, start, elements)) + #print("after reduce scatter gradients") + input_list.append(input) + + rank = dist.get_rank(group=self.get_partition_dp_group(param)) + handle = dist.reduce_scatter(input_list[rank], + input_list, + group=self.get_partition_dp_group(param), + async_op=True) + + return handle, input_list[rank] + + def _partition_gradients(self, param_list, partition_buffers=None, accumulate=False): + if partition_buffers is None: + partition_buffers = [None] * len(param_list) + + for param, partition_buffer in zip(param_list, partition_buffers): + self._partition_gradient(param, partition_buffer=partition_buffer, accumulate=accumulate) + + def _partition_gradient(self, param, partition_buffer=None, accumulate=False): + + #import pdb;pdb.set_trace() + # param.grad=None + # param.grad.test() + print_rank_0( + f"Partitioning param {param.ds_id} gradient of size {param.grad.numel()} type {param.grad.dtype} part_size {param.ds_tensor.ds_numel}" + ) + see_memory_usage("Before partitioning gradients", force=False) + partition_size = param.ds_tensor.ds_numel + + if partition_buffer is None: + assert not accumulate, "No buffer to accumulate to" + partition_buffer = torch.zeros(partition_size, dtype=param.dtype, device=param.device) + else: + assert partition_buffer.numel( + ) >= partition_size, f"The partition buffer size {partition_buffer.numel()} should match the size of param.ds_tensor {partition_size}" + + rank = dist.get_rank(group=self.get_partition_dp_group(param)) + start = partition_size * rank + end = start + partition_size + + dest_tensor_full_buffer = partition_buffer.view(-1).narrow(0, 0, partition_size) + + #print("before partition gradients") + if start < param.ds_numel: + elements = min(param.ds_numel - start, partition_size) + + dest_tensor = dest_tensor_full_buffer.narrow(0, 0, elements) + src_tensor = param.grad.view(-1).narrow(0, start, elements) + + # just copy the grad partition to the buffer + if not accumulate: + dest_tensor.copy_(src_tensor) + + # if source and destination are on same device, + # add to the provided buffer + elif src_tensor.device == dest_tensor.device: + dest_tensor.add_(src_tensor) + + # if source and destination are on different device, copy first to src + # then add and move back to the destination. This seems to run faster + # when src is gpu and dest is cpu + # adding directly to cpu is very slow + else: + acc_tensor = torch.empty(src_tensor.numel(), dtype=param.dtype, device=param.device) + + acc_tensor.copy_(dest_tensor) + acc_tensor.add_(src_tensor) + dest_tensor.copy_(acc_tensor) + + # partition_buffer.view(-1).narrow( + # 0, + # 0, + # elements).copy_(param.grad.view(-1).narrow(0, + # start, + # elements)) + + #print("after partition gradients") + param.grad.data = dest_tensor_full_buffer.data + see_memory_usage("After partitioning gradients", force=False) + + def get_partition_dp_group(self, param): + return param.ds_process_group + + def get_partition_rank(self): + """subclass can overload to specify different relative rank in + parameter partition group""" + return self.rank + + @property + def num_partitions(self): + return self.dp_world_size + + def get_dp_process_group(self): + """ Return the communication group with all data-parallel ranks """ + return self.ds_process_group + + +class GatheredParameters: + + def __init__(self, params, modifier_rank=None, fwd_module=None, enabled=True): + """A context that collects parameters that were partitioned via a + :class:`deepspeed.zero.Init` context. The parameters are partitioned + again upon exit. + + Args: + params (``torch.nn.Parameter``): A single parameter, or an iterable of parameters (list, tuple, generator) of parameters to collect. + It's assumed that all parameters are zero params. + modifier_rank (int, optional): If specified, this rank's parameter will be + broadcasted on exit from the context. This argument is required if ``params`` are + modified, so that all processes have a consistent view of the data. Defaults + to ``None``. + fwd_module (``torch.nn.Module``, optional): If specified, ``params`` will be + registered as external parameters of ``fwd_module``. See :meth:`deepspeed.zero.register_external_parameter`. + enabled (bool, optional): If ``False``, this context is a no-op. Defaults to ``True``. + + Important: Make sure to use ``modifier_rank`` that is not ``None`` (e.g., ``modifier_rank=0``) + if you need the GPU memory allocated by gather to be released upon exit from the context manager. + + Important: if ``params`` isn't an iterable of parameters or a single parameter it'll be silently ignored! + + Examples + ======== + + #. Allocate a partitioned module, initialize its weight on rank 0, and update all + processes. + + .. code-block:: python + + with deepspeed.zero.Init(): + linear = torch.nn.Linear(1000,1000) + + with deepspeed.zero.GatheredParameters(linear.weight, + modifier_rank=0): + if deepspeed.comm.get_rank() == 0: + linear.weight.zero_() + + with deepspeed.zero.GatheredParameters(linear.weight, + modifier_rank=0): + if deepspeed.comm.get_rank() == 0: + linear.weight.zero_() + + #. Collect a partitioned weight to pass to another module during + training. The parameter will be registered as an external parameter + and made available during the backward pass. + + .. code-block:: python + :emphasize-lines: 6 + + def forward(self, input): + x = self.layer1(input) + + # self.layer1.weight is required by self.layer2.forward + with deepspeed.zero.GatheredParameters(self.layer1.weight, + fwd_module=self): + y = self.layer2(x, self.layer1.weight) + return y + + + #. Pretrained model loading + + .. code-block:: python + + with deepspeed.zero.Init(): + model = MyModel() + + state_dict = torch.load(model_path, map_location="cpu") + + def load(module: nn.Module, prefix=""): + # because zero3 puts placeholders in model params, this context + # manager gathers (unpartitions) the params of the current layer, then loads from + # the state dict and then re-partitions them again + with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0): + if deepspeed.comm.get_rank() == 0: + module._load_from_state_dict(state_dict, prefix) + + for name, child in module._modules.items(): + if child is not None: + load(child, prefix + name + ".") + + load(model, prefix="") + + If this approach is not used, then the full model will first be copied to each GPU. For models + bigger than the memory of a single GPU, this method is required. + """ + + self.enabled = enabled + if not enabled: + return + + if isinstance(params, Iterable) and not isinstance(params, torch.Tensor): + # deal with generators like model.parameters() + # must convert to list to be able to iterate more than once if we get a generator + params = list(params) + else: + # single param + params = [params] + # enable if at least one is zero-param, otherwise a noop + if not any(is_zero_param(p) for p in params): + self.enabled = False + return + + self.params = [p for p in params if hasattr(p, "ds_id")] + self.params = sorted( + set(self.params), key=lambda x: x.ds_id + ) # remove the duplicates to prevent racing condition, we must also make sure the order is the same on all ranks otherwise we'll get deadlocks + self.src_rank = None + if modifier_rank is not None: + if self.params[0].ds_process_group == dist.get_world_group(): + self.src_rank = modifier_rank + else: + # A group was specified; convert DP rank to global rank + self.src_rank = dist.get_global_rank(self.params[0].ds_process_group, modifier_rank) + self.fwd_module = fwd_module + if self.fwd_module is not None: + # is a no-op if already registered + for p in self.params: + register_external_parameter(self.fwd_module, p) + + def __enter__(self): + if not self.enabled: + return + self.params[0].all_gather(param_list=self.params) + + def __exit__(self, *exc): + if not self.enabled: + return + if self.src_rank is None: + self.params[0].partition(param_list=self.params, has_been_updated=False) + return + + handles = [dist.broadcast(p.data, self.src_rank, group=p.ds_process_group, async_op=True) for p in self.params] + for h in handles: + h.wait() + self.params[0].partition(param_list=self.params, has_been_updated=True) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py new file mode 100644 index 0000000000000000000000000000000000000000..80b20cb1d3e883d91a5dce95a16a496c42b9faef --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py @@ -0,0 +1,585 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from dataclasses import dataclass +import collections +from collections import UserDict +from typing import Deque, Set + +from deepspeed import comm as dist +from deepspeed.utils import z3_leaf_module +from deepspeed.utils.logging import logger +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.runtime.zero.partition_parameters import * +from deepspeed.runtime.zero.partitioned_param_profiler import PartitionedParameterProfiler +from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus +from deepspeed.utils.debug import debug_param2name_id_shape +from deepspeed.accelerator import get_accelerator +import deepspeed.runtime.compiler as compiler +from deepspeed.runtime.compiler import is_compiling + +import logging + +ENABLE_PROFILER = False + + +def debug_rank0(message: str) -> None: + if dist.get_rank() == 0: + logger.debug(message) + + +@instrument_w_nvtx +def get_all_parameters(sub_module, recurse=False): + return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters()) + + +@compiler.disable +def iter_params(module: Module, recurse=False) -> Iterable[Parameter]: + return map(lambda pair: pair[1], get_all_parameters(module, recurse)) + + +class ZeRoTraceMode(Enum): + # Record trace of the network during a single forward+backward (for training) or forward (for inference) + RECORD = 1 + # Use recorded network trace to optimize current forward+backward or forward + COMPLETE = 2 + # Recorded trace does not match current forward+backward or forward pass. + INVALID = 3 + + +class InflightParamRegistry(UserDict): + """registry for parameters in flight""" + + def __setitem__(self, param: Parameter, handle: AllGatherCoalescedHandle) -> None: + if param in self.data: + raise RuntimeError(f"{param.ds_summary()} already in registry") + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"attempted to add non-inflight parameter to registry {param.ds_summary()}") + self.data[param] = handle + + +class PartitionedParameterCoordinator: + FORWARD_FETCH_SUBMIT = 'forward_fetch_submit' + FORWARD_FETCH_WAIT = 'forward_fetch_wait' + FORWARD_PREFETCH_SUBMIT = 'forward_prefetch_submit' + BACKWARD_FETCH_SUBMIT = 'backward_fetch_submit' + BACKWARD_FETCH_WAIT = 'backward_fetch_wait' + BACKWARD_PREFETCH_SUBMIT = 'backward_prefetch_submit' + FORWARD_ALL_GATHER = 'forward_all_gather' + BACKWARD_ALL_GATHER = 'backward_all_gather' + """Handles partitioning and gathering of parameters.""" + + @dataclass + class __ParamInTrace: + param: Parameter + step_id_last_used_at: int + + def __init__( + self, + prefetch_bucket_sz: int, + max_reuse_distance_in_numel: int, + max_available_parameters_in_numel: int, + allgather_stream: get_accelerator().Stream, + inflight_param_registry: InflightParamRegistry, + prefetch_nvme: bool = False, + timers=None, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + fast_sharding_for_leaf_module=False, + log_trace_cache_warnings=False, + ) -> None: + # mapping of param -> handle for each param that is currently in flight + self.__inflight_param_registry = inflight_param_registry + # keeps track of the number of submodules invoked so far. + self.__step_id: int = 0 + # network tracing mode + self.__trace_mode: ZeRoTraceMode = ZeRoTraceMode.INVALID + # sequence of submodules/parameters in forward pass + backward pass + self.__submodule_order: Iterable[Module] = [] + self.__param_order: Iterable[__class__.__ParamInTrace] = [] + self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) + self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque()) + # number of available params, and max number of available params + self.__n_available_params: int = 0 + self.__max_n_available_params: int = max_available_parameters_in_numel + # max distance between two use of the module beyond which module is released + self.__max_reuse_dist_in_numel: int = max_reuse_distance_in_numel + # queue for parameters to fetch. parameters will be popped off the left + # side of the dequeue as they are fetched + self.__param_queue: Deque[__class__.__ParamInTrace] = None + self.__prefetch_bucket_sz: int = prefetch_bucket_sz + self.__prefetch_nvme: bool = prefetch_nvme + self.hierarchy: int = 0 + self.zero_quantized_weights = zero_quantized_weights + self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights + + # stream that will be used for allgather operations + self.__allgather_stream: get_accelerator().Stream = allgather_stream + + # limit the number of fetch events that can be queued at once + # otherwise, what happens is memory is allocated by the host thread at the + # time of the call, but not used until later by the asynchronous cuda stream. + # allowing an infinite number of these to queue up causes a lot of memory + # pressure that then becomes detrimental to performance. + # this is a much less elegant way of fixing this vs something like using + # cudaMallocAsync/cudaFreeAsync. Choosing to not expose this to the user now + # because ideally in the future its replaced by an async allocation + # mechanism which doesn't require any configuration by the user. + self.__ongoing_fetch_events: Deque[get_accelerator().Event] = collections.deque() + # TODO. make this configurable via JSON + self.__max_ongoing_fetch_events: int = 2 + self.__profiler = PartitionedParameterProfiler(timers if ENABLE_PROFILER else None) + + # Whether to log trace cache warnings, e.g. invalidation events + self.__log_trace_cache_warnings = log_trace_cache_warnings + + # whether to enable fast fetch for the z3 leaf module. + # this will improve fetch speed but will not break down leaf module parameters to alleviate memory pressure. + self.fast_sharding_for_leaf_module = fast_sharding_for_leaf_module + + """Tracing and Tracking + TODO. consider performing trace before initializing PartitionedParameterCoordinator + and passing trace results into constructor. This way all the code in here can + just assume that the trace is complete and the results can be entirely + immutable. + + Bookkeeping operations used to track where we are in the forward/backward pass + """ + + def _clear_trace_structures(self) -> None: + self.__submodule_order = [] + self.__param_order = [] + self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) + self.__param_queue = None + + def is_complete_trace(self) -> bool: + return self.__trace_mode == ZeRoTraceMode.COMPLETE + + def is_invalid_trace(self) -> bool: + return self.__trace_mode == ZeRoTraceMode.INVALID + + def is_record_trace(self) -> bool: + return self.__trace_mode == ZeRoTraceMode.RECORD + + def _clean_inflight_param_registry(self) -> None: + for param, handle in self.__inflight_param_registry.items(): + handle.wait() + self.__release_param(param) + self.__inflight_param_registry.clear() + + def _invalidate_trace(self) -> None: + if self.is_invalid_trace(): + raise RuntimeError("attempted to invalidate already invalid trace") + self.__trace_mode = ZeRoTraceMode.INVALID + self._clear_trace_structures() + self._clean_inflight_param_registry() + + def trace_prologue(self, sub_module: Module) -> None: + if self.is_complete_trace(): + # sub_module must match expectation else invalidate trace cache + if len(self.__submodule_order) <= self.__step_id: + print_rank_0( + f"Invalidate trace cache @ step {self.__step_id} and module {sub_module.ds_id}: " + f"cache has only {len(self.__submodule_order)} modules", + force=self.__log_trace_cache_warnings) + self._invalidate_trace() + return + + if sub_module != self.__submodule_order[self.__step_id]: + expected_module_id = self.__submodule_order[self.__step_id].ds_id + print_rank_0( + f"Invalidate trace cache @ step {self.__step_id}: " + f"expected module {expected_module_id}, but got module {sub_module.ds_id}", + force=self.__log_trace_cache_warnings) + self._invalidate_trace() + + @compiler.disable + def record_module(self, sub_module: Module) -> None: + """adds sub module to trace""" + if is_compiling(): + return + + if not self.is_record_trace(): + raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}") + + self.__submodule_order.append(sub_module) + self.__step_id_module_fetched_for[sub_module.ds_id].append(self.__step_id) + + def record_parameters(self, sub_module: Module) -> None: + if is_compiling(): + return + """adds sub module to trace""" + if not self.is_record_trace(): + raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}") + + step_id = self.__step_id_module_fetched_for[sub_module.ds_id].popleft() + for param in sorted(set(iter_params(sub_module, recurse=z3_leaf_module(sub_module))), key=lambda p: p.ds_id): + self.__param_order.append(__class__.__ParamInTrace(param=param, step_id_last_used_at=step_id)) + + def construct_parameter_trace_from_module_trace(self): + """use module trace to construct parameter trace""" + self.__param_order = [] + for sub_module in self.__submodule_order: + self.record_parameters(sub_module) + + @compiler.disable + def reset_step(self) -> None: + """indicate that we have completed one fwd+bwd for the model""" + if is_compiling(): + return + + self._clean_inflight_param_registry() + + if not self.is_complete_trace(): # not self.trace_complete: + # Make sure that recorded submodule orders are identical across ranks + assert_ints_same_as_other_ranks([m.ds_id for m in self.__submodule_order]) + + if self.is_record_trace(): + # Successfully recorded a trace + self.construct_parameter_trace_from_module_trace() + # Make sure that recorded parameter orders are identical across ranks + assert_ints_same_as_other_ranks([p.param.ds_id for p in self.__param_order]) + assert_ints_same_as_other_ranks([p.step_id_last_used_at for p in self.__param_order]) + + self.__submodule_order = tuple(self.__submodule_order) # freeze + self.__param_order = tuple(self.__param_order) # freeze + self.__trace_mode = ZeRoTraceMode.COMPLETE + print_rank_0( + f"completed record trace of {len(self.__submodule_order)} sub modules: {[m.ds_id for m in self.__submodule_order]}", + force=False) + else: + # Enable trace recording for next forward/backward pass + self.__trace_mode = ZeRoTraceMode.RECORD + + else: + if self.__profiler is not None: + self.__profiler.log_events() + + self.__param_queue = collections.deque(self.__param_order) # reset fetch queue + self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) + self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque()) + self.__step_id = 0 + self.__n_available_params = 0 + self.__profiler.reset_events() + + def _dump_params(self, tag, sub_module, params, step_id=None): + if step_id is None: + step_id = self.__step_id + param_names = [debug_param2name_id_shape(p) for p in params] + print_rank_0(f'{tag} step = {step_id} p_names = {param_names}', force=False) + + def _dump_param_ids(self, tag, mod_id, p_ids, step_id=None): + if step_id is None: + step_id = self.__step_id + print_rank_0(f'{tag} mod = {mod_id}, step = {step_id}, p_ids = {p_ids}', force=False) + + """Fetch and Release + Fetching, prefetching, and releasing parameters + """ + + @compiler.disable + @instrument_w_nvtx + @torch.no_grad() + def fetch_sub_module(self, current_submodule: Module, forward: bool) -> None: + """This method does the following (in order): + 1. kick off fetch for parameters in immediately required sub module + 2. kick off fetch for next few parameters we will need later (prefetch) + 3. block on parameters in immediately required sub module + """ + if logger.isEnabledFor(logging.DEBUG): + debug_rank0( + f"{self.__step_id}: M{current_submodule.ds_id}({type(current_submodule).__name__}) P{[p.ds_id for p in iter_params(current_submodule, recurse=z3_leaf_module(current_submodule))]} " + + str({ + "avail": f"{self.__n_available_params:.1e}", + "queue_sz": f"{len(self.__param_queue or [])}", + "inflight": [p.ds_id for p in self.__inflight_param_registry], + })) + + params_to_fetch = set(iter_params(current_submodule, recurse=z3_leaf_module(current_submodule))) + fetch_numel = sum( + [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + + if fetch_numel > 0: + event_name = __class__.FORWARD_FETCH_SUBMIT if forward else __class__.BACKWARD_FETCH_SUBMIT + self._dump_param_ids(event_name, current_submodule.ds_id, + [(p.ds_id, p.ds_shape) + for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + # self._dump_params(event_name, current_submodule, [p for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + + self.__profiler.start_event(event_name) + # kick off all gather for params in the immediately required submodule + #for param in params_to_fetch: + if logger.isEnabledFor(logging.DEBUG): + for param in params_to_fetch: + debug_rank0(f"-fetch: {param.ds_summary()}") + self.__all_gather_params(params_to_fetch, forward) + self.__profiler.stop_event(event_name, fetch_numel) + + wait_numel = 0 + wait_event_name = __class__.FORWARD_FETCH_WAIT if forward else __class__.BACKWARD_FETCH_WAIT + self.__profiler.start_event(wait_event_name) + fast_fetch = self.fast_sharding_for_leaf_module and z3_leaf_module(current_submodule) + # wait for parameters in the immediately needed submodule to become available + for param in params_to_fetch: + param.ds_active_sub_modules.add(current_submodule.ds_id) + if logger.isEnabledFor(logging.DEBUG): + debug_rank0(f"-wait: {param.ds_summary()}") + if param in self.__inflight_param_registry: + wait_numel += param.partition_numel() + with get_accelerator().stream(self.__allgather_stream): + while self.__ongoing_fetch_events and self.__ongoing_fetch_events[0].query(): + self.__ongoing_fetch_events.popleft() + if len(self.__ongoing_fetch_events) > self.__max_ongoing_fetch_events: + self.__ongoing_fetch_events.popleft().synchronize() + + self.__inflight_param_registry.pop(param).wait(handle_dependency=not fast_fetch) + + if not get_accelerator().handles_memory_backpressure() and not fast_fetch: + event = get_accelerator().Event() + event.record() + self.__ongoing_fetch_events.append(event) + + assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary() + if not get_accelerator().resolves_data_dependency(): + get_accelerator().current_stream().wait_stream(self.__allgather_stream) + if fast_fetch: + AllGatherCoalescedHandle.free_buffer() + self.__profiler.stop_event(wait_event_name, wait_numel) + + # kick off parameter prefetches for upcoming modules + # don't prefetch if we dont have a completed model trace + if self.is_complete_trace(): + # go through the parameters we need for the current module and pop them + # off the fetch queue so that they aren't prefetched later. + # if params have already been popped off the fetch queue by earlier + # prefetches we won't look for them here + discarded_from_prefetch_queue = set() + params_not_already_fetched = set( + filter(lambda p: self.__most_recent_step_id_param_fetched_for[p] < self.__step_id, params_to_fetch)) + while self.__param_queue and len(discarded_from_prefetch_queue) < len(params_not_already_fetched): + param_in_trace = self.__param_queue.popleft() + self.__most_recent_step_id_param_fetched_for[ + param_in_trace.param] = param_in_trace.step_id_last_used_at + discarded_from_prefetch_queue.add(param_in_trace.param) + + if discarded_from_prefetch_queue != params_not_already_fetched: + raise RuntimeError( + f"tracing error at step {self.__step_id}: \n" + f"module id: {current_submodule.ds_id}, training: {current_submodule.training}\n" + f"expected the next {len(params_not_already_fetched)} parameters in the " + f"parameter fetch queue to be {tuple(p.ds_summary(use_debug_name=True) for p in params_not_already_fetched)} \n" + f"but got \n {tuple(p.ds_summary(use_debug_name=True) for p in discarded_from_prefetch_queue)}.") + + def _is_currently_on_nvme(param): + if param.nvme_swapper is None: + return False + + return param.ds_tensor.final_location == OffloadDeviceEnum.nvme \ + and param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE + + # kick off all gather for params in the next few submodules (prefetch) + if self.__prefetch_bucket_sz > 0: + max_params_to_prefetch = min(self.__max_n_available_params - self.__n_available_params, + self.__prefetch_bucket_sz) + params_to_prefetch = set() + numel_prefetching = 0 + while self.__param_queue and numel_prefetching < max_params_to_prefetch: + param_in_trace: __class__.__ParamInTrace = self.__param_queue.popleft() + + if _is_currently_on_nvme(param_in_trace.param): + # nvme prefetch is handled elsewhere. Need to break here to preserve fetch order + self.__param_queue.appendleft(param_in_trace) + break + + do_prefetch = param_in_trace.param.ds_status == ZeroParamStatus.NOT_AVAILABLE + if param_in_trace.param in params_to_prefetch: + # Avoid duplicates + do_prefetch = False + + self.__most_recent_step_id_param_fetched_for[param_in_trace.param] = \ + max(self.__most_recent_step_id_param_fetched_for[param_in_trace.param], + param_in_trace.step_id_last_used_at) + + if do_prefetch: + params_to_prefetch.add(param_in_trace.param) + numel_prefetching += param_in_trace.param.ds_numel + + if numel_prefetching > 0: + event_name = __class__.FORWARD_PREFETCH_SUBMIT if forward else __class__.BACKWARD_PREFETCH_SUBMIT + self.__profiler.start_event(event_name) + if logger.isEnabledFor(logging.DEBUG): + for param in params_to_prefetch: + debug_rank0(f"-prefetch: {param.ds_summary()}") + self.__all_gather_params(params_to_prefetch, forward) + self.__profiler.stop_event(event_name, numel_prefetching) + + if self.__prefetch_nvme: + self.__prefetch_nvme_param_partitions() + + self.__step_id += 1 + + @instrument_w_nvtx + @torch.no_grad() + def release_sub_module(self, submodule: Module, forward=False) -> None: + """release the parameters of a sub module, assuming they meet conditions to + be released.""" + #print_rank_0(f"release_sub_module {'fwd' if forward else 'bwd'}: {debug_module2name_id(submodule)}", force=False) + params_to_release = (self.__params_to_release(submodule, self.__step_id) if self.is_complete_trace() else set( + p.ds_id for p in iter_params(submodule, recurse=z3_leaf_module(submodule)))) + + free_data = not z3_leaf_module(submodule) or not self.fast_sharding_for_leaf_module + if not free_data: + # wait for the computation to finish and launch as early as possible. + empty_buffer = torch.empty(1, device=get_accelerator().current_device()) + + for param in iter_params(submodule, recurse=z3_leaf_module(submodule)): + param.ds_active_sub_modules.discard(submodule.ds_id) + if param.ds_id in params_to_release and not param.is_external_param: + self.__release_param(param, free_data) + if not free_data: + if param.ds_id in params_to_release and not param.is_external_param: + # empty buffer ensures that all computations are complete + param.data = empty_buffer + + @instrument_w_nvtx + @torch.no_grad() + def release_and_reset_all(self, module: Module) -> None: + """release all module parameters""" + for param in iter_params(module, recurse=True): + if param in self.__inflight_param_registry: + self.__inflight_param_registry.pop(param).wait() + + # TODO. make this throw if if there are still active submodules. currently + # there's a hook execution issue + param.ds_active_sub_modules.clear() + self.__release_param(param) + + for param in iter_params(module, recurse=True): + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(f"{param.ds_summary()} expected to be released") + + @instrument_w_nvtx + def __all_gather_params(self, params: Set[Parameter], forward: bool) -> None: + quantized_params = [] + nonquantized_params = [] + for param in params: + if hasattr(param.ds_tensor, 'ds_quant_scale'): + quantized_params.append(param) + else: + nonquantized_params.append(param) + if quantized_params: + self.__all_gather_params_(quantized_params, forward, quantize=True) + if nonquantized_params: + self.__all_gather_params_(nonquantized_params, forward, quantize=self.zero_quantized_weights) + + def __all_gather_params_(self, params: Set[Parameter], forward: bool, quantize: bool = False) -> None: + """for each partitioned parameter, kick off an async allgather and store + the work handle for the in flight parameters.""" + partitioned_params = [] + all_gather_numel = 0 # numel = num of elements + for param in params: + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + partitioned_params.append(param) + all_gather_numel += param.ds_numel + + if partitioned_params: + self.__n_available_params += all_gather_numel + # here we need to handle a special case where some of the parameters have a valid hpz secondary tensor (e.g. they are not trainable so their secondary tensor never expire) but others do not. + partitioned_params_with_secondary_tensors = [ + p for p in partitioned_params if p.ds_secondary_tensor is not None + ] + partitioned_params_without_secondary_tensors = [ + p for p in partitioned_params if p.ds_secondary_tensor is None + ] + for param_group in [ + partitioned_params_with_secondary_tensors, partitioned_params_without_secondary_tensors + ]: + if not param_group: + continue + with get_accelerator().stream(self.__allgather_stream): + event_name = __class__.FORWARD_ALL_GATHER if forward else __class__.BACKWARD_ALL_GATHER + self.__profiler.start_event(event_name) + handle = param_group[0].all_gather_coalesced(param_group, quantize=quantize) + self.__profiler.stop_event(event_name, all_gather_numel) + for param in param_group: + assert param.ds_status == ZeroParamStatus.INFLIGHT, param.ds_summary() + self.__inflight_param_registry[param] = handle + + # Release swap buffers for persisted params on nvme since they will never be partitioned or evicted from GPU + swap_persisted_params = [ + p for p in partitioned_params if p.ds_persist and p.ds_tensor.final_location == OffloadDeviceEnum.nvme + ] + if swap_persisted_params: + swap_persisted_params[0].nvme_swapper.remove_partition_and_release_buffers(swap_persisted_params) + + @compiler.disable + @instrument_w_nvtx + def __release_param(self, param: Parameter, free_data: bool = True) -> None: + if param.ds_status == ZeroParamStatus.AVAILABLE and not param.ds_active_sub_modules: + if logger.isEnabledFor(logging.DEBUG): + debug_rank0(f"-release: {param.ds_summary()}") + print_rank_0(f"release: {debug_param2name_id_shape(param)}", force=False) + param.partition(free_data=free_data) + self.__n_available_params -= param.ds_numel + + @instrument_w_nvtx + @functools.lru_cache(maxsize=None) + def __params_to_release(self, submodule_to_release: Module, step_id: int) -> Set[int]: + if not self.is_complete_trace(): + raise RuntimeError("expected trace to be complete") + + params_to_release = set( + p.ds_id for p in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release)) + if not p.ds_persist) + + # Problem: When prefetcher scans the param trace, it skips AVAILABLE params. + # This creates issues if those params are released before the skipped uses: + # 1) It hurts performance as the skipped uses are never prefetched. + # 2) For nvme params, we run out of swap buffers because the prefetch order + # diverges from the trace. + # Solution: Don't release params whose reuse was skipped by prefetch. This is + # possible because we detect such skips during prefetch and mark those params. + for param in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release)): + if self.__most_recent_step_id_param_fetched_for[param] > step_id: + params_to_release.discard(param.ds_id) + + # examine all modules within `max_reuse_dist_in_numel` of the current step, + # if we see any of the candidate parameters to be released reoccur while + # doing this, remove them from the set of parameters to release. + params_traversed = 0 + for module in self.__submodule_order[step_id:]: + if params_traversed >= self.__max_reuse_dist_in_numel: + break + for param in iter_params(module, recurse=z3_leaf_module(submodule_to_release)): + params_to_release.discard(param.ds_id) + params_traversed += param.ds_numel + + return params_to_release + + @instrument_w_nvtx + def __prefetch_nvme_param_partitions(self) -> None: + """swap in parameter partitions from nvme for those parameters that will be used + after the ones that are already being prefetched into full parameters + """ + if not self.is_complete_trace(): + return + + numel_in_flight = sum(param.ds_numel for param in self.__inflight_param_registry) + + numel_considered = 0 + swap_in_params = [] + for param_in_trace in self.__param_queue: + param = param_in_trace.param + if param.nvme_swapper is None: + continue + if (numel_considered > 2 * numel_in_flight + or len(swap_in_params) >= param.nvme_swapper.available_swap_in_buffers()): + break + if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE: + swap_in_params.append(param) + numel_considered += param.ds_numel + + if swap_in_params: + swap_in_params[0].nvme_swapper.swap_in(swap_in_params, async_op=True) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..b4ea11f3b83632a5f7f74343ef7092af6c5bf3a7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from dataclasses import dataclass +from deepspeed.utils import log_dist + + +class PartitionedParameterProfiler(object): + + @dataclass + class EventCounter: + name: str + count: int + num_elem: int + + def reset(self): + self.count = 0 + self.num_elem = 0 + + def increment(self, numel): + self.count += 1 + self.num_elem += numel + + def __init__(self, timers): + self.timers = timers + self.event_counters = {} + + def reset_events(self): + for event_ctr in self.event_counters.values(): + event_ctr.reset() + + def start_event(self, name): + if self.timers is None: + return + + if name not in self.event_counters: + self.event_counters[name] = __class__.EventCounter(name=name, count=0, num_elem=0) + self.timers(name).start() + + def stop_event(self, name, num_elem): + if self.timers is None: + return + assert name in self.event_counters, f'unknown event {name}' + self.event_counters[name].increment(num_elem) + self.timers(name).stop() + + def _log_timers(self): + if self.timers is None: + return + self.timers.log(names=list(self.event_counters.keys())) + + def _log_event_counters(self): + for event_ctr in self.event_counters.values(): + log_dist( + f'{event_ctr.name}: count = {event_ctr.count}, numel = {event_ctr.num_elem}', + #f'{event_ctr.name}: time = {self._log_timers()},count = {event_ctr.count}, numel = {event_ctr.num_elem}', + ranks=[0]) + + def log_events(self): + self._log_event_counters() + self._log_timers() diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/stage3.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/stage3.py new file mode 100644 index 0000000000000000000000000000000000000000..7ecefc6fdc4173978e12ee46a5b88394d3021164 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/stage3.py @@ -0,0 +1,3245 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import gc +import collections +import itertools +from typing import Deque, Dict, Set, Tuple, Container +from contextlib import contextmanager + +from deepspeed import comm as dist +from deepspeed.utils import groups, z3_leaf_parameter + +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed.runtime.base_optimizer import ZeROOptimizer +from deepspeed.utils import logger +from deepspeed.utils.torch import register_grad_hook +from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler +from deepspeed.runtime.comm.coalesced_collectives import reduce_scatter_coalesced, all_to_all_quant_reduce, all_to_all_loco_quant_reduce +from deepspeed.runtime.utils import inf, is_model_parallel_parameter, get_only_unique_item, mask_nan_or_inf_with_val_inplace +from deepspeed.runtime.zero.partition_parameters import * +from deepspeed.runtime.zero.config import ZeroStageEnum +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum +from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload +from deepspeed.runtime.zero.utils import apply_to_tensors_only, get_mapping_to_flat_buffer +from deepspeed.runtime.zero.offload_states import offload_adam_states, reload_adam_states +from deepspeed.ops.adam import DeepSpeedCPUAdam +from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus +from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper +from deepspeed.runtime.swap_tensor.partitioned_optimizer_swapper import PartitionedOptimizerSwapper +from deepspeed.runtime.swap_tensor.pipelined_optimizer_swapper import PipelinedOptimizerSwapper +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FP32_FLAT_GROUPS, PARTITION_COUNT, ZERO_STAGE, LOSS_SCALER +from deepspeed.accelerator import get_accelerator + +# Toggle this to true to enable correctness test +# with gradient partitioning and without +pg_correctness_test = False + +OPTIMIZER_SWAP_IN_STATE_TIMER = 'optimizer_swap_in_state' +INIT_OPTIMIZER_TIMER = 'init_optimizer_state' +OPTIMIZER_SWAP_OUT_STATE_TIMER = 'optimizer_swap_out_state' +OPTIMIZER_STEP_TIMER = 'optimizer_step' + + +def print_rank_0(message, debug=False, force=False): + rank = dist.get_rank() + if rank == 0 and (debug or force): + logger.info(message) + # other variations + # - print for all ranks w/o interleaving + # printflock(f"[{rank}] {message}") + # - print to log file per rank + # log_rank_file(rank, message) + + +def input(msg): + return + + +def isclose(a, b, rtol=1e-09, atol=0.0): + return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol) + + +def lcm(x, y): + from fractions import gcd # or can import gcd from `math` in Python 3 + return x * y // gcd(x, y) + + +def move_to_cpu(tensor_list): + for tensor in tensor_list: + tensor.data = tensor.data.cpu() + + +@contextmanager +def unwrap_model_for_generation(model): + """ + For ZeRO-3 models, we gather the weights once to speed up generation. + """ + with GatheredParameters(model.parameters()): + # Removes the optimizer hooks from a DeepSpeed ZeRO-3 model. + + # Remove hooks + if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): + optimizer_offload = model.optimizer.parameter_offload + elif model.optimizer is not None: + optimizer_offload = model.optimizer + + for hook in optimizer_offload.forward_hooks: + hook.remove() + for hook in optimizer_offload.backward_hooks: + hook.remove() + + optimizer_offload.forward_hooks = [] + optimizer_offload.backward_hooks = [] + + yield model + + # Adds the optimizer hooks from a DeepSpeed ZeRO-3 model. + if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): + optimizer_offload = model.optimizer.parameter_offload + elif model.optimizer is not None: + optimizer_offload = model.optimizer + optimizer_offload._register_deepspeed_module(optimizer_offload.module) + return + + +INITIAL_MICRO_STEP_ID = -1 + + +class DeepSpeedZeroOptimizer_Stage3(ZeROOptimizer): + """ + DeepSpeedZeroOptimizer designed to reduce the memory footprint + required for training large deep learning models. + + For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models + https://arxiv.org/abs/1910.02054 + + For usage examples, refer to TODO: DeepSpeed Tutorial + + """ + + def __init__( + self, + module, + init_optimizer, + timers, + ds_config, + static_loss_scale=1.0, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + contiguous_gradients=True, + reduce_bucket_size=500000000, + prefetch_bucket_size=50000000, + max_reuse_distance=1000000000, + max_live_parameters=1000000000, + param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, + dp_process_group=None, + reduce_scatter=True, + overlap_comm=False, + offload_optimizer_config=None, + offload_param_config=None, + sub_group_size=1000000000000, + offload_ratio=0.0, + mpu=None, + clip_grad=0.0, + gradient_accumulation_dtype=torch.float32, + communication_data_type=torch.float16, + postscale_gradients=True, + gradient_predivide_factor=1.0, + gradient_accumulation_steps=1, + elastic_checkpoint=False, + aio_config=None, + all2all_process_group=None, + zero_hpz_partition_size=1, + zero_quantized_weights=False, + zero_quantized_nontrainable_weights=False, + zero_module_granularity_threshold=0, + zeropp_loco_param=None, + log_trace_cache_warnings=False, + ): + see_memory_usage("Stage 3 initialize beginning", force=True) + + print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False) + + if dist.get_rank() == 0: + logger.info(f"Reduce bucket size {reduce_bucket_size}") + logger.info(f"Prefetch bucket size {prefetch_bucket_size}") + # The fused optimizer does all the work. We need this layer for two reason: + # 1. maintain same user API from apex.fp16_utils + # 2. keep common stuff here in case we need to add ne552w fused optimizer later + + # differences from apex.fp16_utils: + # - assume all model params in fp16 + # - assume all params requires grad + # - flat by groups, not keeping state. TODO: remove state explicitly? + # - master grad and unflat master weight never exist. TODO: a way to save out unflat master? + if not get_accelerator().is_available(): + raise SystemError("Cannot use fp16 without accelerator.") + + self.optimizer = init_optimizer + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + self.dtype = self.optimizer.param_groups[0]['params'][0].dtype + self.gradient_accumulation_dtype = gradient_accumulation_dtype + self._global_grad_norm = 0. + + self.custom_loss_scaler = False + self.external_loss_scale = None + + self.optimizer_swapper = None + self.swap_optimizer = False + + self.offload_optimizer = False + self.offload_optimizer_pin_memory = False + self.offload_optimizer_fast_init = False + self.offload_param = False + self.offload_param_pin_memory = False + self.params_in_nvme_and_cpu = False + self.max_params_in_cpu = 0 + self.partial_offload = offload_ratio + + #num of ranks in a ZeRO param partitioning group + self.zero_hpz_partition_size = zero_hpz_partition_size + + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + print_rank_0( + f"ZeRO Stage 3 param partitioning group {self.zero_hpz_partition_size} {zero_param_parallel_group}", + force=False) + if self.zero_hpz_partition_size > 1 and zero_param_parallel_group is None: + self._set_zero_group_parallelism() + zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() + + self.parameter_offload = self.initialize_ds_offload( + module=module, + timers=timers, + ds_config=ds_config, + overlap_comm=overlap_comm, + prefetch_bucket_size=prefetch_bucket_size, + max_reuse_distance=max_reuse_distance, + max_live_parameters=max_live_parameters, + param_persistence_threshold=param_persistence_threshold, + model_persistence_threshold=model_persistence_threshold, + dp_process_group=dp_process_group, + offload_param_config=offload_param_config, + mpu=mpu, + zero_param_parallel_group=zero_param_parallel_group, + zero_quantized_weights=zero_quantized_weights, + zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights, + zero_module_granularity_threshold=zero_module_granularity_threshold, + log_trace_cache_warnings=log_trace_cache_warnings, + ) + + self.persistent_parameters = self.parameter_offload.persistent_parameters + self._configure_offloading(offload_optimizer_config, offload_param_config) + + # backup fused_adam optimizer init + if self.offload_optimizer and self.partial_offload != 1.0: + backup_gpu_tensor = torch.randn(1, device=get_accelerator().device_name()).to(self.dtype) + backup_gpu_param = torch.nn.Parameter(backup_gpu_tensor) + assert type(init_optimizer) == DeepSpeedCPUAdam, 'Hybrid Optimizer Only Supports DeepSpeedCPUAdam' + self.backup_optimizer = torch.optim.AdamW([backup_gpu_param], + lr=self.optimizer.param_groups[0]["lr"], + betas=self.optimizer.param_groups[0]["betas"], + eps=self.optimizer.param_groups[0]["eps"], + weight_decay=self.optimizer.param_groups[0]["weight_decay"], + amsgrad=self.optimizer.param_groups[0]["amsgrad"]) + # Multiple param_groups configs for back-up optimizer + if len(self.optimizer.param_groups) > 1: + for i in range(1, len(self.optimizer.param_groups)): + self.backup_optimizer.add_param_group(self.optimizer.param_groups[i]) + + self.module = module + self.elastic_checkpoint = elastic_checkpoint + + self.device = get_accelerator().current_device_name() if not self.offload_optimizer else OffloadDeviceEnum.cpu + + self.inf_or_nan_tracker: Tensor = torch.zeros(1, dtype=torch.bool, device=self.device, requires_grad=False) + + self.deepspeed_adam_offload = (self.offload_optimizer and type(init_optimizer) == DeepSpeedCPUAdam) + + ### streams used for overlapping computation with communication + self.reduce_and_partition_stream = None if get_accelerator().is_synchronized_device() else get_accelerator( + ).Stream() if overlap_comm else get_accelerator().default_stream() + + ############################################################################ + + self.n_caching_allocator_flushes = 0 + + #-------------Stage 3 Setup-------------------# + + self.timers = timers + + self.all2all_process_group = all2all_process_group + + self.reduce_scatter = reduce_scatter + + self.dp_process_group = self.parameter_offload.dp_process_group + self.sequence_parallel_size = groups._get_sequence_parallel_world_size() + + self.all2all_process_group = all2all_process_group + + self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights + + self.partition_count = dist.get_world_size(group=self.dp_process_group) + + self.zeropp_loco_param = zeropp_loco_param + + if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): + self.model_parallel_group = None + self.model_parallel_rank = 0 + else: + self.model_parallel_group = mpu.get_model_parallel_group() + self.model_parallel_rank = mpu.get_model_parallel_rank() + + self.overflow = False + self.clip_grad = clip_grad + self.communication_data_type = communication_data_type + self.gradient_predivide_factor = gradient_predivide_factor + self.postscale_gradients = postscale_gradients + self.gradient_accumulation_steps = gradient_accumulation_steps + self.micro_step_id = 0 + self.reduce_bucket_size = int(reduce_bucket_size) + + if self.all2all_process_group is not None: + assert self.all2all_process_group is not None and self.reduce_scatter == True, "when enable all_to_all_reduce, reduce_scatter should also be enabled for data type checks." + + if self.reduce_scatter: + valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32) + assert self.communication_data_type in valid_reduce_scatter_dtypes, f"ZeRO-3 supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'" + assert self.gradient_predivide_factor == 1.0, "gradient_predivide_factor != 1.0 is not yet supported with ZeRO-3 with reduce scatter enabled" + assert self.postscale_gradients, "pre-scale gradients is not yet supported with ZeRO-3 with reduce scatter enabled" + + # Holds the mode parameter + # The param.data may not hold any meaningful data + # when param's status is NOT_AVAILABLE or IN_FLGHT + self.fp16_groups = [] + + # Hold partitioned parameters + self.fp16_partitioned_groups = [] + + # Holds a fused and flattened copy of the parameters + self.fp16_partitioned_groups_flat = [] + self.fp16_partitioned_groups_flat_numel = [] + self.fp16_partitioned_groups_flat_id = [] + + #defragmented pinned memory + self.param_groups_fp16_flat_cpu_memory = [] + + #a single 32-bit partition of the parallel partitioned parameters + #that this process will update + self.fp32_partitioned_groups_flat = [] + self.next_swappable_fp32_partitioned_groups = [] + + # number of elements per partition in each group + self.partition_size = [] + + self.all_reduce_print = False + + self.prefetch_elements = int(prefetch_bucket_size) + + self.contiguous_gradients = contiguous_gradients + + # padding on each partition for alignment purposes + self.groups_padding = [] + + self.sub_group_size = sub_group_size + + self.sub_group_to_group_id = {} + + # Trainable parameters + self.trainable_param_groups = self._get_trainable_parameter_groups() + + see_memory_usage("Before creating fp16 partitions", force=True) + self._create_fp16_partitions_with_defragmentation(self.trainable_param_groups) + num_fp16_subgroups = len(self.fp16_partitioned_groups_flat) + see_memory_usage(f"After creating fp16 partitions: {num_fp16_subgroups}", force=True) + + # Optimizer tensor swapping + if self.swap_optimizer: + self._configure_tensor_swapping(offload_optimizer_config, aio_config) + + self.is_gradient_accumulation_boundary: bool = True + + self.param_reduce_events: Deque[get_accelerator().Event] = collections.deque() + # TODO. make this configurable via JSON + self.max_param_reduce_events: int = 2 + + self.param_dict = {} + + # map between param_id and bool to specify if a param is in this partition + self.is_param_in_current_partition = {} + + self.extra_large_param_to_reduce = None + self.grads_in_ipg_bucket = [] + self.params_in_ipg_bucket = [] + + self.params_already_reduced = {} + self._release_ipg_buffers() + self.previous_reduced_grads = None + + # model parameter traversal-based param id that's stable across runs + for params_group in self.fp16_groups: + for param in params_group: + param_id = self.get_param_id(param) + self.param_dict[param_id] = param + self.params_already_reduced[param_id] = False + + #Largest partitioned param + largest_partitioned_param_numel = 0 + for fp16_partitioned_group in self.fp16_partitioned_groups: + if len(fp16_partitioned_group) > 0: + largest_partitioned_param_numel = max( + largest_partitioned_param_numel, + max([max(tensor.numel(), tensor.ds_numel) for tensor in fp16_partitioned_group])) + + print_rank_0(f'Largest partitioned param numel = {largest_partitioned_param_numel}', force=False) + + self._setup_for_real_optimizer() + self.grad_position = {} + self.set_grad_positions() + + if self.offload_optimizer: + self.norm_for_param_grads = {} + + # stores if a partition has been reduced in this step + self.is_partition_reduced = {} + + # stores if a grad in a partition has been computed or not + self.is_grad_computed = {} + + # will store the averaged gradients required by this partition + self.averaged_gradients = {} + + #creates backward hooks for gradient partitioning + ###Calls all gather param + self._grad_acc_hooks = [] + self._leaf_module_hooks = [] + self.create_reduce_and_remove_grad_hooks() + + #exit(0) + + # we may have a way of fusing dynamic scale. Do not support for now + self.loss_scaler = CreateLossScaler(dtype=self.dtype, + static_loss_scale=static_loss_scale, + dynamic_scaling=dynamic_loss_scale, + dynamic_loss_args=dynamic_loss_args) + self.dynamic_loss_scale = self.loss_scaler.dynamic + + self.debug_fp16_grads = [{} for _ in self.fp16_groups] + + self._link_all_hp_params() + + self.offloaded_states: Set(OffloadDeviceEnum) = set() + + if dist.get_rank(group=self.dp_process_group) == 0: + see_memory_usage(f"After initializing ZeRO optimizer", force=True) + + def destroy(self): + self.parameter_offload.destroy() + for hook in self._grad_acc_hooks: + hook.remove() + for hook in self._leaf_module_hooks: + hook.remove() + print_rank_0("Removed grad acc hooks", force=False) + self._release_ipg_buffers() + + def initialize_ds_offload( + self, + module, + timers, + ds_config, + overlap_comm, + prefetch_bucket_size, + max_reuse_distance, + max_live_parameters, + param_persistence_threshold, + model_persistence_threshold, + dp_process_group, + offload_param_config, + mpu, + zero_param_parallel_group, + zero_quantized_weights, + zero_quantized_nontrainable_weights, + zero_module_granularity_threshold, + log_trace_cache_warnings, + ): + return DeepSpeedZeRoOffload(module=module, + timers=timers, + ds_config=ds_config, + overlap_comm=overlap_comm, + prefetch_bucket_size=prefetch_bucket_size, + max_reuse_distance=max_reuse_distance, + max_live_parameters=max_live_parameters, + param_persistence_threshold=param_persistence_threshold, + model_persistence_threshold=model_persistence_threshold, + dp_process_group=dp_process_group, + offload_param_config=offload_param_config, + mpu=mpu, + zero_param_parallel_group=zero_param_parallel_group, + zero_quantized_weights=zero_quantized_weights, + zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights, + zero_module_granularity_threshold=zero_module_granularity_threshold, + log_trace_cache_warnings=log_trace_cache_warnings) + + def _get_trainable_parameter_groups(self): + param_groups = [] + PARAMS_KEY = "params" + for param_group in self.optimizer.param_groups: + trainable_params = [p for p in param_group[PARAMS_KEY] if p.requires_grad] + if len(trainable_params) == 0: + continue + + trainable_param_group = {} + for key in param_group.keys(): + if key == PARAMS_KEY: + trainable_param_group[PARAMS_KEY] = trainable_params + else: + trainable_param_group[key] = param_group[key] + param_groups.append(trainable_param_group) + + return param_groups + + def _set_zero_group_parallelism(self): + groups._create_zero_param_parallel_group(self.zero_hpz_partition_size) + + def invalidate_secondary_tensor(self): + for fpg in self.fp16_groups: + for param in fpg: + if param.ds_secondary_tensor is not None: + param.ds_secondary_tensor = None + + def _setup_for_real_optimizer(self): + see_memory_usage("Before creating fp32 partitions", force=True) + self._create_fp32_partitions() + see_memory_usage("After creating fp32 partitions", force=True) + dist.barrier() + + # To support pipelined optimizer swapping + self._create_next_swappable_fp32_groups() + + see_memory_usage("Before initializing optimizer states", force=True) + + self.initialize_optimizer_states() + see_memory_usage("After initializing optimizer states", force=True) + dist.barrier() + + if dist.get_rank() == 0: + logger.info(f"optimizer state initialized") + + # IPG + if self.contiguous_gradients: + self.__ipg_bucket_flat_buffer: Tensor = torch.empty(self.reduce_bucket_size, + dtype=self.dtype, + device=get_accelerator().current_device_name()) + + self.grad_partitions_flat_buffer = None + self.__param_id_to_grad_partition: Dict[int, Tensor] = {} + + all_params = list(itertools.chain.from_iterable(self.fp16_groups)) + + self.grad_partitions_flat_buffer: Tensor = torch.zeros(sum(p.partition_numel() for p in all_params), + dtype=self.gradient_accumulation_dtype, + device=self.device) + if self.offload_optimizer_pin_memory: + self.grad_partitions_flat_buffer = get_accelerator().pin_memory(self.grad_partitions_flat_buffer) + + offset = 0 + for param in all_params: + self.__param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( + 0, offset, param.partition_numel()) + offset += param.partition_numel() + + def _link_all_hp_params(self): + for p in self.module.parameters(): + p._z3_optimizer = self + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + # TODO. factor out to a utility outside of stage3 + @staticmethod + def defragment(tensors: List[Tensor]) -> Tensor: + """move provided tensors into a contiguous flat buffer, with some additional + measures taken to reduce memory fragmentation""" + assert len(set(t.dtype for t in tensors)) == 1 + assert len(set(t.device for t in tensors)) == 1 + + cpu_buffer = torch.empty(sum(p.numel() for p in tensors), + dtype=get_only_unique_item(t.dtype for t in tensors), + device="cpu") + tensor_infos: List[Tuple[Tensor, int, int]] = get_mapping_to_flat_buffer(tensors) + orig_device = get_only_unique_item(t.device for t in tensors) + + offset = 0 + for tensor, offset, tensor_numel in tensor_infos: + # move the tensor from device memory to host memory + cpu_buffer.narrow(0, offset, tensor_numel).copy_(tensor) + tensor.data = torch.empty(0, dtype=tensor.dtype, device=tensor.device) + + gc.collect() + get_accelerator().empty_cache() + + # copy tensors (now flattened and contiguous) back to GPU + device_buffer = cpu_buffer.to(orig_device) + + # restore device tensors + for tensor, offset, tensor_numel in tensor_infos: + tensor.data = device_buffer.narrow(0, offset, tensor_numel) + + return device_buffer + + def _get_param_coordinator(self): + return self.parameter_offload.get_param_coordinator() + + def _configure_offloading(self, offload_optimizer_config, offload_param_config): + ###################### offload optimizer setup ################################## + if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none: + self.offload_optimizer = True + self.offload_optimizer_pin_memory = offload_optimizer_config.pin_memory + self.swap_optimizer = offload_optimizer_config.device == OffloadDeviceEnum.nvme + self.offload_optimizer_fast_init = offload_optimizer_config.fast_init + + ###################### offload param setup ################################## + if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none: + self.offload_param = True + self.offload_param_pin_memory = offload_param_config.pin_memory + self.params_in_nvme_and_cpu = offload_param_config.device == OffloadDeviceEnum.nvme + self.max_params_in_cpu = offload_param_config.max_in_cpu + print_rank_0( + f"FP16 params swapping is {self.params_in_nvme_and_cpu}, Max params in CPU is {self.max_params_in_cpu}", + force=False) + + def _configure_tensor_swapping(self, offload_optimizer_config, aio_config): + nvme_swap_folder = os.path.join(offload_optimizer_config.nvme_path, 'zero_stage_3') + os.makedirs(nvme_swap_folder, exist_ok=True) + if dist.get_rank() == 0: + logger.info(f'Tensor Swapping: Adding optimizer tensors') + + swapper_type = PipelinedOptimizerSwapper if offload_optimizer_config.pipeline else PartitionedOptimizerSwapper + + self.optimizer_swapper = swapper_type(swap_config=offload_optimizer_config, + aio_config=aio_config, + base_folder=nvme_swap_folder, + optimizer=self.optimizer, + largest_numel=max(self.fp16_partitioned_groups_flat_numel), + device=self.device, + dtype=torch.float32, + timers=self.timers) + + @property + def elements_in_ipg_bucket(self): + return sum(p.ds_numel for p in self.params_in_ipg_bucket) + + def _move_to_flat_buffer(self, param_list, flat_buffer, avoid_copy=False): + '''If flat buffer is None then the parameters in the param_list are + not copied to the flat buffer. This is because they exceed the number of max_params_in_cpu + Some of these parameters may already be in CPU in unflattened buffers + or they maybe in GPU, or they maybe in NVME. If they are in NVME, then + they will be marked as NOT_AVAILABLE, and will be moved to CPU when they are + needed during training.''' + if flat_buffer is None: + # this dst buffer is on NVMe, so skip this + return + + start = 0 + for param in param_list: + src = param.ds_tensor + dest = flat_buffer.narrow(0, start, src.ds_numel) + start = start + src.ds_numel + '''if the parameter was initialized in nvme then bring it to the destination buffer directly''' + if src.status == PartitionedParamStatus.NOT_AVAILABLE: + print_rank_0( + f"Swapping in {param.ds_id} with partition size {param.partition_numel()} permanently to CPU") + param.nvme_swapper.swap_into_buffer(param, dest) + src.data = dest.data + src.status = PartitionedParamStatus.AVAILABLE + else: + assert src.status == PartitionedParamStatus.AVAILABLE, "Partitioned Param must be available here" + if not avoid_copy: + dest.data.copy_(src.data) + src.data = dest.data + + # Final location must be gpu/cpu in this case + param.ds_tensor.final_location = 'not-nvme' + + def _create_param_groups_fp16_flat_cpu_memory(self): + + aggregate_params_count = 0 + + for j, param_group in enumerate(self.trainable_param_groups): + params_in_group = sum([p.partition_numel() for p in param_group['params']]) + + flat_buffer_size = params_in_group + + if self.params_in_nvme_and_cpu and \ + aggregate_params_count + params_in_group > self.max_params_in_cpu: + + flat_buffer_size = max(0, self.max_params_in_cpu - aggregate_params_count) + + aggregate_params_count += params_in_group + + if flat_buffer_size > 0: + print_rank_0(f"group {j} flat buffer size {flat_buffer_size}", force=False) + self.param_groups_fp16_flat_cpu_memory.append(get_accelerator().pin_memory( + torch.empty(int(flat_buffer_size), dtype=self.dtype))) + else: + print_rank_0(f"No flat buffer size. Param group size was {params_in_group}", force=False) + + self.param_groups_fp16_flat_cpu_memory.append(torch.empty(1, dtype=self.dtype)) + + def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): + dist.barrier() + + param_groups: List[List[Parameter]] = tuple( + self._create_fp16_sub_groups(param_group["params"]) for param_group in fp16_param_groups) + + # bookkeeping related to param groups + for param_group_idx, param_group in enumerate(param_groups): + for sub_group in param_group: + sub_group_idx = len(self.fp16_groups) + + # record sub group and partitions + self.fp16_groups.append(sub_group) + self.fp16_partitioned_groups.append([param.ds_tensor for param in sub_group]) + + # record sub group -> group mapping + self.sub_group_to_group_id[sub_group_idx] = param_group_idx + + # record total elements of parameter partitions in sub group + self.fp16_partitioned_groups_flat_numel.append(sum(p.partition_numel() for p in sub_group)) + + # record ds_ids of parameter partitions in sub group + self.fp16_partitioned_groups_flat_id.append([p.ds_id for p in sub_group]) + + # record padding required to align group to world size (only applies to last rank) + rank_requires_padding = dist.get_rank( + self.dp_process_group) == dist.get_world_size(self.dp_process_group) - 1 + self.groups_padding.append([p.padding_size() if rank_requires_padding else 0 for p in sub_group]) + + # move parameters to flattened buffer + if not self.offload_param: # partitioned params remain in GPU during training + # move parameter partitions into a single contiguous flat buffer + parameter_partitions = self._get_parameter_partitions() + + # We need to keep the reference to this buffer to make sure you can free it in `offload_states` + self.lp_param_buffer = __class__.defragment(parameter_partitions) + self._set_fp16_partitioned_groups_flat() + + else: # partitioned params offloaded to CPU when not in use + # create a flat CPU memory allocation for each param group + self._create_param_groups_fp16_flat_cpu_memory() + for param_group_idx, param_group in enumerate(param_groups): + flat_offset = 0 + for i, sub_group in enumerate(param_group): + total_elements = sum(p.partition_numel() for p in sub_group) + print_rank_0(f"Params in nvme and cpu {self.params_in_nvme_and_cpu}") + #Flat buffer may not be available for parameters that reside in NVME + if not self.params_in_nvme_and_cpu or flat_offset + total_elements <= self.param_groups_fp16_flat_cpu_memory[ + param_group_idx].numel(): + fp16_partitioned_group_flat = self.param_groups_fp16_flat_cpu_memory[param_group_idx].narrow( + 0, flat_offset, total_elements) + print_rank_0( + f"Creating a flat buffer for subgroup {i} requiring {total_elements} elements, and cumulative CPU elements {flat_offset + total_elements}", + force=False) + + elif self.params_in_nvme_and_cpu: + fp16_partitioned_group_flat = None + print_rank_0(f"No flat buffer for sub group {i} of {total_elements} elements", force=False) + else: + assert False, "Either params are in nvme, or they are in CPU memory. This code path should not be triggered. Please see you max_params_in_cpu and params_in_nvme configs" + + self.fp16_partitioned_groups_flat.append(fp16_partitioned_group_flat) + flat_offset += total_elements + + self._move_to_flat_buffer(sub_group, + fp16_partitioned_group_flat, + avoid_copy=not self.offload_param) + + # if necessary, create a pinned memory buffer to be used for swapping out + # params to NVME after optimizer step + should_create_fp16_flat_reuse_buffer = any(flattened_partition_group is None + for flattened_partition_group in self.fp16_partitioned_groups_flat) + if should_create_fp16_flat_reuse_buffer: + max_partition_numel, largest_partition_numel = 0, None + for sub_group in self.fp16_groups: + total_elements = sum(t.partition_numel() for t in sub_group) + if total_elements > max_partition_numel: + largest_partition_numel = [t.ds_numel for t in sub_group] + max_partition_numel = total_elements + + assert len(largest_partition_numel) > 0, f'Unexpected that largest partition is empty' + self.fp16_groups[0][0].nvme_swapper.reserve_partitioned_swap_space(largest_partition_numel) + + def _get_parameter_partitions(self) -> List[Tensor]: + return [param.ds_tensor for sub_group in self.fp16_groups for param in sub_group] + + def _swap_in_sub_group_to_flat_buffer(self, flat_buffer, sub_group_id): + offset = 0 + elements_in_sub_group = sum([t.ds_numel for t in self.fp16_partitioned_groups[sub_group_id]]) + assert (flat_buffer.numel() == elements_in_sub_group) + for param, partitioned_param in zip(self.fp16_groups[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]): + dest = flat_buffer.narrow(0, offset, partitioned_param.ds_numel) + if partitioned_param.status == PartitionedParamStatus.NOT_AVAILABLE: + print_rank_0( + f"Swapping in {param.ds_id} with elements {param.ds_numel} and partition {param.partition_numel()}" + ) + param.nvme_swapper.swap_in([param], async_op=False) + dest.data.copy_(partitioned_param.data) + param.nvme_swapper.remove_partition_and_release_buffers([param]) + print_rank_0(f"Swapping in {param.ds_id} done") + else: + dest.data.copy_(partitioned_param.data) + offset += partitioned_param.ds_numel + + def _create_next_swappable_fp32_groups(self): + reverse_order_indices = [i for i in range(len(self.fp32_partitioned_groups_flat))] + reverse_order_indices.reverse() + + next_group = None + for i in reverse_order_indices: + self.next_swappable_fp32_partitioned_groups.append(next_group) + if self._swappable_optimizer_subgroup(i): + next_group = self.fp32_partitioned_groups_flat[i] + + self.next_swappable_fp32_partitioned_groups.reverse() + + def _get_sub_group_partitions(self, sub_group_id): + sub_group_partitions = [] + for param, partitioned_param in zip(self.fp16_groups[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]): + if partitioned_param.status == PartitionedParamStatus.NOT_AVAILABLE: + swap_path = param.nvme_swapper.get_path(param, True) + sub_group_partitions.append((partitioned_param, param.partition_numel(), swap_path)) + else: + sub_group_partitions.append((partitioned_param, partitioned_param.ds_numel, None)) + + return sub_group_partitions + + def _create_fp32_partitions(self): + cpu_memory_usage = 0 + cpu_memory_sub_groups = 0 + nvme_memory_usage = 0 + num_swappable_partitions = 0 + num_swap_from_nvme_partitions = 0 + num_swap_from_cpu_partitions = 0 + swap_from_nvme_memory_usage = 0 + swap_from_cpu_memory_usage = 0 + GIGA_BYTES = (1024**3) + + swappable_fp32_tensors = [] + swappable_fp16_src_tensors = [] + nvme_fp16_partitions_info = [] + nvme_fp16_num_elems = [] + nvme_fp32_dest_tensors = [] + fp32_element_size = torch.tensor([], dtype=torch.float32).element_size() + + # Assign portion of subgroup to cpu, the other to gpu. + if self.offload_optimizer: + self.subgroup_to_device = {} + sub_group_size = len(self.fp16_partitioned_groups_flat) + # print(f"Partial offload sub_group_size is {sub_group_size}, ratio is {self.partial_offload}\n") + for i in range(sub_group_size): + if i < int(self.partial_offload * sub_group_size): + self.subgroup_to_device[i] = 'cpu' + else: + self.subgroup_to_device[i] = get_accelerator()._name + + for i, tensor in enumerate(self.fp16_partitioned_groups_flat): + num_elements = self.fp16_partitioned_groups_flat_numel[i] + ds_id_begin = str(self.fp16_partitioned_groups_flat_id[i][0]) + ds_id_end = str(self.fp16_partitioned_groups_flat_id[i][-1]) + ds_id = ds_id_begin + '_' + ds_id_end + + # a partition of the fp32 master weights that will be updated by this process + if self._swappable_optimizer_subgroup(i): + self.fp32_partitioned_groups_flat.append(torch.Tensor()) + self.fp32_partitioned_groups_flat[i].ds_id = ds_id + nvme_memory_usage += (fp32_element_size * num_elements) + num_swappable_partitions += 1 + + if self.params_in_nvme_and_cpu and tensor is None: + num_swap_from_nvme_partitions += 1 + swap_from_nvme_memory_usage += (fp32_element_size * num_elements) + if self.offload_optimizer_fast_init: + sub_group_partitions = self._get_sub_group_partitions(i) + nvme_fp16_partitions_info.append(sub_group_partitions) + nvme_fp16_num_elems.append(num_elements) + nvme_fp32_dest_tensors.append(self.fp32_partitioned_groups_flat[i]) + else: + unpinned_fp32_buffer = torch.empty(num_elements, device=self.device, dtype=torch.float) + self._swap_in_sub_group_to_flat_buffer(unpinned_fp32_buffer, i) + self.optimizer_swapper.initialize_parameters(parameters=[self.fp32_partitioned_groups_flat[i]], + src_tensors=[unpinned_fp32_buffer]) + else: + num_swap_from_cpu_partitions += 1 + swap_from_cpu_memory_usage += (fp32_element_size * num_elements) + swappable_fp32_tensors.append(self.fp32_partitioned_groups_flat[i]) + swappable_fp16_src_tensors.append(self.fp16_partitioned_groups_flat[i]) + else: + cpu_memory_usage += (fp32_element_size * num_elements) + cpu_memory_sub_groups += 1 + + if self.params_in_nvme_and_cpu and tensor is None: + unpinned_fp32_buffer = torch.empty(num_elements, device=self.device, dtype=torch.float) + self._swap_in_sub_group_to_flat_buffer(unpinned_fp32_buffer, i) + self.fp32_partitioned_groups_flat.append(unpinned_fp32_buffer) + else: + if self.offload_optimizer: + self.fp32_partitioned_groups_flat.append(self.fp16_partitioned_groups_flat[i].to( + self.subgroup_to_device[i]).clone().float().detach()) + else: + self.fp32_partitioned_groups_flat.append(self.fp16_partitioned_groups_flat[i].to( + self.device).clone().float().detach()) + self.fp32_partitioned_groups_flat[i].ds_id = ds_id + + self.fp32_partitioned_groups_flat[i].requires_grad = True # keep this in case internal optimizer uses it + + if len(swappable_fp32_tensors) > 0: + self.optimizer_swapper.initialize_parameters(parameters=swappable_fp32_tensors, + src_tensors=swappable_fp16_src_tensors) + + if len(nvme_fp32_dest_tensors) > 0: + fp16_pinned_buffers = self.fp16_groups[0][0].nvme_swapper.reserve_available_buffers() + assert len(fp16_pinned_buffers) > 0 + self.optimizer_swapper.initialize_from_swapped_fp16_params(fp16_partitions_info=nvme_fp16_partitions_info, + fp16_num_elems=nvme_fp16_num_elems, + fp16_pinned_buffers=fp16_pinned_buffers, + fp32_parameters=nvme_fp32_dest_tensors) + self.fp16_groups[0][0].nvme_swapper.release_reserved_buffers() + + nvme_gigabytes = nvme_memory_usage / GIGA_BYTES + print_rank_0(f'Swappable FP32 Partitions: count={num_swappable_partitions} size={nvme_gigabytes:5.2f} GB', + force=False) + if self.params_in_nvme_and_cpu: + print_rank_0( + f'Swap from NVMe Partitions: count = {num_swap_from_nvme_partitions}, size = {swap_from_nvme_memory_usage/GIGA_BYTES:5.2f}GB', + force=False) + print_rank_0( + f'Swap from CPU Partitions: count = {num_swap_from_cpu_partitions}, size = {swap_from_cpu_memory_usage/GIGA_BYTES:5.2f}GB', + force=False) + + cpu_memory_gigabytes = cpu_memory_usage / GIGA_BYTES + print_rank_0(f'In-Memory FP32 Partitions: count={cpu_memory_sub_groups} size={cpu_memory_gigabytes:5.2f} GB', + force=False) + + # Clear for on-the-fly population before the optimizer step + for param_group in self.optimizer.param_groups: + param_group['params'] = [] + + def _create_fp16_sub_groups(self, params_group): + + params_group_numel = sum([param.partition_numel() for param in params_group]) + sub_group_size = self.sub_group_size + + if sub_group_size is None or sub_group_size >= params_group_numel: + return [params_group] + + sub_groups = [] + sub_group = [] + local_sub_group_size = 0 + for param in params_group: + + sub_group.append(param) + local_sub_group_size += param.partition_numel() + + if local_sub_group_size >= sub_group_size or id(param) == id(params_group[-1]): + + sub_groups.append(sub_group) + + sub_group = [] + local_sub_group_size = 0 + + return sub_groups + + def _release_ipg_buffers(self): + if self.contiguous_gradients: + self.__ipg_bucket_flat_buffer = None + + def _optimizer_step(self, sub_group_id): + param_group_id = self.sub_group_to_group_id[sub_group_id] + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + if self.offload_optimizer: + cur_device = self.subgroup_to_device[sub_group_id] + if cur_device == 'cpu': + self.optimizer.param_groups[param_group_id]['params'] = [fp32_param] + cpu_loss = self.optimizer.step() + self.optimizer.param_groups[param_group_id]['params'] = [] + else: + self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] + gpu_loss = self.backup_optimizer.step() + self.backup_optimizer.param_groups[param_group_id]['params'] = [] + else: + self.optimizer.param_groups[param_group_id]['params'] = [fp32_param] + self.optimizer.step() + self.optimizer.param_groups[param_group_id]['params'] = [] + + def _swappable_optimizer_subgroup(self, sub_group_id): + if not self.swap_optimizer: + return False + + return self.optimizer_swapper.is_swappable_tensor(None, + numel=self.fp16_partitioned_groups_flat_numel[sub_group_id]) + + def _partitioned_params_swap_out(self, i): + offset = 0 + fp32_param = self.fp32_partitioned_groups_flat[i] + assert fp32_param is not None, \ + f'fp32 parameters of sub_group {i} is None' + + swap_fp16_params = [] + swap_fp32_params = [] + for param, partitioned_param in zip(self.fp16_groups[i], self.fp16_partitioned_groups[i]): + src = fp32_param.narrow(0, offset, partitioned_param.ds_numel) + if partitioned_param.status == PartitionedParamStatus.AVAILABLE: + partitioned_param.data.copy_(src.data) + else: + swap_fp32_params.append(src) + swap_fp16_params.append(param) + offset += partitioned_param.ds_numel + + if len(swap_fp16_params): + swap_fp16_params[0].nvme_swapper.swap_out_partitioned_params(dst_fp16_params=swap_fp16_params, + src_fp32_params=swap_fp32_params) + + def _set_fp16_partitioned_groups_flat(self): + # setup flat buffers per subgroup, these are each just sections of the + # contiguous flat buffer for all parameters that we created earlier + offset = 0 + for sub_group in self.fp16_groups: + sub_group_numel = sum(param.partition_numel() for param in sub_group) + self.fp16_partitioned_groups_flat.append(self.lp_param_buffer.narrow(0, offset, sub_group_numel)) + offset += sub_group_numel + + def initialize_optimizer_states(self): + num_subgroups = len(self.fp16_groups) + + largest_numel = max([sum([p.ds_numel for p in psg]) for psg in self.fp16_partitioned_groups]) + gradient_dtype = self.fp32_partitioned_groups_flat[0].dtype + gradient_buffer = torch.zeros(int(largest_numel), dtype=gradient_dtype, device=self.device) + + timer_names = set() + + # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers + # which do lazy initialization of the state at the first call to step. + is_adagrad = isinstance(self.optimizer, torch.optim.Adagrad) + + if self.swap_optimizer: + self.optimizer_swapper.init_timers() + + timer_names.add(INIT_OPTIMIZER_TIMER) + self.timers(INIT_OPTIMIZER_TIMER).start() + + for i, group in enumerate(self.fp16_groups): + swappable_optimizer_subgroup = self._swappable_optimizer_subgroup(i) + swappable_param_subgroup = self.fp16_partitioned_groups_flat[i] is None + + num_elements = int(self.fp16_partitioned_groups_flat_numel[i]) + + see_memory_usage( + f'[Begin] Initialize optimizer states {i} / {num_subgroups} subgroups, num_elems: {num_elements}, swappable opt/param:{swappable_optimizer_subgroup}/{swappable_param_subgroup}', + force=False) + + if swappable_optimizer_subgroup: + self._optimizer_states_and_gradient_swap_in(i, timer_names) + + if self.offload_optimizer and not swappable_optimizer_subgroup: + subgroup_gradient_buffer = torch.zeros(num_elements, dtype=gradient_dtype, device=self.device) + if self.offload_optimizer_pin_memory: + subgroup_gradient_buffer = get_accelerator().pin_memory(subgroup_gradient_buffer) + + self.fp32_partitioned_groups_flat[i].grad = subgroup_gradient_buffer.to(self.subgroup_to_device[i]) + else: + self.fp32_partitioned_groups_flat[i].grad = gradient_buffer.narrow(0, 0, num_elements) + + if swappable_param_subgroup: + self._partitioned_params_swap_out(i) + + if swappable_optimizer_subgroup: + self._optimizer_states_and_gradient_swap_out(i, timer_names) + + see_memory_usage( + f'[End] Initialize optimizer states {i} / {num_subgroups} subgroups, num_elems: {num_elements}, swappable opt/param:{swappable_optimizer_subgroup}/{swappable_param_subgroup}', + force=False) + + # Initialize the optimizer states with the flattened fp32 partition. + if is_adagrad: + self.optimizer = torch.optim.Adagrad(self.fp32_partitioned_groups_flat, **self.optimizer.defaults) + + self.timers(INIT_OPTIMIZER_TIMER).stop() + self.timers.log(timer_names) + + if self.swap_optimizer: + self.optimizer_swapper.log_timers() + + if not self.offload_optimizer: + for group in self.fp32_partitioned_groups_flat: + group.grad = None + + # Reset steps + return + + ######################################################################### + #########################ZeRO Partition Gradients######################## + ######################################################################### + + def get_first_param_index(self, group_id, param_group, partition_id): + for index, param in enumerate(param_group): + param_id = self.get_param_id(param) + if partition_id in self.param_to_partition_ids[group_id][param_id]: + return index + return None + + def initialize_gradient_partitioning_data_structures(self): + + total_partitions = dist.get_world_size(group=self.dp_process_group) + + for i, param_group in enumerate(self.fp16_groups): + + self.param_to_partition_ids[i] = {} + self.is_partition_reduced[i] = {} + self.total_grads_in_partition[i] = {} + self.remaining_grads_in_partition[i] = {} + self.is_grad_computed[i] = {} + self.grad_partition_insertion_offset[i] = {} + self.grad_start_offset[i] = {} + self.first_param_index_in_partition[i] = {} + + for partition_id in range(total_partitions): + self.is_grad_computed[i][partition_id] = {} + self.grad_partition_insertion_offset[i][partition_id] = {} + self.grad_start_offset[i][partition_id] = {} + self.initialize_gradient_partition(i, param_group, partition_id) + self.is_partition_reduced[i][partition_id] = False + self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index( + i, param_group, partition_id) + + @instrument_w_nvtx + def independent_gradient_partition_epilogue(self): + self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) + self.__reduce_and_partition_ipg_grads() + self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) + + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + for param_id in self.params_already_reduced.keys(): + self.params_already_reduced[param_id] = False + + #in case of cpu offload, averaged gradients are already in fp32_partitioned_groups_flat.grad + #TODO: use a similar code path for both cpu_offload and non-cpu offload + if not self.offload_optimizer: + for i, sub_group in enumerate(self.fp16_groups): + #TODO: This is redundant + self.averaged_gradients[i] = [ + self.__param_id_to_grad_partition[param.ds_id] + if param.requires_grad else torch.zeros_like(param.ds_tensor) for param in sub_group + ] + # this method gets called after every backward. need to increment + # here because if it gets incremented in backward() the micro step + # id will be off by one when we do the reduce and partition at the. + # start of this method. + # TODO. make this less error prone + self.micro_step_id += 1 + + def overlapping_partition_gradients_reduce_epilogue(self): + self.independent_gradient_partition_epilogue() + + def create_reduce_and_remove_grad_hooks(self): + print_rank_0(f'[Begin] Create gradient reduction hooks') + self.leaf_parameters = defaultdict(list) + for i, param_group in enumerate(self.fp16_groups): + for param in param_group: + if param.requires_grad: + #print_rank_0(f" Before all gather {param.device}, {param.shape}") + print_rank_0(f"Before all gather {param.device}, {param.shape}", force=False) + + # The hook must be created in un-partitioned parameter + param.all_gather() + + #print(f"After all gather {param.device}, {param.shape}") + def wrapper(param): + + @instrument_w_nvtx + def reduce_partition_and_remove_grads(*notneeded): + self.reduce_ready_partitions_and_remove_grads(param) + + self._grad_acc_hooks.append(register_grad_hook(param, reduce_partition_and_remove_grads)) + + #print(f"param grad fn {param.expand_as(param).grad_fn}") + if z3_leaf_parameter(param): + self.leaf_parameters[param.ds_z3_leaf_module].append(param) + else: + wrapper(param) + + # Partition the parameter after creating the hook + param.partition() + + # We delay reduce-scatter for all gradients in the leaf modules until the backward pass of the leaf module is done + for leaf_module, leaf_parameters in self.leaf_parameters.items(): + + def wrapper_pre_hook(params): + + def forward_pre_hook(module, input): + """Pre-forward hook to set backward hook on input tensors to the leaf module""" + module._leaf_module_inputs_remaining = 0 + + @instrument_w_nvtx + def reduce_leaf_module_grads(grad): + module._leaf_module_inputs_remaining -= 1 + # Make sure everything is done in the leaf module + if module._leaf_module_inputs_remaining == 0: + for param in params: + if param.grad is None: + param.grad = torch.zeros_like(param) + self.reduce_ready_partitions_and_remove_grads(param) + + def set_module_bwd_hook(tensor): + if tensor.requires_grad: + module._leaf_module_inputs_remaining += 1 + tensor.register_hook(reduce_leaf_module_grads) + return tensor + + output = apply_to_tensors_only(set_module_bwd_hook, input) + + return output + + return forward_pre_hook + + def wrapper_post_hook(): + + def forward_post_hook(module, input, output): + """Pre-forward hook to set backward hook on input tensors to the leaf module""" + module._leaf_output_required_grad_num = 0 + + def increment_rg_count_bwd_hook(tensor): + if tensor.requires_grad: + module._leaf_output_required_grad_num += 1 + return tensor + + apply_to_tensors_only(increment_rg_count_bwd_hook, output) + + if module._leaf_module_inputs_remaining == 0 and module._leaf_output_required_grad_num > 0: + raise RuntimeError( + "A module cannot be set as a leaf module when it does not have any input tensors that require gradients and has output tensors that require gradients. This is because the gradient reduction hook will not be called in this case." + ) + + return forward_post_hook + + self._leaf_module_hooks.append(leaf_module.register_forward_pre_hook(wrapper_pre_hook(leaf_parameters))) + self._leaf_module_hooks.append(leaf_module.register_forward_hook(wrapper_post_hook())) + + print_rank_0(f'[End] Create gradient reduction hooks') + + def get_param_id(self, param): + return OptimizerSwapper.parameter_id(param) + + def report_ipg_memory_usage(self, tag, param_elems): + elem_count = self.elements_in_ipg_bucket + param_elems + percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size + see_memory_usage( + f"{tag}: elems in_bucket {self.elements_in_ipg_bucket} param {param_elems} max_percent {percent_of_bucket_size}", + force=False) + + ###############Independent Partition Gradient ######################## + def reduce_independent_p_g_buckets_and_remove_grads(self, param): + #print_rank_0(f"Inside reduce ipg buckets. {debug_param2name_id_shape(param)}, ipg elements {self.elements_in_ipg_bucket}, reduce bucket size {self.reduce_bucket_size}", force=True) + + # Because the ipg bucket is initialized with a random place holder tensor, we must + # explicitly check that the bucket has any real data in it (self.elements_in_ipg_bucket > + # 0). Otherwise if the incoming param.ds_numel is large, this branch may get triggered on a + # garbage data and `self.average_tensor()` will crash because its params_to_reduce will be + # empty, while reduction_list will have that garbage data. + if self.elements_in_ipg_bucket + param.ds_numel > self.reduce_bucket_size and self.elements_in_ipg_bucket > 0: + self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.ds_numel) + + self.__reduce_and_partition_ipg_grads() + + # deal with a use-case of transient grads that will be generated in a loop for the same computation involving some model params - e.g. when performing a tiled memory calculation that shards the normal single sub-module call into a loop over a shards. + if getattr(param, "ds_grad_is_ready", True): + self.__add_grad_to_ipg_bucket(param) + + @instrument_w_nvtx + @torch.no_grad() + def __add_grad_to_ipg_bucket(self, param: Parameter) -> None: + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.wait_stream(get_accelerator().default_stream()) + + if self.contiguous_gradients and self.elements_in_ipg_bucket + param.grad.numel() <= self.reduce_bucket_size: + # move the gradient to a contiguous buffer + with get_accelerator().stream(self.reduce_and_partition_stream): + # move the parameter's gradient to the contiguous flat buffer + new_grad_tensor = self.__ipg_bucket_flat_buffer.narrow(0, self.elements_in_ipg_bucket, + param.grad.numel()).view_as(param.grad) + new_grad_tensor.copy_(param.grad, non_blocking=True) + if not get_accelerator().is_synchronized_device(): + param.grad.record_stream(get_accelerator().current_stream()) + param.grad.data = new_grad_tensor + + self.params_in_ipg_bucket.append(param) + + @instrument_w_nvtx + @torch.no_grad() + def __reduce_and_partition_ipg_grads(self, safe_mode: bool = False) -> None: + if not self.params_in_ipg_bucket: + return + + for param in self.params_in_ipg_bucket: + if param.grad.numel() != param.ds_numel: + raise RuntimeError(f"{param.grad.numel()} != {param.ds_numel} Cannot reduce scatter " + f"gradients whose size is not same as the params") + + assert len(set(p.ds_id for p in self.params_in_ipg_bucket)) == len(self.params_in_ipg_bucket) + + while self.param_reduce_events and self.param_reduce_events[0].query(): + self.param_reduce_events.popleft() + if len(self.param_reduce_events) > self.max_param_reduce_events: + self.param_reduce_events.popleft().synchronize() + + with get_accelerator().stream(self.reduce_and_partition_stream): + if safe_mode: + assert_ints_same_as_other_ranks([p.ds_id for p in self.params_in_ipg_bucket]) + + if self.contiguous_gradients and self.elements_in_ipg_bucket <= self.reduce_bucket_size and not self.reduce_scatter: + grad_bucket = self.__ipg_bucket_flat_buffer.narrow(0, 0, self.elements_in_ipg_bucket) + grad_partitions = self.__avg_scatter_contiguous_grads(grad_bucket) + else: + self.params_in_ipg_bucket.sort(key=lambda p: p.ds_id) + grad_partitions = self.__avg_scatter_grads(self.params_in_ipg_bucket) + + self.partition_grads(self.params_in_ipg_bucket, grad_partitions) + + self.params_in_ipg_bucket.clear() + + if not get_accelerator().handles_memory_backpressure(): + event = get_accelerator().Event() + event.record() + self.param_reduce_events.append(event) + + @instrument_w_nvtx + def __avg_scatter_contiguous_grads(self, buffer_to_reduce: Tensor) -> List[Tensor]: + dtype = buffer_to_reduce.dtype + if self.communication_data_type != dtype: + buffer_to_reduce = buffer_to_reduce.to(self.communication_data_type) + if self.postscale_gradients and self.gradient_predivide_factor != 1.0: + buffer_to_reduce = buffer_to_reduce.div_(self.gradient_predivide_factor) + + world_sz = dist.get_world_size(self.dp_process_group) + rank = dist.get_rank(self.dp_process_group) + buffer_to_reduce.div_(world_sz / float(self.sequence_parallel_size)) + + dist.all_reduce(buffer_to_reduce, group=self.dp_process_group) + + if self.postscale_gradients and self.gradient_predivide_factor != world_sz: + buffer_to_reduce = buffer_to_reduce.mul(self.gradient_predivide_factor) + + if self.communication_data_type != self.dtype: + buffer_to_reduce = buffer_to_reduce.to(self.dtype) + + grad_partitions = [] + grad_offset_in_buffer = 0 + for param in self.params_in_ipg_bucket: + grad = param.grad + chunk_sz = math.ceil(grad.numel() / world_sz) + + start_offset = grad_offset_in_buffer + min(rank * chunk_sz, grad.numel()) + end_offset = grad_offset_in_buffer + min(rank * chunk_sz + chunk_sz, grad.numel()) + + partition = buffer_to_reduce[start_offset:end_offset] + if param.partition_numel() != partition.numel(): + padded_partition = torch.zeros(param.partition_numel(), device=grad.device, dtype=grad.dtype) + if partition.numel() > 0: + padded_partition[:partition.numel()] = partition + grad_partitions.append(padded_partition) + else: + grad_partitions.append(partition) + grad_offset_in_buffer += grad.numel() + + return grad_partitions + + @instrument_w_nvtx + def __avg_scatter_grads(self, params_to_reduce: List[Parameter]) -> List[Tensor]: + """average gradients and scatter partitions across ranks""" + + full_grads_for_rank = [p.grad for p in params_to_reduce] + if self.communication_data_type != self.dtype: + full_grads_for_rank = [g.to(self.communication_data_type) for g in full_grads_for_rank] + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0: + full_grads_for_rank = [g.div(self.gradient_predivide_factor) for g in full_grads_for_rank] + + local_world_size = get_accelerator().device_count() + global_world_size = dist.get_world_size() + num_nodes = global_world_size // local_world_size + if self.all2all_process_group is not None and num_nodes > 1: + grad_partitions_for_rank = (all_to_all_loco_quant_reduce(params_to_reduce, self.all2all_process_group, + self.zeropp_loco_param) + if self.zeropp_loco_param is not None else all_to_all_quant_reduce( + full_grads_for_rank, self.all2all_process_group)) + else: + grad_partitions_for_rank = reduce_scatter_coalesced(full_grads_for_rank, self.dp_process_group) + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0 and self.gradient_predivide_factor != dist.get_world_size( + self.dp_process_group): + grad_partitions_for_rank = [g.mul(self.gradient_predivide_factor) for g in grad_partitions_for_rank] + + if self.communication_data_type != self.dtype: + grad_partitions_for_rank = [g.to(self.dtype) for g in grad_partitions_for_rank] + + return grad_partitions_for_rank + + def set_grad_positions(self): + for i, group in enumerate(self.fp16_groups): + current_offset = 0 + for param in group: + param_id = self.get_param_id(param) + num_elements = param.partition_numel() + + self.grad_position[param_id] = [int(i), int(current_offset), int(num_elements)] + #print(f"param id {param_id} i:{i}, ds_tensor {num_elements} numel {param.numel()}") + current_offset += num_elements + see_memory_usage(f"After Set Grad positions", force=False) + + def _constant_buffered_norm2(self, input, buffer_size=250000000): + norm = None + for part in input.view(-1).split(buffer_size): + if norm is None: + norm = part.data.double().norm(2)**2.0 + else: + norm += part.data.double().norm(2)**2.0 + return norm**0.5 + + def set_norm_for_param_grad_in_gpu(self, param): + param_id = self.get_param_id(param) + #self.norm_for_param_grads[param_id] = param.grad.data.double().norm(2) + #Using a more memory efficient version + self.norm_for_param_grads[param_id] = self._constant_buffered_norm2(param.grad) + + def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param, fp32_grad_tensor): + with get_accelerator().stream(self.copy_grad_stream): + param_id = self.get_param_id(param) + src_tensor = param.grad.view(-1).float() + #print(f"src_tensor {src_tensor.size()} and fp32 grad {fp32_grad_tensor.size()}") + fp32_grad_tensor.copy_(src_tensor, non_blocking=True) + param.grad = None + + def complete_grad_norm_calculation_for_cpu_offload(self, params): + total_norm = 0.0 + norm_type = 2.0 + for p in params: + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + param_id = self.get_param_id(p) + if param_id in self.norm_for_param_grads.keys(): + param_norm = self.norm_for_param_grads[param_id] + total_norm += param_norm**2 + + # Sum across all model parallel GPUs. + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) + + total_norm = total_norm_cuda[0]**(1. / norm_type) + + mask_nan_or_inf_with_val_inplace(total_norm, device=total_norm.device) + + return total_norm.cpu() + + @instrument_w_nvtx + def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: + offload_fp32_gradients = {} + offload_fp32_offsets = {} + buffers = [] + for param, grad_partition in zip(params_to_release, grad_partitions): + + contains_real_data = param.partition_numel() * dist.get_rank(self.dp_process_group) < param.ds_numel + if not contains_real_data: + # this grad partition is empty - don't need to do anything + param.grad = None + continue + + # move or accumulate gradient partition to target buffer + grad_buffer = self.__param_id_to_grad_partition[param.ds_id].narrow(0, 0, grad_partition.numel()) + buffers.append(grad_buffer) + if self.micro_step_id == 0: # don't accumulate + grad_buffer.copy_(grad_partition, non_blocking=True) + # ensure grad buffer is a CUDA buffer to speed up the next few + # operations and so it can be used asynchronously + grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) + elif get_accelerator().on_accelerator(grad_buffer): + grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(grad_buffer.shape)) + else: + # if dst is CPU, copy first to src device, do the addition + # there, then move back to dst. adding directly to cpu is very slow + cuda_grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) + cuda_grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(cuda_grad_buffer.shape)) + grad_buffer.copy_(cuda_grad_buffer, non_blocking=True) + # ensure grad buffer is a CUDA buffer to speed up the next few + # operations and so it can be used asynchronously + grad_buffer = cuda_grad_buffer + + # offload the gradient partition if applicable + if self.offload_optimizer: + i, dest_offset, _ = self.grad_position[self.get_param_id(param)] + + if self.is_gradient_accumulation_boundary: + self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) + + if self._swappable_optimizer_subgroup(i): + if not i in offload_fp32_gradients.keys(): + offload_fp32_gradients[i] = [] + offload_fp32_offsets[i] = [] + + offload_fp32_gradients[i].append(grad_buffer.float()) + offload_fp32_offsets[i].append(dest_offset) + else: + fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow( + 0, dest_offset, grad_buffer.numel()) + fp32_grad_tensor.copy_(grad_buffer.float()) + + # free the gradient + if not get_accelerator().is_synchronized_device(): + if param.grad is not None: + param.grad.record_stream(get_accelerator().current_stream()) + param.grad = None + + if self.offload_optimizer and self.swap_optimizer: + for i in offload_fp32_gradients.keys(): + self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i], + gradient_offsets=offload_fp32_offsets[i], + gradient_tensors=offload_fp32_gradients[i]) + return buffers + + def reduce_ready_partitions_and_remove_grads(self, param): + #print_rank_0(f"Backward {debug_param2name_id_shape(param)}", force=True) + self.reduce_independent_p_g_buckets_and_remove_grads(param) + + def zero_reduced_gradients(self, partition_id, i): + + def are_all_related_partitions_reduced(params_id): + for partition_id in self.param_to_partition_ids[i][params_id]: + if not self.is_partition_reduced[i][partition_id]: + return False + return True + + for params_id in self.is_grad_computed[i][partition_id]: + if are_all_related_partitions_reduced(params_id): + self.param_dict[params_id].grad = None + + def quantize_nontrainable_params(self): + """ In ZeRO-3, when the zero_quantized_nontrainable_weights flag is set, we quantize the non-trainable weights and also store them in quantized format. However, this check for trainable/non-trainable is done when deepspeed initializes the partitioning. So, if the user changes the trainable/non-trainable status of a parameter after the partitioning is done (e.g. LoRA), the user needs to re-quantize the non-trainable weights by calling this function. + """ + if not self.zero_quantized_nontrainable_weights: + print_rank_0( + f"Warning: quantize_nontrainable_params() called with zero_quantized_nontrainable_weights disabled, return without doing anything", + force=True) + return + quantizer_module = CUDAQuantizer() + + def quantize_dstensor(tensor): + assert tensor.dtype == torch.float16, f"quantize_dstensor() expects tensor.dtype == torch.float16, got {tensor.dtype}" + partition_size = tensor.ds_numel + ds_status = tensor.status + final_location = tensor.final_location + tensor, tensor.ds_quant_scale = quantizer_module.quantize(tensor) + tensor.ds_numel = partition_size + tensor.status = ds_status + tensor.final_location = final_location + tensor.requires_grad = False + return tensor + + for param in self.module.parameters(): + if hasattr(param, "ds_tensor") and (param.ds_tensor.numel() <= 2048 or param.ds_numel <= 500000): + # skip small parameters + continue + if hasattr(param, + "ds_tensor") and not param.requires_grad and not hasattr(param.ds_tensor, "ds_quant_scale"): + param.ds_tensor = quantize_dstensor(param.ds_tensor) + if hasattr(param, "ds_secondary_tensor") and not param.requires_grad and not hasattr( + param.ds_secondary_tensor, "ds_quant_scale") and param.ds_secondary_tensor is not None: + param.ds_secondary_tensor = quantize_dstensor(param.ds_secondary_tensor) + get_accelerator().synchronize() + + def flatten_and_print(self, message, tensors, start=0, n=5): + flatten_tensor = self.flatten(tensors) + + def print_func(): + logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n)) + + self.sequential_execution(print_func, message) + + def get_grads_to_reduce(self, i, partition_id): + + def get_reducible_portion(key): + grad = self.param_dict[key].grad + total_elements = grad.numel() + start = self.grad_start_offset[i][partition_id][key] + num_elements = min(total_elements - start, + self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key]) + if not pg_correctness_test: + if num_elements == total_elements: + return grad + else: + return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements)) + else: + if num_elements == total_elements: + return grad.clone() + else: + return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements)) + + grads_to_reduce = [] + for key in self.is_grad_computed[i][partition_id]: + grad = get_reducible_portion(key) + grads_to_reduce.append(grad) + return grads_to_reduce + + def sequential_execution(self, function, message, group=None): + if group is None: + group = self.dp_process_group + if dist.get_rank(group=group) == 0: + logger.info(message) + for id in range(dist.get_world_size(group=group)): + if id == dist.get_rank(group=group): + function() + dist.barrier(group=group) + + def set_none_gradients_to_zero(self, i, partition_id): + for param_id in self.is_grad_computed[i][partition_id]: + param = self.param_dict[param_id] + if param.grad is None: + param.grad = torch.zeros_like(param) + + ######################Reduction Related Methods############################## + + def allreduce_bucket(self, bucket, rank=None, log=None): + rank = None + tensor = self.flatten(bucket) + + tensor_to_allreduce = tensor + + if pg_correctness_test: + communication_data_type = torch.float32 + else: + communication_data_type = self.communication_data_type + + if communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(communication_data_type) + + tensor_to_allreduce.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) + + if rank is None: + # "All Reducing" + dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) + else: + global_rank = dist.get_global_rank(self.dp_process_group, rank) + dist.reduce(tensor_to_allreduce, global_rank, group=self.dp_process_group) + + if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + if rank is None or rank == dist.get_rank(group=self.dp_process_group): + tensor.copy_(tensor_to_allreduce) + + return tensor + + # if rank is specified do a reduction instead of an allreduce + def allreduce_and_copy(self, small_bucket, rank=None, log=None): + with get_accelerator().stream(self.reduction_stream): + allreduced = self.allreduce_bucket(small_bucket, rank=rank, log=log) + if rank is None or rank == dist.get_rank(group=self.dp_process_group): + for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): + buf.copy_(synced) + + def allreduce_no_retain(self, bucket, numel_per_bucket=500000000, rank=None, log=None): + small_bucket = [] + numel = 0 + for tensor in bucket: + small_bucket.append(tensor) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy(small_bucket, rank=rank, log=None) + small_bucket = [] + if len(small_bucket) > 0: + self.allreduce_and_copy(small_bucket, rank=rank, log=log) + + ############################################################################# + ############################################################################# + ############################################################################# + + # views the tensor as multiple partitions and returns + # those partitions + def get_data_parallel_partitions(self, tensor): + partitions = [] + + dp = dist.get_world_size(group=self.dp_process_group) + dp_id = dist.get_rank(group=self.dp_process_group) + + total_num_elements = tensor.numel() + + base_size = total_num_elements // dp + remaining = total_num_elements % dp + + start = 0 + for id in range(dp): + partition_size = base_size + if id < remaining: + partition_size = partition_size + 1 + partitions.append(tensor.narrow(0, start, partition_size)) + start = start + partition_size + return partitions + + def get_partition_info(self, tensor_list, partition_size, partition_id): + params_in_partition = [] + params_not_in_partition = [] + + start_index = partition_size * partition_id + end_index = partition_size * (partition_id + 1) + + current_index = 0 + first_offset = 0 + + for tensor in tensor_list: + + tensor_size = tensor.numel() + + if start_index <= current_index < end_index: + params_in_partition.append(tensor) + + elif current_index < start_index < (current_index + tensor_size): + params_in_partition.append(tensor) + + assert (first_offset == 0 + ), "This can happen either zero or only once as this must be the first tensor in the partition" + first_offset = start_index - current_index + + else: + params_not_in_partition.append(tensor) + + current_index = current_index + tensor_size + + return params_in_partition, params_not_in_partition, first_offset + + @instrument_w_nvtx + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + self.micro_step_id = 0 + + # FP32 grad should never exist. + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if set_to_none: + if p.grad is not None and get_accelerator().on_accelerator(p.grad): + p.grad.record_stream(get_accelerator().current_stream()) + p.grad = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def _model_parallel_all_reduce(self, tensor, op): + """ Perform all reduce within model parallel group, if any. + """ + if self.model_parallel_group is None: + pass + else: + dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group) + + @instrument_w_nvtx + def get_grad_norm_direct(self, gradients, params, norm_type=2): + """Clips gradient norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float or int): max norm of the gradients + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + """ + norm_type = float(norm_type) + if norm_type == inf: + total_norm = max(g.data.abs().max() for g in gradients) + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + # Take max across all GPUs. + self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX) + total_norm = total_norm_cuda[0] + else: + # if dist.get_rank() == 0: + # logger.info(f"Total Norm beginning {total_norm}") + grad_norms = [] + for g, p in zip(gradients, params): + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + grad_norms.append(g.to(get_accelerator().device_name(), non_blocking=True).double().norm(2)) + + # Sum across all model parallel GPUs. + if len(grad_norms) == 0: + # FIX https://github.com/deepspeedai/DeepSpeed/issues/3564 + total_norm_cuda = torch.tensor(0, + dtype=gradients[0].dtype).to(get_accelerator().device_name()).double() + else: + total_norm_cuda = torch.sum(torch.pow(torch.stack(grad_norms), 2)) + + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) + + total_norm = total_norm_cuda**(1. / norm_type) + + norm_is_inf = total_norm.isinf() + norm_is_nan = total_norm.isnan() + inf_or_nan = norm_is_nan.logical_or(norm_is_inf) + + err = torch.tensor(-1.0, device=self.device, dtype=torch.float) + total_norm = torch.where(inf_or_nan, err, total_norm) + + return total_norm + + # creates a flat fused tensor from the tensor list starting at the first_offset + # in the first tensor of the list. If there are not enough elements in the tensor + # list then the flat tensor will be padded with zeros + def get_flat_partition(self, tensor_list, first_offset, partition_size, return_tensor_list=False): + flat_tensor_list = [] + current_size = 0 + for i, tensor in enumerate(tensor_list): + if tensor.grad is None: + tensor.grad = torch.zeros_like(tensor) + + tensor = tensor.grad + num_elements = tensor.numel() + tensor_offset = 0 + + # we need to offset to get to the right element + if i == 0 and first_offset > 0: + tensor_offset = first_offset + num_elements = num_elements - tensor_offset + + # we dont need all elements of the tensor + if num_elements > (partition_size - current_size): + num_elements = partition_size - current_size + + # we need a narrow view of the tensor based on the tensor offset and number of elements that + # we need from this tensor + if tensor_offset > 0 or num_elements < tensor.numel(): + flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements))) + else: + flat_tensor_list.append(tensor) + + current_size = current_size + num_elements + + # this means its the last partition and does not align with the dp boundary. We need to pad before flattening + if current_size < partition_size: + flat_tensor_list.append( + torch.zeros(int(partition_size - current_size), + dtype=tensor_list[0].dtype, + device=tensor_list[0].device)) + + if return_tensor_list: + return flat_tensor_list + + return self.flatten(flat_tensor_list) + + def free_grad_in_param_list(self, param_list): + for p in param_list: + p.grad = None + + def reset_cpu_buffers(self): + self.norm_for_param_grads = {} + + def _pre_step(self): + self.micro_step_id = 0 + + print_rank_0(f"Inside Step function") + see_memory_usage(f"In step before checking overflow", force=False) + + print_rank_0("Finished Tracing at Beginning of Step") + self._get_param_coordinator().hierarchy = 0 + + print_rank_0("Finished Tracing at Beginning of Step") + + @instrument_w_nvtx + def _get_norm_groups(self): + norm_groups = [] + for i, group in enumerate(self.fp16_groups): + if self.offload_optimizer: + norm_groups.append(self.complete_grad_norm_calculation_for_cpu_offload(self.fp16_groups[i])) + else: + norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.fp16_groups[i])) + return norm_groups + + @instrument_w_nvtx + def _prepare_fp32_grad_for_sub_group(self, sub_group_id): + partition_id = dist.get_rank(group=self.dp_process_group) + + single_grad_partition = self.flatten(self.averaged_gradients[sub_group_id]).to( + self.fp32_partitioned_groups_flat[sub_group_id].dtype) + + assert single_grad_partition.numel() == self.fp32_partitioned_groups_flat[sub_group_id].numel(), \ + "averaged gradients have different number of elements that partition size {} {} {} {}".format( + single_grad_partition.numel(), self.fp32_partitioned_groups_flat[sub_group_id].numel(), sub_group_id, partition_id) + + self.fp32_partitioned_groups_flat[sub_group_id].grad = single_grad_partition + + # release all the gradient since we have already created a necessary copy in dp_grad_partition + self.zero_grad(set_to_none=True) + + if not get_accelerator().is_synchronized_device(): + for grad in filter(lambda g: get_accelerator().on_accelerator(g), self.averaged_gradients[sub_group_id]): + grad.record_stream(get_accelerator().current_stream()) + + self.averaged_gradients[sub_group_id] = None + + @instrument_w_nvtx + def _prepare_sub_group(self, sub_group_id, timer_names): + see_memory_usage(f'Before prepare optimizer sub group {sub_group_id}', force=False) + if self._swappable_optimizer_subgroup(sub_group_id): + self._optimizer_states_and_gradient_swap_in(sub_group_id, timer_names) + elif not self.offload_optimizer: + self._prepare_fp32_grad_for_sub_group(sub_group_id) + see_memory_usage(f'After prepare optimizer sub group {sub_group_id}', force=False) + + def _optimizer_states_and_gradient_swap_in(self, sub_group_id, timer_names=None): + param_length = self.fp16_partitioned_groups_flat_numel[sub_group_id] + fp32_param_id = self.get_param_id(self.fp32_partitioned_groups_flat[sub_group_id]) + assert self._swappable_optimizer_subgroup(sub_group_id), \ + f'Parameter {fp32_param_id} of numel={param_length} is not swappable' + + see_memory_usage(f'pre-step Before swapping in optimizer tensors {sub_group_id}', force=False) + if timer_names is not None: + timer_names.add(OPTIMIZER_SWAP_IN_STATE_TIMER) + self.timers(OPTIMIZER_SWAP_IN_STATE_TIMER).start() + + self.optimizer_swapper.swap_in_optimizer_state( + parameter=self.fp32_partitioned_groups_flat[sub_group_id], + async_parameter=self.next_swappable_fp32_partitioned_groups[sub_group_id]) + + if timer_names is not None: + self.timers(OPTIMIZER_SWAP_IN_STATE_TIMER).stop() + see_memory_usage(f'pre-step After swapping in optimizer tensors {sub_group_id}', force=False) + + @instrument_w_nvtx + def _release_sub_group(self, sub_group_id, timer_names): + see_memory_usage(f'Before release optimizer sub group {sub_group_id}', force=False) + # get rid of the fp32 gradients. Not needed anymore + if not self.offload_optimizer: + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + if self._swappable_optimizer_subgroup(sub_group_id): + self._optimizer_states_and_gradient_swap_out(sub_group_id, timer_names) + see_memory_usage(f'After release optimizer sub group {sub_group_id}', force=False) + + # create a flat tensor aligned at the alignment boundary + @instrument_w_nvtx + def flatten_dense_tensors_aligned(self, tensor_list, alignment): + num_elements = 0 + for tens in tensor_list: + num_elements = num_elements + tens.numel() + + remaining = num_elements % alignment + + if remaining: + elements_to_add = alignment - remaining + pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype) + padded_tensor_list = tensor_list + [pad_tensor] + + num_elements = num_elements + elements_to_add + else: + padded_tensor_list = tensor_list + + return self.flatten(padded_tensor_list) + + def _optimizer_states_and_gradient_swap_out(self, sub_group_id, timer_names=None): + param_length = self.fp16_partitioned_groups_flat_numel[sub_group_id] + fp32_param_id = self.get_param_id(self.fp32_partitioned_groups_flat[sub_group_id]) + assert self._swappable_optimizer_subgroup(sub_group_id), \ + f'Parameter {fp32_param_id} of numel={param_length} is not swappable' + + see_memory_usage(f'post-step Before swapping out optimizer tensors {sub_group_id}', force=False) + if timer_names is not None: + timer_names.add(OPTIMIZER_SWAP_OUT_STATE_TIMER) + self.timers(OPTIMIZER_SWAP_OUT_STATE_TIMER).start() + + self.optimizer_swapper.swap_out_optimizer_state( + parameter=self.fp32_partitioned_groups_flat[sub_group_id], + async_swap=self.next_swappable_fp32_partitioned_groups[sub_group_id] is not None) + + if timer_names is not None: + self.timers(OPTIMIZER_SWAP_OUT_STATE_TIMER).stop() + see_memory_usage(f'post-step After swapping out optimizer tensors {sub_group_id}', force=False) + + # get rid of the fp32 gradients. Not needed anymore + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + def _release_swap_buffers(self, sub_group_id): + self.optimizer_swapper.release_swap_buffers(parameter=self.fp32_partitioned_groups_flat[sub_group_id]) + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + def _writeback_swap_state(self, sub_group_id, write_opt_state, write_gradients): + self.optimizer_swapper.writeback_optimizer_state_and_gradients(self.fp32_partitioned_groups_flat[sub_group_id], + write_opt_state, write_gradients) + self.fp32_partitioned_groups_flat[sub_group_id].grad = None + + def _unflatten_partitioned_parameters(self, sub_group_id): + updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]) + + for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): + partitioned_param.data = q.data + + def _overflow_clean_up(self, prev_scale): + see_memory_usage('After overflow before clearing gradients', force=False) + self.zero_grad(set_to_none=True) + + if self.offload_optimizer: + self.reset_cpu_buffers() + else: + self.averaged_gradients = {} + + see_memory_usage('After overflow after clearing gradients', force=False) + + def _loco_err_buf_update(self, overflow: bool, scale=1.0): + """ + Loco Error Buffer update. + """ + if not overflow and scale == 1.0: return + if dist.get_rank() == 0: + logger.info(f"update loco-zero++ error buffer with overflow: {overflow}") + # FP32 grad should never exist. + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if hasattr(p, 'intra_ef_buf'): + if overflow: + del p.intra_ef_buf + del p.inter_ef_buf + else: + p.intra_ef_buf[1] *= scale + p.inter_ef_buf[1] *= scale + + @instrument_w_nvtx + def _overflow_check_and_loss_scale_update(self): + + # First compute norm for all group so we know if there is overflow + if self.dtype == torch.float16: + self.check_overflow() + + #loss scaling related computation + prev_scale = self.loss_scale + self._update_scale(self.overflow) + + if self.overflow: + self._overflow_clean_up(prev_scale) + + #update loco error buf + self._loco_err_buf_update(self.overflow, self.loss_scale / prev_scale) + + return self.overflow + + @instrument_w_nvtx + def _post_step(self, timer_names): + if self.offload_optimizer: + self.reset_cpu_buffers() + + #Gathering persisting parameters + if len(self.persistent_parameters) > 0: + self.persistent_parameters[0].all_gather(self.persistent_parameters) + + if self.swap_optimizer: + self.optimizer_swapper.log_timers() + + self.invalidate_secondary_tensor() + + self.timers.log(timer_names) + + see_memory_usage('After zero_optimizer step', force=False) + print_rank_0(f"------------------Finishing Step-----------------------") + + @instrument_w_nvtx + def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): + if self.fp16_partitioned_groups_flat[sub_group_id] is not None: + self.fp16_partitioned_groups_flat[sub_group_id].data.copy_( + self.fp32_partitioned_groups_flat[sub_group_id].data) + + #unflatten fp16 parameter subgroup + self._unflatten_partitioned_parameters(sub_group_id) + else: + self._partitioned_params_swap_out(sub_group_id) + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + @instrument_w_nvtx + def step(self, closure=None): + """ + Not supporting closure. + """ + self._pre_step() + self._partition_all_parameters() + + #checks for overflow, adjust the loss scale accordingly + if self._overflow_check_and_loss_scale_update(): + if self.swap_optimizer: + self.optimizer_swapper.log_timers() + return + + norm_groups = self._get_norm_groups() + scaled_global_grad_norm = torch.linalg.vector_norm(torch.stack(norm_groups)) + + # Stash unscaled gradient norm + self._global_grad_norm = scaled_global_grad_norm / self.loss_scale + + timer_names = set() + + timer_names.add(OPTIMIZER_STEP_TIMER) + self.timers(OPTIMIZER_STEP_TIMER).start() + + #update parameters one sub group at a time + for sub_group_id, group in enumerate(self.fp16_groups): + + #prepare optimizer states, gradients and fp32 parameters for update + self._prepare_sub_group(sub_group_id, timer_names) + + #scale the fp32 gradients + self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) + + #apply the optimizer step on the sub group and copy fp32 parameters to fp16 + self._optimizer_step(sub_group_id) + + #put fp16 parameters in appropriate location + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + + #release memory or swap out optimizer states of fp32 parameters + self._release_sub_group(sub_group_id, timer_names) + + self.timers(OPTIMIZER_STEP_TIMER).stop() + + self._post_step(timer_names) + + # warn user about caching allocator flushes + memory_stats = get_accelerator().memory_stats() + alloc_retries = memory_stats.get("num_alloc_retries") + if alloc_retries is None: + alloc_retries = 0 + if alloc_retries > self.n_caching_allocator_flushes: + if dist.get_rank() == 0: + logger.warning( + "%d pytorch allocator cache flushes since last step. this happens " + "when there is high memory pressure and is detrimental to " + "performance. if this is happening frequently consider adjusting " + "settings to reduce memory consumption. If you are unable to " + "make the cache flushes go away consider adding " + "get_accelerator().empty_cache() calls in your training loop to ensure " + "that all ranks flush their caches at the same time", + alloc_retries - self.n_caching_allocator_flushes) + self.n_caching_allocator_flushes = alloc_retries + + def dump_pre_step_gradients(self, debug_fp32_grads): + # Dump gradient norms for debugging + for i, _ in enumerate(self.fp16_groups): + print(f'Pre-Step Dump Norms for Group {i} FP16P, FP16G, FP32G, FP32GUC') + for fp16_param, fp32_grad in zip(self.fp16_groups[i], debug_fp32_grads[i]): + param_id = self.get_param_id(fp16_param) + fp16_grad_norm = self.debug_fp16_grads[i][param_id] + + fp32_grad_norm = [float(t.data.float().norm(2)) for t in fp32_grad] + norm_list = [fp16_grad_norm, fp32_grad_norm] + print(f'Pre-Step Norms {i} {param_id} = {norm_list}') + + def dump_post_step_gradients(self): + # Dump gradient norms for debugging + for i, group in enumerate(self.fp16_groups): + print(f'Post-Step Dump Norms for Group {i} FP16P, FP16DS, FP16FLAT, FP32FLAT') + unflat_fp16 = self.unflatten(self.fp16_groups_flat[i], self.fp16_groups[i]) + unflat_fp32 = self.unflatten(self.fp32_partitioned_groups_flat[i], self.fp16_groups[i]) + for j, p in enumerate(self.fp16_groups[i]): + param_id = self.get_param_id(p) + param_norm = float(p.data.float().norm(2)) + ds_norm = float(p.ds_tensor.data.float().norm(2)) + + unflat_norm = [float(t.data.float().norm(2)) for t in [unflat_fp16[j], unflat_fp32[j]]] + norm_list = [param_norm, ds_norm] + unflat_norm + print(f'Post-Step Norms {i} {param_id} = {norm_list}') + + @instrument_w_nvtx + def unscale_and_clip_grads(self, sub_group_id, total_norm): + # compute combined scale factor for this group + combined_scale = self.loss_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad + clip = torch.clamp(clip, min=1.0) + combined_scale = clip * self.loss_scale + + self.fp32_partitioned_groups_flat[sub_group_id].grad.mul_(1. / combined_scale) + + def _check_overflow(self, partition_gradients=True): + self.overflow = self.has_overflow(partition_gradients) + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params, is_grad_list=False): + for p in params: + if p.grad is not None and self._has_inf_or_nan(p.grad.data): + return True + + return False + + def has_overflow_partitioned_grads_serial(self): + for i in range(len(self.fp16_groups)): + for j, grad in enumerate(self.averaged_gradients[i]): + if grad is not None and self._has_inf_or_nan(grad.data, j): + return True + return False + + @instrument_w_nvtx + def has_overflow(self, partition_gradients=True): + if partition_gradients: + with get_accelerator().stream(self.reduce_and_partition_stream): + if hasattr(self.inf_or_nan_tracker, "logical_or_"): + self.inf_or_nan_tracker.logical_or_(torch.isinf(self.grad_partitions_flat_buffer).any()) + self.inf_or_nan_tracker.logical_or_(torch.isnan(self.grad_partitions_flat_buffer).any()) + else: + # logical_or_ not available in older versions of pytorch + self.inf_or_nan_tracker += torch.isinf(self.grad_partitions_flat_buffer).any() + self.inf_or_nan_tracker += torch.isnan(self.grad_partitions_flat_buffer).any() + self.inf_or_nan_tracker = self.inf_or_nan_tracker > 0 + + overflow_gpu = self.inf_or_nan_tracker.clone().to(get_accelerator().current_device_name()).to( + torch.uint8) + self.inf_or_nan_tracker.zero_() + + if not get_accelerator().resolves_data_dependency(): + get_accelerator().default_stream().wait_stream(self.reduce_and_partition_stream) + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + else: + params = [] + for group in self.fp16_groups: + for param in group: + params.append(param) + + overflow = self.has_overflow_serial(params, is_grad_list=partition_gradients) + overflow_gpu = get_accelerator().ByteTensor([overflow]) + + # Since each model parallel GPU carries only part of the model, + # make sure overflow flag is synced across all the model parallel GPUs + self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX) + + overflow = overflow_gpu[0].item() + return bool(overflow) + + # `x` is a torch.Tensor + @staticmethod + def _has_inf_or_nan(x, j=None): + try: + # if x is half, the .float() incurs an additional deep copy, but it's necessary if + # Pytorch's .sum() creates a one-element tensor of the same type as x + # (which is true for some recent version of pytorch). + cpu_sum = float(x.float().sum()) + # More efficient version that can be used if .sum() returns a Python scalar + # cpu_sum = float(x.sum()) + except RuntimeError as instance: + # We want to check if inst is actually an overflow exception. + # RuntimeError could come from a different error. + # If so, we still want the exception to propagate. + if "value cannot be converted" not in instance.args[0]: + raise + return True + else: + if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: + return True + return False + + @instrument_w_nvtx + def backward(self, loss, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + if self.swap_optimizer: + self.optimizer_swapper.pre_backward() + + see_memory_usage(f"Before backward", force=False) + + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward() + else: + self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) + + if self.swap_optimizer: + self.optimizer_swapper.post_backward() + + def get_fp32_grad_partitions(self) -> Dict[int, Dict[int, Tensor]]: + """get fp32 gradient partition dictionary + accessed as grad_dict[parameter_group_index][parameter_index] + """ + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + grad_dict = collections.defaultdict(dict) + if self.offload_optimizer: + for group in self.fp16_groups: + for param_idx, param in enumerate(group): + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + grad_dict[group_idx][param_idx] = fp32_grad + else: + for group_idx, group in self.averaged_gradients.items(): + for param_idx, gradient in enumerate(group): + grad_dict[group_idx][param_idx] = gradient.float() + + return grad_dict + + def _fp32_state_allgather(self, param, fp32_state_partition): + reduce_buffer = torch.empty(self.partition_count * fp32_state_partition.numel(), + dtype=torch.float32, + device=param.device) + my_rank = dist.get_rank(group=self.dp_process_group) + partition = reduce_buffer.narrow(0, fp32_state_partition.numel() * my_rank, fp32_state_partition.numel()) + partition.data.copy_(fp32_state_partition.data, non_blocking=False) + dist.all_gather_into_tensor(reduce_buffer, partition, group=self.dp_process_group) + return reduce_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape) + + def _get_fp32_grad_state_partition(self, param, release_swap_buffers): + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + if self.offload_optimizer: + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_in(group_idx) + + fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + + if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: + self._release_swap_buffers(group_idx) + else: + fp32_grad = self.__param_id_to_grad_partition[param.ds_id] + + return fp32_grad, group_idx + + def get_fp32_grad_for_param(self, param) -> Tensor: + if not param.requires_grad: + return None + + fp32_grad, _ = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=True) + fp32_grad = fp32_grad.to(get_accelerator().current_device_name()).float() + return self._fp32_state_allgather(param, fp32_grad) + + def set_fp32_grad_for_param(self, value, param): + if not param.requires_grad: + return + + # if not get_accelerator().resolves_data_dependency(): + # self.reduce_and_partition_stream.synchronize() + + # if self.offload_optimizer: + # group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + # fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + # else: + # fp32_grad = self.__param_id_to_grad_partition[param.ds_id] + + fp32_grad, group_idx = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=False) + # import pdb; pdb.set_trace() + my_rank = dist.get_rank(group=self.dp_process_group) + value_partition = value.flatten().narrow(0, fp32_grad.numel() * my_rank, fp32_grad.numel()) + fp32_grad.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._writeback_swap_state(group_idx, write_opt_state=False, write_gradients=True) + + def _get_fp32_opt_state_partition(self, param, release_swap_buffers, optim_state_key=None): + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_in(group_idx) + + fp32_param = self.fp32_partitioned_groups_flat[group_idx] + if optim_state_key is None: + fp32_opt_state = fp32_param.narrow(0, dest_offset, num_elements) + else: + fp32_opt_state = self.optimizer.state[fp32_param][optim_state_key].narrow(0, dest_offset, num_elements) + + if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: + self._release_swap_buffers(group_idx) + + return fp32_opt_state, group_idx + + def get_full_hp_param(self, param, optim_state_key=None) -> Tensor: + if not param.requires_grad: + return None + + # import pdb; pdb.set_trace() + fp32_opt_state, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=True, + optim_state_key=optim_state_key) + fp32_opt_state = fp32_opt_state.to(get_accelerator().current_device_name()) + hp_param = self._fp32_state_allgather(param, fp32_opt_state) + + return hp_param + + def set_full_hp_param(self, value, param, optim_state_key=None): + if not param.requires_grad: + return + + assert value.numel( + ) == param.ds_numel, f" Number of elements do not match: {value.numel()} != {param.ds_numel}" + + fp32_opt_state_partition, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=False, + optim_state_key=optim_state_key) + # print(f'{dist.get_rank()=} {fp32_opt_state_partition.shape=} -------- {value.shape=}') + # import pdb; pdb.set_trace() + my_rank = dist.get_rank(group=self.dp_process_group) + value_partition = value.flatten().narrow(0, + fp32_opt_state_partition.numel() * my_rank, + fp32_opt_state_partition.numel()) + fp32_opt_state_partition.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_out(group_idx) + + ### Local API START ### + def get_local_fp32_grad_for_param(self, param) -> Tensor: + if not param.requires_grad: + return None + + fp32_grad, _ = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=True) + fp32_grad = fp32_grad.to(get_accelerator().current_device_name()).float() + return fp32_grad + + def set_local_grad_for_param(self, value, param): + if not param.requires_grad: + return + + assert value.numel() == param.ds_tensor.numel( + ), f" Number of elements do not match: {value.numel()} != {param.ds_tensor.ds_numel}" + + # if not get_accelerator().resolves_data_dependency(): + # self.reduce_and_partition_stream.synchronize() + + # if self.offload_optimizer: + # group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + # fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) + # else: + # fp32_grad = self.__param_id_to_grad_partition[param.ds_id] + + if self.offload_optimizer: + self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(value) + + fp32_grad, group_idx = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=False) + fp32_grad.data.copy_(value.flatten().data) + + if self._swappable_optimizer_subgroup(group_idx): + self._writeback_swap_state(group_idx, write_opt_state=False, write_gradients=True) + + def get_local_fp32_param(self, param, optim_state_key=None) -> Tensor: + if not param.requires_grad: + return None + fp32_opt_state, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=True, + optim_state_key=optim_state_key) + fp32_opt_state = fp32_opt_state.to(get_accelerator().current_device_name()) + return fp32_opt_state + + def set_local_hp_param(self, value, param, optim_state_key=None): + if not param.requires_grad: + return + + assert hasattr(param, "ds_tensor"), f" The parameter does not contain the partitioned copy of the tensor." + assert value.numel() == param.ds_tensor.numel( + ), f" Number of elements do not match: {value.numel()} != {param.ds_tensor.ds_numel}" + + fp32_opt_state_partition, group_idx = self._get_fp32_opt_state_partition(param, + release_swap_buffers=False, + optim_state_key=optim_state_key) + value_partition = value.flatten() + fp32_opt_state_partition.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_out(group_idx) + # logger.info(f"[set_local_hp_param][update the params' value successfully]") + + ### Local API END ### + + ### Vectorized API BEGIN ### + def update_fp32_grad_for_param_vectorized(self, update_func, param_list): + params_with_grad = [p for p in param_list if p.requires_grad] + if not params_with_grad: + return + + if not get_accelerator().resolves_data_dependency(): + self.reduce_and_partition_stream.synchronize() + + subgroups = {} + for p in params_with_grad: + group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(p)] + param_entry = (p, dest_offset, num_elements) + if group_idx in subgroups.keys(): + subgroups[group_idx].append(param_entry) + else: + subgroups[group_idx] = [param_entry] + + for group_idx, group_params in subgroups.items(): + if self._swappable_optimizer_subgroup(group_idx): + self._optimizer_states_and_gradient_swap_in(group_idx) + + for param, dest_offset, num_elements in group_params: + if self.offload_optimizer: + fp32_grad_part = self.fp32_partitioned_groups_flat[group_idx].grad.narrow( + 0, dest_offset, num_elements) + else: + fp32_grad_part = self.__param_id_to_grad_partition[param.ds_id] + + fp32_grad_full = self._fp32_state_allgather(param, fp32_grad_part) + new_fp32_grad_full = update_func(fp32_grad_full, param) + my_rank = dist.get_rank(group=self.dp_process_group) + value_partition = new_fp32_grad_full.flatten().narrow(0, + fp32_grad_part.numel() * my_rank, + fp32_grad_part.numel()) + fp32_grad_part.data.copy_(value_partition.data) + + if self._swappable_optimizer_subgroup(group_idx): + self._writeback_swap_state(sub_group_id=group_idx, write_opt_state=False, write_gradients=True) + + ### Vectorized API END ### + + ### Device API BEGIN ### + def get_hp_param_device(self, param, optim_state_key=None) -> torch.device: + if not param.requires_grad: + return None + + fp32_opt_state, _ = self._get_fp32_opt_state_partition(param, + release_swap_buffers=True, + optim_state_key=optim_state_key) + return fp32_opt_state.device + + ### Device API END ### + + @instrument_w_nvtx + def _partition_all_parameters(self): + self.parameter_offload.partition_all_parameters() + + def check_overflow(self, partition_gradients=True): + self._check_overflow(partition_gradients) + + def _update_scale(self, has_overflow=False): + self.loss_scaler.update_scale(has_overflow) + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + self.trainable_param_groups = self._get_trainable_parameter_groups() + + param_groups = property(_get_param_groups, _set_param_groups) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.loss_scaler.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) + cur_scale = property(_get_loss_scale, _set_loss_scale) + + def _get_lean_tensors(self, padded_flattened_tensor, group_tensors, paddings): + # Remove paddings from flattened tensor + individual_tensors = self.unflatten(padded_flattened_tensor, group_tensors) + lean_lengths = [t.numel() - pad for t, pad in zip(group_tensors, paddings)] + lean_tensors = [t[:len] for t, len in zip(individual_tensors, lean_lengths)] + #logger.info(f'rank {dist.get_rank()}: lean_tensors = {[t.numel() for t in lean_tensors]}') + return lean_tensors + + #TODO REVISIT this for stage 3 + def get_lean_optimizer_state(self): + # Return optimizer states after removing paddings. + # This method assumes that each param group contains a single flattened tensor. + optimizer_groups_state = [] + + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + lean_state = {} + for key, value in self.optimizer.state[p].items(): + if torch.is_tensor(value): + padded_lens = [t.numel() for t in self.fp16_partitioned_groups[i]] + lean_state[key] = self._get_lean_tensors(value, self.fp16_partitioned_groups[i], + self.groups_padding[i]) + lean_flat_len = sum([t.numel() for t in lean_state[key]]) + else: + lean_state[key] = value + + optimizer_groups_state.append(lean_state) + + return optimizer_groups_state + + def get_groups_without_padding(self, groups_with_padding): + # Return group tensor after removing paddings added for alignment to DP world size. + groups_without_padding = [] + for i, group in enumerate(groups_with_padding): + lean_group = self._get_lean_tensors(group, self.fp16_partitioned_groups[i], self.groups_padding[i]) + groups_without_padding.append(lean_group) + + return groups_without_padding + + def _set_fp32_optimizer_param_groups(self): + for sub_group_id, _ in enumerate(self.fp16_groups): + param_group_id = self.sub_group_to_group_id[sub_group_id] + self.optimizer.param_groups[param_group_id]['params'].append( + self.fp32_partitioned_groups_flat[sub_group_id]) + + def _clear_fp32_optimizer_param_groups(self): + for param_group in self.optimizer.param_groups: + param_group['params'] = [] + + def _rigid_state_dict(self): + state_dict = {} + state_dict[ZERO_STAGE] = ZeroStageEnum.weights + state_dict[LOSS_SCALER] = self.loss_scaler + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['overflow'] = self.overflow + state_dict[PARTITION_COUNT] = self.partition_count + + self._set_fp32_optimizer_param_groups() + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + state_dict[FP32_FLAT_GROUPS] = self.fp32_partitioned_groups_flat + self._clear_fp32_optimizer_param_groups() + + return state_dict + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + if self.elastic_checkpoint: + raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.") + + return self._rigid_state_dict() + + +# Restore base optimizer fp32 weights from checkpoint by: +# 1) Merging fp32 weights from checkpoints of all partitions +# 2) Extracting fp32 weights for current partition from merged weights +# 3) Using extracted weights to update base optimizer weights directly. + + def _restore_from_fp32_weights(self, all_state_dict): + + flat_local_partition = [] + for i in range(len(self.fp32_partitioned_groups_flat)): + merged_partitions = [sd['fp32_groups'][i] for sd in all_state_dict] + flat_local_partition.append(self._get_flattened_partition(merged_partitions)) + + for current, saved in zip(self.fp32_partitioned_groups_flat, flat_local_partition): + current.data.copy_(saved.data) + + # Restore base optimizer fp32 weights from ZeRO fp16 weights + def _restore_from_bit16_weights(self): + for fp16_partitions, fp32_partition in zip(self.fp16_partitioned_groups_flat, + self.fp32_partitioned_groups_flat): + fp32_partition.data.copy_(fp16_partitions.data) + + # Refresh the fp32 master params from the fp16 copies. + def refresh_fp32_params(self): + self._restore_from_bit16_weights() + + # Extract flattened partition for current rank from all partitions + def _get_flattened_partition(self, all_partition_states): + partition_id = dist.get_rank(group=self.dp_process_group) + alignment = dist.get_world_size(group=self.dp_process_group) + + param_partitions = [[] for _ in range(len(all_partition_states[0]))] + for i, partition in enumerate(all_partition_states): + for j, param in enumerate(partition): + param_partitions[j].append(param) + + local_state_partitions = [] + for param_index, param_slices in enumerate(param_partitions): + flattened_merged_tensor = self.flatten_dense_tensors_aligned(param_slices, alignment) + new_partitions = self.get_data_parallel_partitions(flattened_merged_tensor) + local_state_partitions.append(new_partitions[partition_id]) + + if torch.is_tensor(local_state_partitions[0]): + return self.flatten_dense_tensors_aligned(local_state_partitions, alignment) + + # Assume non-tensor states are not partitioned and equal across ranks, so return first one + return local_state_partitions[0] + + # Restore base optimizer state from checkpoint by + # 1) Merging optimizer state from checkpoints of all partitions + # 2) Extracting optimizer state for current partition from the merged state + # 3) Using the extracted value to directly update the base optimizer. + def _restore_base_optimizer_state(self, all_state_dict): + base_optimizer_group_states = [] + for i in range(len(self.optimizer.param_groups)): + partition_states = {} + all_partition_group_states = [sd['base_optimizer_state'][i] for sd in all_state_dict] + for key in all_partition_group_states[0].keys(): + all_partition_states = [all_states[key] for all_states in all_partition_group_states] + partition_states[key] = self._get_flattened_partition(all_partition_states) + base_optimizer_group_states.append(partition_states) + + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + for key, saved in base_optimizer_group_states[i].items(): + if torch.is_tensor(self.optimizer.state[p][key]): + self.optimizer.state[p][key].data.copy_(saved.data) + else: + self.optimizer.state[p][key] = saved + + def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True): + # I think it should actually be ok to reload the optimizer before the model. + self.loss_scaler = state_dict[LOSS_SCALER] + self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] + self.overflow = state_dict['overflow'] + + if load_optimizer_states: + self._set_fp32_optimizer_param_groups() + self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) + self._clear_fp32_optimizer_param_groups() + + if self.swap_optimizer: + # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint + self.optimizer_swapper.purge_state() + + if self.swap_optimizer: + # Touch all parameters to synchronize all buffers + timer_names = set() + self._partition_all_parameters() + for sub_group_id, group in enumerate(self.fp16_groups): + self._prepare_sub_group(sub_group_id, timer_names) + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._release_sub_group(sub_group_id, timer_names) + self._post_step(timer_names) + + # restore fp32 partitions + for curr_param, saved_param in zip(self.fp32_partitioned_groups_flat, state_dict[FP32_FLAT_GROUPS]): + curr_param.data.copy_(saved_param.data) + + # restore fp16 partitions from fp32 + for sub_group_id in range(len(self.fp32_partitioned_groups_flat)): + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + if sum(fp32_param.size()) > 0: + fp16_param = self.fp16_partitioned_groups_flat[sub_group_id] + fp16_param.data.copy_(fp32_param.data) + + # update fp16 unflattened params + for sub_group_id in range(len(self.fp16_partitioned_groups_flat)): + updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]) + + for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): + partitioned_param.data = q.data + + # TODO: Support different/changing load/save DP degree. + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None, + param_shapes=None): + r"""Loading a ZeRO checkpoint + Arguments: + state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition. + Note that the number of saved partitions may differ from number of loading partitions to support + changing GPU count, specifically DP world size, between saving and loading checkpoints. + load_optimizer_states: Boolean indicating whether or not to load base optimizer states + load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32 + copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss). + """ + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + + if self.elastic_checkpoint: + raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.") + + if checkpoint_folder: + self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights, + param_shapes) + else: + self._rigid_load_state_dict(state_dict_list[dist.get_rank(group=self.dp_process_group)], + load_optimizer_states=load_optimizer_states) + + # when use loading checkpoint serial, after finish loading, we need to + # delete the temp state_dict_list variable to save memory, then trigger + # the next rank's loading + if load_serial is not None: + load_serial += 1 + rank = dist.get_rank(group=self.dp_process_group) + local_rank = dist.get_local_rank() + del state_dict_list[rank] + rank_end = dist.get_world_size() - 1 + if local_rank != rank_end: + dist.send(tensor=load_serial, dst=rank + 1) + + if len(self.persistent_parameters) > 0: + self.persistent_parameters[0].partition(self.persistent_parameters) + # self.persistent_parameters[0].all_gather(self.persistent_parameters) # this will be done in checkpoint_event_epilogue() so remove it to prevent double all_gather + + def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights, + param_shapes): + self.load_hp_checkpoint_state_from_checkpoint_dir_stage3(checkpoint_folder, param_shapes) + + def load_hp_checkpoint_state_from_checkpoint_dir_stage3(self, checkpoint_dir, param_shapes): + """ Load optimizer and model states from the checkpoint directory. """ + checkpoint_dir = os.path.join(checkpoint_dir, "zero") + optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") + assert os.path.isfile( + optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' + + optim_sd = torch.load(optim_state_path, weights_only=False) + self._load_global_state_stage3(optim_sd) + + key_list = ["fp32", "exp_avg", "exp_avg_sq"] + + for key in key_list: + key_tensor = torch.empty(0) + for layer in param_shapes[0].keys(): + key_layer_state_partition = self.load_hp_checkpoint_state(os.path.join(checkpoint_dir, layer), key) + key_tensor = torch.cat((key_tensor, key_layer_state_partition)) + if key == "fp32": + self.fp32_partitioned_groups_flat[0].data.copy_(key_tensor) + self.optimizer.param_groups[0]['params'].append(self.fp32_partitioned_groups_flat[0]) + else: + optim_sd[OPTIMIZER_STATE_DICT]['state'][0][key] = key_tensor + + if self.swap_optimizer: + # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint + self.optimizer_swapper.purge_state() + + if self.swap_optimizer: + # Touch all parameters to synchronize all buffers + timer_names = set() + self._partition_all_parameters() + for sub_group_id, group in enumerate(self.fp16_groups): + self._prepare_sub_group(sub_group_id, timer_names) + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._release_sub_group(sub_group_id, timer_names) + self._post_step(timer_names) + + self.optimizer.load_state_dict(optim_sd[OPTIMIZER_STATE_DICT]) + for param_group in self.optimizer.param_groups: + param_group['params'] = [] + + for sub_group_id in range(len(self.fp32_partitioned_groups_flat)): + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + if sum(fp32_param.size()) > 0: + fp16_param = self.fp16_partitioned_groups_flat[sub_group_id] + fp16_param.data.copy_(fp32_param.data) + + for sub_group_id in range(len(self.fp16_partitioned_groups_flat)): + updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], + self.fp16_partitioned_groups[sub_group_id]) + + for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): + partitioned_param.data = q.data + + def _load_global_state_stage3(self, sd): + self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler) + self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale) + self.overflow = sd.get('overflow', self.overflow) + + def load_hp_checkpoint_state(self, folder, key): + local_rank = dist.get_local_rank() + + # Load tensors from files and reshape them to flat vectors + loaded_checkpoint_state = torch.load(os.path.join(folder, f"{key}.pt"), weights_only=False).view(-1) + + # Partition the loaded data according to the local rank + world_size = dist.get_world_size(group=self.dp_process_group) + unpartitioned_numel = loaded_checkpoint_state.numel() + partitioned_numel = math.ceil(unpartitioned_numel / world_size) + + if world_size * partitioned_numel != unpartitioned_numel: + padding_size = world_size * partitioned_numel - unpartitioned_numel + padding_tensor = torch.zeros(padding_size, dtype=loaded_checkpoint_state.dtype) + loaded_checkpoint_state = torch.cat([loaded_checkpoint_state, padding_tensor]) + checkpoint_state_partition = loaded_checkpoint_state.narrow(0, local_rank * partitioned_numel, + partitioned_numel) + + return checkpoint_state_partition + + def reset_swap_buffers(self): + timer_names = set() + for sub_group_id, group in enumerate(self.fp16_groups): + self._prepare_sub_group(sub_group_id, timer_names) + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._release_sub_group(sub_group_id, timer_names) + + def checkpoint_event_prologue(self): + self._partition_all_parameters() + + def checkpoint_event_epilogue(self): + if len(self.persistent_parameters) > 0: + self.persistent_parameters[0].all_gather(self.persistent_parameters) + + def empty_partition_cache(self): + self.parameter_offload.empty_partition_cache() + + def offload_states(self, + include: Container[OffloadStateTypeEnum] = None, + device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, + pin_memory: bool = True, + non_blocking: bool = False): + device = device.value + + self.empty_partition_cache() + + assert self.optimizer.__class__ == deepspeed.ops.adam.fused_adam.FusedAdam, f"Offloading is supported only for DeepSpeed FusedAdam." + + def needs_offload(target): + # return True + return target not in self.offloaded_states and (include == None or target in include) + + # HP param + if needs_offload(OffloadStateTypeEnum.hp_params): + if pin_memory: + if not hasattr(self, "hp_params_pin_buffers"): + self.hp_params_pin_buffers = [ + get_accelerator().pin_memory(torch.empty_like(t, device=device)) + for t in self.fp32_partitioned_groups_flat + ] + + for src_tensor, dest_buf in zip(self.fp32_partitioned_groups_flat, self.hp_params_pin_buffers): + dest_buf.copy_(src_tensor, non_blocking=non_blocking) + src_tensor.data = dest_buf + else: + for buf in self.fp32_partitioned_groups_flat: + buf.data = buf.data.to(device, non_blocking=non_blocking) + self.offloaded_states.add(OffloadStateTypeEnum.hp_params) + + # LP param + if needs_offload(OffloadStateTypeEnum.lp_params): + if pin_memory: + if not hasattr(self, "lp_param_contiguous_pin_buffer"): + self.lp_param_contiguous_pin_buffer = get_accelerator().pin_memory( + torch.empty_like(self.lp_param_buffer, device=device)) + self.lp_param_contiguous_pin_buffer.copy_(self.lp_param_buffer, non_blocking=non_blocking) + cpu_buffer = self.lp_param_contiguous_pin_buffer + else: + cpu_buffer = self.lp_param_buffer.to(device, non_blocking=non_blocking) + + self.lp_param_buffer.data = cpu_buffer + for tensor, offset, tensor_numel in get_mapping_to_flat_buffer( + [p.ds_tensor for p in self.module.parameters()]): + tensor.data = cpu_buffer.narrow(0, offset, tensor_numel) + + self.fp16_partitioned_groups_flat.clear() + self.offloaded_states.add(OffloadStateTypeEnum.lp_params) + + # LP grad + if needs_offload(OffloadStateTypeEnum.lp_grads): + if pin_memory: + if not hasattr(self, "lp_grad_partitions_flat_pin_buffers"): + self.lp_grad_partitions_flat_pin_buffers = get_accelerator().pin_memory( + torch.empty_like(self.grad_partitions_flat_buffer, device=device)) + self.lp_grad_partitions_flat_pin_buffers.copy_(self.grad_partitions_flat_buffer, + non_blocking=non_blocking) + self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers + else: + self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to(device) + self.averaged_gradients = {} + + self.__param_id_to_grad_partition = {} + + self.offloaded_states.add(OffloadStateTypeEnum.lp_grads) + + # contiguous bucket + if needs_offload(OffloadStateTypeEnum.contiguous_grad_buffer): + if hasattr(self, "_DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer" + ) and self.__ipg_bucket_flat_buffer is not None: + # Record properties like shape, strides, etc. as a meta tensor + self.grad_buffer_meta = self.__ipg_bucket_flat_buffer.to("meta") + self.__ipg_bucket_flat_buffer = None + self.offloaded_states.add(OffloadStateTypeEnum.contiguous_grad_buffer) + + # Adam + if needs_offload(OffloadStateTypeEnum.optim_states): + offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) + self.offloaded_states.add(OffloadStateTypeEnum.optim_states) + + gc.collect() + get_accelerator().empty_cache() + + def reload_states(self, non_blocking: bool = False): + + device = get_accelerator().current_device_name() + + # HP param + if OffloadStateTypeEnum.hp_params in self.offloaded_states: + if hasattr(self, "hp_params_pin_buffers"): + for src, dest in zip(self.hp_params_pin_buffers, self.fp32_partitioned_groups_flat): + dest.data = src.to(device, non_blocking=non_blocking) + else: + for buf in self.fp32_partitioned_groups_flat: + buf.data = buf.data.to(device, non_blocking=non_blocking) + self.offloaded_states.remove(OffloadStateTypeEnum.hp_params) + + # LP Param + if OffloadStateTypeEnum.lp_params in self.offloaded_states: + cpu_buffer = self.lp_param_contiguous_pin_buffer if hasattr( + self, "lp_param_contiguous_pin_buffer") else self.lp_param_buffer + self.lp_param_buffer.data = cpu_buffer.data.to(device, non_blocking=non_blocking) + self._set_fp16_partitioned_groups_flat() + + parameter_partitions = self._get_parameter_partitions() + for tensor, offset, tensor_numel in get_mapping_to_flat_buffer(parameter_partitions): + tensor.data = self.lp_param_buffer.narrow(0, offset, tensor_numel) + self.offloaded_states.remove(OffloadStateTypeEnum.lp_params) + + # LP grad + if OffloadStateTypeEnum.lp_grads in self.offloaded_states: + if hasattr(self, "lp_grad_partitions_flat_pin_buffers"): + self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers.to( + device, non_blocking=non_blocking) + else: + self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to( + device, non_blocking=non_blocking) + self.averaged_gradients = {} + + offset = 0 + all_params = list(itertools.chain.from_iterable(self.fp16_groups)) + for param in all_params: + self.__param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( + 0, offset, param.partition_numel()) + offset += param.partition_numel() + + self.offloaded_states.remove(OffloadStateTypeEnum.lp_grads) + + # contiguous bucket + if OffloadStateTypeEnum.contiguous_grad_buffer in self.offloaded_states: + self.__ipg_bucket_flat_buffer = torch.empty_like(self.grad_buffer_meta, device=device) + # self.__ipg_bucket_flat_buffer.data = self.__ipg_bucket_flat_buffer.data.to(device) + self.offloaded_states.remove(OffloadStateTypeEnum.contiguous_grad_buffer) + + # Adam + if OffloadStateTypeEnum.optim_states in self.offloaded_states: + reload_adam_states(self.optimizer, device, non_blocking=non_blocking) + self.offloaded_states.remove(OffloadStateTypeEnum.optim_states) + + if non_blocking: + get_accelerator().synchronize() + + +def _handle_overflow(cpu_sum, x, i): + import math + rank = dist.get_rank() + if rank == 0: + t_i = -1 + for v_i, v in enumerate(x.data.contiguous().view(-1)): + if not math.isfinite(float(v)): + t_i = v_i + break + logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") + + +def estimate_zero3_model_states_mem_needs(total_params, + largest_layer_params, + num_gpus_per_node=1, + num_nodes=1, + cpu_offload=True, + cpu_offload_params=True, + zero_init=True, + additional_buffer_factor=1.5): + + total_gpus = num_nodes * num_gpus_per_node + gpus_factor = 1 / num_nodes + largest_layer_memory = (4 * largest_layer_params) + + if cpu_offload: + if cpu_offload_params: + gpu_mem = largest_layer_memory + + if zero_init: + cpu_mem = total_params * 18 * gpus_factor * additional_buffer_factor + else: + cpu_mem = total_params * max(4 * num_gpus_per_node, 18 * gpus_factor) * additional_buffer_factor + + else: + gpu_mem = largest_layer_memory + int(2 * total_params / total_gpus) + + if zero_init: + cpu_mem = total_params * 16 * gpus_factor * additional_buffer_factor + else: + cpu_mem = total_params * max(4 * num_gpus_per_node, 16 * gpus_factor) * additional_buffer_factor + else: + gpu_mem = largest_layer_memory + int(18 * total_params / total_gpus) + if zero_init: + cpu_mem = largest_layer_params * 4 * num_gpus_per_node * additional_buffer_factor + else: + cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor + + return int(cpu_mem), int(gpu_mem), largest_layer_memory + + +def model_to_params(model): + # shared params calculated only once + total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) + + largest_layer_params = 0 + for m in model.modules(): + # assuming no shared params within a single layer + layer_params = sum(p.numel() for p in m.parameters(recurse=False)) + largest_layer_params = max(largest_layer_params, layer_params) + + return total_params, largest_layer_params + + +def estimate_zero3_model_states_mem_needs_all_live(model, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients + for a given ``model`` and hardware setup. + + If you have an actual model object, use this function and everything will be derived + automatically. + + If it's a hypothetical model, use ``estimate_zero3_model_states_mem_needs_all_cold`` where you have to pass + the ``total_params`` and ``largest_layer_params`` explicitly. + + Args: + - ``model``: ``nn.Module`` object + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + total_params, largest_layer_params = model_to_params(model) + + estimate_zero3_model_states_mem_needs_all_cold(total_params=total_params, + largest_layer_params=largest_layer_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + additional_buffer_factor=additional_buffer_factor) + + +def estimate_zero3_model_states_mem_needs_all_cold(total_params, + largest_layer_params, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients + for a given ``model`` and hardware setup. + + If it's a hypothetical model, use this function where you have to pass + the ``total_params`` and ``largest_layer_params`` explicitly. + + If you have an actual model object, use ``estimate_zero3_model_states_mem_needs_all_live`` and everything + will be derived automatically. + + Args: + - ``total_params``: total model params + - ``largest_layer_params``: largest layer's params + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + def format_options(cpu_offload, cpu_offload_params, zero_init): + enabled = [] + padded_cpu_str = f'{OffloadDeviceEnum.cpu:4}' + param_device = padded_cpu_str if cpu_offload_params else "none" + enabled.append(f"offload_param={param_device}") + optimizer_device = padded_cpu_str if cpu_offload else "none" + enabled.append(f"offload_optimizer={optimizer_device}") + enabled.append(f"zero_init={1 if zero_init else 0}") + return ", ".join(enabled) + + nodes_str = "nodes" if num_nodes > 1 else "node" + gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU" + print( + "Estimated memory needed for params, optim states and gradients for a:\n" + f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n" + f"SW: Model with {int(total_params/1e6)}M total params, {int(largest_layer_params/1e6)}M largest layer params." + ) + print(" per CPU | per GPU | Options") + for cpu_offload in [True, False]: + for cpu_offload_params in [True, False]: + if not cpu_offload and cpu_offload_params: + continue + for zero_init in [True, False]: + cpu_mem, gpu_mem, largest_layer_memory = estimate_zero3_model_states_mem_needs( + total_params=total_params, + largest_layer_params=largest_layer_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + cpu_offload=cpu_offload, + cpu_offload_params=cpu_offload_params, + zero_init=zero_init, + additional_buffer_factor=additional_buffer_factor) + + options_str = format_options(cpu_offload=cpu_offload, + cpu_offload_params=cpu_offload_params, + zero_init=zero_init) + print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}") diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/stage_1_and_2.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/stage_1_and_2.py new file mode 100644 index 0000000000000000000000000000000000000000..30c1d9525802839772f7ec7a14c8457c13a30ac9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/stage_1_and_2.py @@ -0,0 +1,2557 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed import comm as dist +from packaging import version as pkg_version +from collections import OrderedDict +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + +from deepspeed.runtime.base_optimizer import ZeROOptimizer +from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler +from deepspeed.runtime.utils import (empty_cache, see_memory_usage, inf, is_model_parallel_parameter, + align_dense_tensors, all_gather_dp_groups, mask_nan_or_inf_with_val_inplace) +from deepspeed.runtime.zero.config import ZeroStageEnum +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.ops.adam import DeepSpeedCPUAdam +from deepspeed.utils import logger +from deepspeed.utils.torch import register_grad_hook +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank +from deepspeed.moe.utils import is_moe_param +from deepspeed.git_version_info import version + +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.accelerator import get_accelerator + +from deepspeed.checkpoint.constants import (DS_VERSION, GROUP_PADDINGS, PARTITION_COUNT, LOSS_SCALER, + SINGLE_PARTITION_OF_FP32_GROUPS, BASE_OPTIMIZER_STATE, + BASE_OPTIMIZER_STATE_STEP, CLIP_GRAD, ZERO_STAGE, PARAM_SLICE_MAPPINGS) +from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state +from deepspeed.checkpoint import enable_universal_checkpoint + +from deepspeed.utils import groups +# Toggle this to true to enable correctness test +# with gradient partitioning and without +pg_correctness_test = False + +OPTIMIZER_ALLGATHER_TIMER = 'optimizer_allgather' +OPTIMIZER_GRADIENTS_TIMER = 'optimizer_gradients' +OPTIMIZER_STEP_TIMER = 'optimizer_step' +OPTIMIZER_TIMERS = [OPTIMIZER_ALLGATHER_TIMER, OPTIMIZER_GRADIENTS_TIMER, OPTIMIZER_STEP_TIMER] +INITIAL_MICRO_STEP_ID = -1 + + +def input(msg): + return + + +def split_half_float_double(tensors): + device_type = get_accelerator().device_name() + dtypes = [ + "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type), + "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type) + ] + buckets = [] + for i, dtype in enumerate(dtypes): + bucket = [t for t in tensors if t.type() == dtype] + if bucket: + buckets.append(bucket) + return buckets + + +def isclose(a, b, rtol=1e-09, atol=0.0): + return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol) + + +def lcm(x, y): + from fractions import gcd # or can import gcd from `math` in Python 3 + return x * y // gcd(x, y) + + +def get_alignment_padding(tensor_list, alignment): + num_elements = sum([tensor.numel() for tensor in tensor_list]) + remainder = num_elements % alignment + return (alignment - remainder) if remainder else remainder + + +def print_rank_msg(msg): + print(f"rank {dist.get_rank()} - {msg}") + + +def _get_padded_tensor(src_tensor, size): + if src_tensor.numel() >= size: + return src_tensor + padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) + slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) + slice_tensor.data.copy_(src_tensor.data) + return padded_tensor + + +def _pad_tensor_by_size(src_tensor, pad_size, dtype, device): + padded_tensor = torch.zeros(src_tensor.numel() + pad_size, dtype=dtype, device=device) + padded_tensor.data[:src_tensor.numel()].copy_(src_tensor.data) + return padded_tensor + + +class DeepSpeedZeroOptimizer(ZeROOptimizer): + """ + DeepSpeedZeroOptimizer designed to reduce the memory footprint + required for training large deep learning models. + + For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models + https://arxiv.org/abs/1910.02054 + + For usage examples, refer to TODO: DeepSpeed Tutorial + + """ + + def __init__(self, + init_optimizer, + param_names, + timers, + static_loss_scale=1.0, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + contiguous_gradients=True, + reduce_bucket_size=500000000, + use_multi_rank_bucket_allreduce=True, + allgather_bucket_size=5000000000, + dp_process_group=None, + expert_parallel_group=None, + expert_data_parallel_group=None, + reduce_scatter=True, + overlap_comm=False, + offload_optimizer_config=None, + mpu=None, + clip_grad=0.0, + gradient_accumulation_dtype=torch.float32, + communication_data_type=torch.float16, + postscale_gradients=True, + gradient_predivide_factor=1.0, + gradient_accumulation_steps=1, + ignore_unused_parameters=True, + partition_grads=True, + round_robin_gradients=False, + has_moe_layers=False, + fp16_master_weights_and_gradients=False, + elastic_checkpoint=False): + + if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none: + self.cpu_offload = True + self.cpu_offload_pin_memory = offload_optimizer_config.pin_memory + else: + self.cpu_offload = False + self.cpu_offload_pin_memory = False + + if dist.get_rank() == 0: + logger.info(f"Reduce bucket size {reduce_bucket_size}") + logger.info(f"Allgather bucket size {allgather_bucket_size}") + logger.info(f"CPU Offload: {self.cpu_offload}") + logger.info(f'Round robin gradient partitioning: {round_robin_gradients}') + # The fused optimizer does all the work. We need this layer for two reason: + # 1. maintain same user API from apex.fp16_utils + # 2. keep common stuff here in case we need to add ne552w fused optimizer later + + self.elastic_checkpoint = elastic_checkpoint + self.param_names = param_names + self.mpu = mpu + # differences from apex.fp16_utils: + # - assume all model params in fp16 + # - assume all params requires grad + # - flat by groups, not keeping state. TODO: remove state explicitly? + # - master grad and unflat master weight never exist. TODO: a way to save out unflat master? + if not get_accelerator().is_available(): + raise SystemError("Accelerator is not detected, cannot perform low precision training (e.g., fp16, bf16).") + self.optimizer = init_optimizer + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + + # ZeRO stage 1 (False) or 2 (True) + self.partition_gradients = partition_grads + self.zero_stage_string = "ZeRO-2" if partition_grads else "ZeRO-1" + + self.timers = timers + + self.reduce_scatter = reduce_scatter + + self.overlap_comm = overlap_comm + + self.deepspeed_adam_offload = self.cpu_offload + + self.device = get_accelerator().current_device_name() if not self.cpu_offload else 'cpu' + + self.dp_process_group = dp_process_group + self.sequence_parallel_size = groups._get_sequence_parallel_world_size() + #expert parallel group + self.ep_process_group = expert_parallel_group + + #data parallel group for experts + self.expert_dp_process_group = expert_data_parallel_group + + #data parallel size for non-experts + dp_size = dist.get_world_size(group=self.dp_process_group) + + #For MoE models this maybe different for different param group + #It will be modified during MoE setup later in the init + self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] + self.partition_count = [dp_size for i in range(len(self.optimizer.param_groups))] + + self.is_gradient_accumulation_boundary = True + + # CPU-Offload requires contiguous gradients + self.contiguous_gradients = contiguous_gradients or self.cpu_offload + + self.has_moe_layers = has_moe_layers + if self.has_moe_layers: + self._configure_moe_settings() + self._global_grad_norm = 0. + + if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): + self.model_parallel_group = None + self.model_parallel_world_size = 1 + self.model_parallel_rank = 0 + else: + self.model_parallel_group = mpu.get_model_parallel_group() + self.model_parallel_world_size = mpu.get_model_parallel_world_size() + self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu) + + self.overflow = False + self.clip_grad = clip_grad + self.communication_data_type = communication_data_type + self.gradient_predivide_factor = gradient_predivide_factor + self.postscale_gradients = postscale_gradients + self.gradient_accumulation_steps = gradient_accumulation_steps + self.micro_step_id = INITIAL_MICRO_STEP_ID + self.ignore_unused_parameters = ignore_unused_parameters + self.round_robin_gradients = round_robin_gradients + + self.extra_large_param_to_reduce = None + self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients + + if self.fp16_master_weights_and_gradients: + assert self.cpu_offload and type(self.optimizer) in [DeepSpeedCPUAdam], \ + f"fp16_master_and_gradients requires optimizer to support keeping fp16 master and gradients while keeping the optimizer states in fp32."\ + f"Currently only supported using ZeRO-Offload with DeepSpeedCPUAdam. But current setting is ZeRO-Offload:{self.cpu_offload} and optimizer type {type(self.optimizer)}." \ + f"Either disable fp16_master_weights_and_gradients or enable {self.zero_stage_string} Offload with DeepSpeedCPUAdam." + + if self.reduce_scatter and self.partition_gradients: + valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32) + assert self.communication_data_type in valid_reduce_scatter_dtypes, f"{self.zero_stage_string} supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'" + assert self.gradient_predivide_factor == 1.0, f"gradient_predivide_factor != 1.0 is not yet supported with {self.zero_stage_string} with reduce scatter enabled" + assert self.postscale_gradients, f"pre-scale gradients is not yet supported with {self.zero_stage_string} with reduce scatter enabled" + + # param flattened by groups + self.bit16_groups = [] + self.bit16_groups_flat = [] + + # param partitioned by data parallel degree + # this will contain a list of equal sized tensors + # each of which will be updated by a different process + self.parallel_partitioned_bit16_groups = [] + + # a single 32-bit partition of the parallel partitioned parameters + # that this process will update + self.single_partition_of_fp32_groups = [] + + # a 16-bit CPU param buffer for cpu offload + if self.cpu_offload: + self.param_buffer_of_bit16_for_cpu_offload_groups = [] + + # param partition info + + # These are the parameters in each group that will not be updated by this process directly + self.params_not_in_partition = [] + + # These are the parameters that will be updated by this process directly + self.params_in_partition = [] + + # Offset from the first parameter in the self.params_in_partition + # the parameter boundaries may not align with partition boundaries + # so we need to keep track of the offset + self.first_offset = [] + + # number of elements per partition in each group + self.partition_size = [] + + # align nccl all-gather send buffers to 4-byte boundary + self.nccl_start_alignment_factor = 2 # 4-byte alignment/sizeof(fp16) = 2 + + assert ( + allgather_bucket_size % self.nccl_start_alignment_factor == 0 + ), f"allgather_bucket_size must be a multiple of nccl_start_alignment_factor, {self.nccl_start_alignment_factor} " + + self.all_reduce_print = False + self.dtype = self.optimizer.param_groups[0]['params'][0].dtype + self.gradient_accumulation_dtype = gradient_accumulation_dtype + + if self.dtype != self.gradient_accumulation_dtype: + self.use_separate_grad_accum = True + else: + self.use_separate_grad_accum = False + if self.use_separate_grad_accum and not self.partition_gradients: + self.use_grad_accum_attribute = True + else: + self.use_grad_accum_attribute = False + + self.round_robin_bit16_groups = [] + self.round_robin_bit16_indices = [] + self.round_robin_bit16_meta = [] + + # Use different parallel to do all_to_all_reduce related things + # padding on each partition for alignment purposes + self.groups_padding = [] + # loop to deal with groups + for i, param_group in enumerate(self.optimizer.param_groups): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + + # push this group to list before modify + # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group + trainable_parameters = [] + for param in param_group['params']: + if param.requires_grad: + param.grad_accum = None + param.param_idx_in_group = len(trainable_parameters) + trainable_parameters.append(param) + self.bit16_groups.append(trainable_parameters) + + # not sure why apex was cloning the weights before flattening + # removing cloning here + + see_memory_usage(f"Before moving param group {i} to CPU") + # move all the parameters to cpu to free up GPU space for creating flat buffer + + # Create temp CPU param copies, free accelerator tensors + orig_group_numel = 0 + for param in self.bit16_groups[i]: + orig_group_numel += param.numel() + param.cpu_data = param.data.cpu() + param.data = torch.empty(1).to(param.device) + + empty_cache() + see_memory_usage(f"After moving param group {i} to CPU", force=False) + + # Reorder group parameters for load balancing of gradient partitioning during backward among ranks. + # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks. + # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging + # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m). + if self.round_robin_gradients: + round_robin_tensors, round_robin_indices = self._round_robin_reorder( + self.bit16_groups[i], dist.get_world_size(group=self.real_dp_process_group[i])) + else: + round_robin_tensors = self.bit16_groups[i] + round_robin_indices = list(range(len(self.bit16_groups[i]))) + + self.round_robin_bit16_groups.append(round_robin_tensors) + self.round_robin_bit16_indices.append(round_robin_indices) + + # Create meta tensors list, ordered according to round_robin_tensors + meta_tensors = [] + for param in round_robin_tensors: + meta_tensors.append(torch.zeros_like(param.cpu_data, device="meta")) + self.round_robin_bit16_meta.append(meta_tensors) + + # create flat buffer in CPU + flattened_buffer = self.flatten_dense_tensors_aligned( + self.round_robin_bit16_groups[i], + self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i]), + use_cpu_data=True) + + # free temp CPU params + for param in self.bit16_groups[i]: + del param.cpu_data + + # Move CPU flat tensor to the accelerator memory. + self.bit16_groups_flat.append(flattened_buffer.to(get_accelerator().current_device_name())) + del flattened_buffer + + see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False) + + if dist.get_rank(group=self.real_dp_process_group[i]) == 0: + see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False) + + # set model bit16 weight to slices of flattened buffer + self._update_model_bit16_weights(i) + + # divide the flat weights into near equal partition equal to the data parallel degree + # each process will compute on a different part of the partition + data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i) + self.parallel_partitioned_bit16_groups.append(data_parallel_partitions) + + # Record padding required for alignment + left_boundary = sum([t.numel() for t in data_parallel_partitions[:partition_id]]) + curr_partition_size = data_parallel_partitions[partition_id].numel() + + if orig_group_numel <= left_boundary: + padding = curr_partition_size + elif orig_group_numel < left_boundary + curr_partition_size: + padding = left_boundary + curr_partition_size - orig_group_numel + else: + padding = 0 + self.groups_padding.append(padding) + + # verify that data partition start locations are 4-byte aligned + for partitioned_data in data_parallel_partitions: + assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0) + + # A partition of the fp32 master weights that will be updated by this process. + # Note that the params in single_partition_of_fp32_groups is cloned and detached + # from the origin params of the model. + if not fp16_master_weights_and_gradients: + weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to( + self.device).clone().float().detach() + else: + weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to( + self.device).clone().half().detach() + + if self.cpu_offload: + weights_partition = get_accelerator().pin_memory(weights_partition) + temp_dtype = self.parallel_partitioned_bit16_groups[i][partition_id].dtype + temp_buffer_bit16 = torch.full(weights_partition.shape, + fill_value=0.0, + dtype=temp_dtype, + device=weights_partition.device) + if self.cpu_offload_pin_memory: + temp_pinned = get_accelerator().pin_memory(temp_buffer_bit16) + self.param_buffer_of_bit16_for_cpu_offload_groups.append(temp_pinned) + else: + self.param_buffer_of_bit16_for_cpu_offload_groups.append(temp_buffer_bit16) + + self.single_partition_of_fp32_groups.append(weights_partition) + + # Set local optimizer to have flat params of its own partition. + # After this, the local optimizer will only contain its own partition of params. + # In that case, the local optimizer only saves the states(momentum, variance, etc.) related to its partition's params(zero stage1). + self.single_partition_of_fp32_groups[ + i].requires_grad = True # keep this in case internal optimizer uses it + param_group['params'] = [self.single_partition_of_fp32_groups[i]] + + partition_size = len(self.bit16_groups_flat[i]) / dist.get_world_size(group=self.real_dp_process_group[i]) + params_in_partition, params_not_in_partition, first_offset = self.get_partition_info( + self.round_robin_bit16_groups[i], partition_size, partition_id) + + self.partition_size.append(partition_size) + self.params_in_partition.append(params_in_partition) + self.params_not_in_partition.append(params_not_in_partition) + self.first_offset.append(first_offset) + + self.reduce_bucket_size = int(reduce_bucket_size) + self.use_multi_rank_bucket_allreduce = use_multi_rank_bucket_allreduce + self.allgather_bucket_size = int(allgather_bucket_size) + + self.reduction_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream() + #self.copy_grad_stream = get_accelerator().Stream() + self.callback_queued = False + + self.param_dict = {} + + # map between param_id and bool to specify if a param is in this partition + self.is_param_in_current_partition = {} + + self.grads_in_ipg_bucket = [] + self.params_in_ipg_bucket = [] + self.elements_in_ipg_bucket = 0 + self.params_already_reduced = [] + self._release_ipg_buffers() + self.previous_reduced_grads = None + self.ipg_bucket_has_moe_params = False + + # simplified param id + self.param_id = {} + + #interesting code: unique ids being assigned to individual parameters + largest_param_numel = 0 + count = 0 + for i, params_group in enumerate(self.bit16_groups): + for param in params_group: + unique_id = id(param) + self.param_id[unique_id] = count + self.param_dict[count] = param + self.params_already_reduced.append(False) + if param.numel() > largest_param_numel: + largest_param_numel = param.numel() + count = count + 1 + + for param_group in self.params_in_partition: + for param in param_group: + self.is_param_in_current_partition[self.get_param_id(param)] = True + + for param_group in self.params_not_in_partition: + for param in param_group: + self.is_param_in_current_partition[self.get_param_id(param)] = False + + if self.cpu_offload: + self.accumulated_grads_in_cpu = {} + self.norm_for_param_grads = {} + self.local_overflow = False + self.grad_position = {} + self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel, + device=self.device, + dtype=self.dtype) + if self.cpu_offload_pin_memory: + self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory( + self.temp_grad_buffer_for_cpu_offload) + self.temp_grad_buffer_for_gpu_offload = torch.zeros(largest_param_numel, + device=get_accelerator().current_device_name(), + dtype=self.dtype) + for i, params_group in enumerate(self.bit16_groups): + self.get_grad_position(i, self.params_in_partition[i], self.first_offset[i], self.partition_size[i]) + + # mapping from parameter to partition that it belongs to + self.param_to_partition_ids = {} + + # stores if a partition has been reduced in this step + self.is_partition_reduced = {} + + # number of grads in partition that still need to be computed + self.remaining_grads_in_partition = {} + + # total number of grads in partition + self.total_grads_in_partition = {} + + # stores if a grad in a partition has been computed or not + self.is_grad_computed = {} + + # stores the offset at which a parameter gradient needs to be inserted in a partition + self.grad_partition_insertion_offset = {} + + # the offset in the gradient at which it must be inserted at the beginning of the partition + self.grad_start_offset = {} + + # will store the averaged gradients required by this partition + self.averaged_gradients = {} + + # For cpu_offload, will store the averaged gradients required by this partition + self.offload_gradient_dict = {} + + # store index of first parameter in each partition + self.first_param_index_in_partition = {} + + # initializes all data structures for implementing gradient partitioning + self.initialize_gradient_partitioning_data_structures() + + # resets the data structure value for the next backward propagation + self.reset_partition_gradient_structures() + + # creates backward hooks for the following special handling of gradients + # 1. upcasting for fp32 gradient accumulation + # 2. gradient partitioning + # 3. overlapping backward and reduction + self._grad_acc_hooks = [] + + if (self.partition_gradients or self.overlap_comm or self.use_grad_accum_attribute + or self.contiguous_gradients): + self.create_gradient_handling_hooks() + + self.ready_for_gradients = False + self.custom_loss_scaler = False + self.external_loss_scale = None + + # we may have a way of fusing dynamic scale. Do not support for now + self.loss_scaler = CreateLossScaler(dtype=self.dtype, + static_loss_scale=static_loss_scale, + dynamic_scaling=dynamic_loss_scale, + dynamic_loss_args=dynamic_loss_args) + self.dynamic_loss_scale = self.loss_scaler.dynamic + + if self.dtype != torch.float16: + # Only fp16 should use dynamic loss scaling + assert self.loss_scaler.cur_scale == 1.0 + assert not self.dynamic_loss_scale + + see_memory_usage("Before initializing optimizer states", force=True) + self.initialize_optimizer_states() + see_memory_usage("After initializing optimizer states", force=True) + + if dist.get_rank() == 0: + logger.info(f"optimizer state initialized") + + if dist.get_rank(group=self.dp_process_group) == 0: + see_memory_usage(f"After initializing ZeRO optimizer", force=True) + + self._link_all_hp_params() + self._hp_optimizer_states_linked = False + + self._enable_universal_checkpoint() + self._param_slice_mappings = self._create_param_mapping() + + def destroy(self): + for i, _ in enumerate(self.optimizer.param_groups): + for p in self.bit16_groups[i]: + if getattr(p, '_hp_mapping', None): + p._hp_mapping = None + for hook in self._grad_acc_hooks: + hook.remove() + self.print_rank_0("Removed grad acc hooks") + + def _enable_universal_checkpoint(self): + for lp_param_group in self.bit16_groups: + enable_universal_checkpoint(param_list=lp_param_group) + + def _create_param_mapping(self): + param_mapping = [] + for i, _ in enumerate(self.optimizer.param_groups): + param_mapping_per_group = OrderedDict() + for lp in self.bit16_groups[i]: + if lp._hp_mapping is not None: + lp_name = self.param_names[lp] + param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() + param_mapping.append(param_mapping_per_group) + + return param_mapping + + def _link_all_hp_params(self): + if self.cpu_offload: + self._get_offload_gradient_dict() + + for i, _ in enumerate(self.optimizer.param_groups): + # Link bit16 and fp32 params in partition + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + partition_size = self.bit16_groups_flat[i].numel() // dist.get_world_size( + group=self.real_dp_process_group[i]) + flat_hp_partition = self.single_partition_of_fp32_groups[i] + link_hp_params(lp_param_list=self.bit16_groups[i], + flat_hp_partition=flat_hp_partition, + gradient_dict=self.averaged_gradients, + offload_gradient_dict=self.offload_gradient_dict, + use_offload=self.cpu_offload, + param_group_index=i, + partition_start=partition_id * partition_size, + partition_size=partition_size, + dp_group=self.real_dp_process_group[i]) + + def _lazy_init_hp_params_optimizer_state(self): + if not self._hp_optimizer_states_linked: + for i, _ in enumerate(self.optimizer.param_groups): + lazy_init_hp_params_optimizer_state(self.bit16_groups[i], self.single_partition_of_fp32_groups[i], + self.optimizer.state) + self._hp_optimizer_states_linked = True + + def is_moe_group(self, group): + return 'moe' in group and group['moe'] + + def _configure_moe_settings(self): + # if we're using ZeRO stage 2, ensure contiguous gradients are used + if self.partition_gradients: + assert self.contiguous_gradients, "Contiguous Gradients in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" + # NOTE: To run ZeRO stage 1 with MoE, we need to set self.contiguous_gradients to True or ignore the assertion + if not self.partition_gradients and not self.contiguous_gradients: + logger.warning( + "ZeRO Stage 1 has not been thoroughly tested with MoE. This configuration is still experimental.") + assert self.reduce_scatter, "Reduce Scatter in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" + + assert any( + [self.is_moe_group(group) for group in self.optimizer.param_groups] + ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer" + self.is_moe_param_group = [] + for i, group in enumerate(self.optimizer.param_groups): + if self.is_moe_group(group): + assert all([is_moe_param(param) + for param in group['params']]), "All params in MoE group must be MoE params" + self.real_dp_process_group[i] = self.expert_dp_process_group[group['name']] + self.partition_count[i] = dist.get_world_size(group=self.expert_dp_process_group[group['name']]) + self.is_moe_param_group.append(True) + else: + self.is_moe_param_group.append(False) + + assert self.expert_dp_process_group is not None, "Expert data parallel group should be configured with MoE" + assert self.ep_process_group is not None, "Expert parallel group should be configured with MoE" + + def _update_model_bit16_weights(self, group_index): + updated_params = self.unflatten(self.bit16_groups_flat[group_index], self.round_robin_bit16_meta[group_index]) + for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params): + p.data = q.data + + # set model fp16 weight to slices of reordered flattened buffer + for param_index, param in enumerate(self.bit16_groups[group_index]): + new_index = self.round_robin_bit16_indices[group_index][param_index] + param.data = self.round_robin_bit16_groups[group_index][new_index].data + + def _round_robin_reorder(self, tensor_list, num_partitions): + + # disable round robin if need to debug something + # return tensor_list, list(range(len(tensor_list))) + + partition_tensors = {} + + for i, tensor in enumerate(tensor_list): + j = i % num_partitions + if not j in partition_tensors: + partition_tensors[j] = [] + partition_tensors[j].append((i, tensor)) + + reordered_tensors = [] + reordered_indices = {} + + for partition_index in partition_tensors.keys(): + for i, (original_index, tensor) in enumerate(partition_tensors[partition_index]): + reordered_indices[original_index] = len(reordered_tensors) + reordered_tensors.append(tensor) + + return reordered_tensors, reordered_indices + + def _release_ipg_buffers(self): + if self.contiguous_gradients: + self.ipg_buffer = None + self.grads_in_partition = None + self.grads_in_partition_offset = 0 + self.ready_for_gradients = False + + def initialize_optimizer_states(self): + + for i, group in enumerate(self.bit16_groups): + single_grad_partition = torch.zeros(int(self.partition_size[i]), + dtype=self.single_partition_of_fp32_groups[i].dtype, + device=self.device) + self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory( + single_grad_partition) if self.cpu_offload_pin_memory else single_grad_partition + + # Initialize the optimizer states with the flattened fp32 partition. + # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers + # which do lazy initialization of the state at the first call to step. + if isinstance(self.optimizer, torch.optim.Adagrad): + self.optimizer = torch.optim.Adagrad(self.single_partition_of_fp32_groups, **self.optimizer.defaults) + + if not self.cpu_offload: + for group in self.single_partition_of_fp32_groups: + group.grad = None #class init + + return + + ######################################################################### + #################### ZeRO Stage 1 - reduce gradients #################### + ######################################################################### + def reduce_gradients(self, pipeline_parallel=False): + world_size = dist.get_world_size(self.dp_process_group) + my_rank = dist.get_rank(self.dp_process_group) + + # with PP we must create ipg buffer, since backward is handled outside zero + if pipeline_parallel and self.contiguous_gradients: + self.ipg_buffer = [] + buf_0 = torch.empty(int(self.reduce_bucket_size), + dtype=self.dtype, + device=get_accelerator().current_device_name()) + self.ipg_buffer.append(buf_0) + self.ipg_index = 0 + + if not self.overlap_comm: + for i, group in enumerate(self.bit16_groups): + for param in group: + grad_reduc = self.get_gradient_for_reduction(param) + if grad_reduc is not None: + self.reduce_ready_partitions_and_remove_grads(param, i) + # reduce any pending grads in either hook/non-hook case + self.overlapping_partition_gradients_reduce_epilogue() + + ######################################################################### + #########################ZeRO Partition Gradients######################## + ######################################################################### + + def get_first_param_index(self, group_id, param_group, partition_id): + for index, param in enumerate(param_group): + param_id = self.get_param_id(param) + if group_id in self.param_to_partition_ids and param_id in self.param_to_partition_ids[group_id]: + if partition_id in self.param_to_partition_ids[group_id][param_id]: + return index + return None + + def initialize_gradient_partitioning_data_structures(self): + + for i, param_group in enumerate(self.round_robin_bit16_groups): + total_partitions = dist.get_world_size(group=self.real_dp_process_group[i]) + + self.param_to_partition_ids[i] = {} + self.is_partition_reduced[i] = {} + self.total_grads_in_partition[i] = {} + self.remaining_grads_in_partition[i] = {} + self.is_grad_computed[i] = {} + self.grad_partition_insertion_offset[i] = {} + self.grad_start_offset[i] = {} + self.first_param_index_in_partition[i] = {} + + for partition_id in range(total_partitions): + self.is_grad_computed[i][partition_id] = {} + self.grad_partition_insertion_offset[i][partition_id] = {} + self.grad_start_offset[i][partition_id] = {} + self.total_grads_in_partition[i][partition_id] = 0 + self.initialize_gradient_partition(i, param_group, partition_id) + self.is_partition_reduced[i][partition_id] = False + self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index( + i, param_group, partition_id) + + def independent_gradient_partition_epilogue(self): + self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) + self.reduce_ipg_grads() + self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) + + # if dist.get_rank() == 0: + # logger.info("Params already reduced %s", self.params_already_reduced) + for i in range(len(self.params_already_reduced)): + self.params_already_reduced[i] = False + + if self.overlap_comm: + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + # It is safe to clear previously reduced grads of other partitions + self._clear_previous_reduced_grads() + + if self.cpu_offload is False: + for i, _ in enumerate(self.bit16_groups): + + if not i in self.averaged_gradients or self.averaged_gradients[i] is None: + self.averaged_gradients[i] = self.get_flat_partition( + self.params_in_partition[i], + self.first_offset[i], + self.partition_size[i], + dtype=self.gradient_accumulation_dtype, + device=get_accelerator().current_device_name(), + return_tensor_list=True) + else: + avg_new = self.get_flat_partition(self.params_in_partition[i], + self.first_offset[i], + self.partition_size[i], + dtype=self.gradient_accumulation_dtype, + device=get_accelerator().current_device_name(), + return_tensor_list=True) + + for accumulated_grad, new_avg_grad in zip(self.averaged_gradients[i], avg_new): + accumulated_grad.add_(new_avg_grad) + + self._release_ipg_buffers() + + # No need to keep the gradients anymore. + # All gradients required by the step + # are in self.averaged_gradients + self.zero_grad(set_to_none=True) + see_memory_usage(f"End ipg_epilogue") + + # resets all partition to no reduced + # sets remaining grads to the total number of grads in each partition + # set is grad computed to false for all grads in partition + def reset_partition_gradient_structures(self): + for i, _ in enumerate(self.bit16_groups): + total_partitions = dist.get_world_size(group=self.real_dp_process_group[i]) + for partition_id in range(total_partitions): + self.is_partition_reduced[i][partition_id] = False + self.remaining_grads_in_partition[i][partition_id] = self.total_grads_in_partition[i][partition_id] + + for param_id in self.is_grad_computed[i][partition_id]: + self.is_grad_computed[i][partition_id][param_id] = False + + def initialize_gradient_partition(self, i, param_group, partition_id): + + def set_key_value_list(dictionary, key, value): + if key in dictionary: + dictionary[key].append(value) + else: + dictionary[key] = [value] + + def increment_value(dictionary, key): + if key in dictionary: + dictionary[key] += 1 + else: + dictionary[key] = 1 + + partition_size = self.partition_size[i] + + start_index = partition_size * partition_id + end_index = partition_size * (partition_id + 1) + + current_index = 0 + first_offset = 0 + + for param in param_group: + + param_size = param.numel() + param_id = self.get_param_id(param) + + if start_index <= current_index < end_index: + set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) + increment_value(self.total_grads_in_partition[i], partition_id) + + self.is_grad_computed[i][partition_id][param_id] = False + + self.grad_partition_insertion_offset[i][partition_id][param_id] = current_index - start_index + self.grad_start_offset[i][partition_id][param_id] = 0 + + elif current_index < start_index < (current_index + param_size): + assert (first_offset == 0 + ), "This can happen either zero or only once as this must be the first tensor in the partition" + first_offset = start_index - current_index + + set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) + increment_value(self.total_grads_in_partition[i], partition_id) + + self.is_grad_computed[i][partition_id][param_id] = False + + self.grad_partition_insertion_offset[i][partition_id][param_id] = 0 + self.grad_start_offset[i][partition_id][param_id] = first_offset + + current_index = current_index + param_size + + def overlapping_partition_gradients_reduce_epilogue(self): + self.independent_gradient_partition_epilogue() + + def _fill_param_grad_accum_attribute(self, param): + if param.grad is not None: + if param.grad_accum is None: + param.grad_accum = param.grad.to(self.gradient_accumulation_dtype) + else: + param.grad_accum.add_(param.grad.to(self.gradient_accumulation_dtype).view(param.grad_accum.shape)) + param.grad = None + + def fill_grad_accum_attribute(self): + for group in self.bit16_groups: + for param in group: + self._fill_param_grad_accum_attribute(param) + + def get_gradient_for_reduction(self, param): + if self.use_grad_accum_attribute: + return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None + else: + return param.grad + + def get_param_gradient_attribute(self, param): + return param.grad_accum if self.use_grad_accum_attribute else param.grad + + # Clear the tensor the reduction gradient attribute is pointing to + def clear_grad_attribute(self, param): + if self.use_grad_accum_attribute: + param.grad_accum = None + else: + param.grad = None + + def create_gradient_handling_hooks(self): + for i, param_group in enumerate(self.bit16_groups): + for param in param_group: + if param.requires_grad: + + def wrapper(param, i): + + def grad_handling_hook(*notneeded): + self.process_gradients(param, i) + + self._grad_acc_hooks.append(register_grad_hook(param, grad_handling_hook)) + + wrapper(param, i) + + def get_param_id(self, param): + unique_id = id(param) + return self.param_id[unique_id] + + def report_ipg_memory_usage(self, tag, param_elems): + elem_count = self.elements_in_ipg_bucket + param_elems + percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size + see_memory_usage( + f"{tag}: elems in_bucket {self.elements_in_ipg_bucket} param {param_elems} max_percent {percent_of_bucket_size}" + ) + + # create a flat tensor aligned at the alignment boundary + def flatten_dense_tensors_aligned(self, tensor_list, alignment, use_cpu_data=False): + tensor_list = [param.cpu_data for param in tensor_list] if use_cpu_data else tensor_list + return self.flatten(align_dense_tensors(tensor_list, alignment)) + + ############### Independent Partition Gradient ######################## + def reduce_independent_p_g_buckets_and_remove_grads(self, param, i): + + grad_reduc = self.get_gradient_for_reduction(param) + if self.elements_in_ipg_bucket + param.numel() > self.reduce_bucket_size: + self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.numel()) + self.reduce_ipg_grads() + if self.contiguous_gradients and self.overlap_comm: + # Swap ipg_index between 0 and 1 + self.ipg_index = 1 - self.ipg_index + self.report_ipg_memory_usage("In ipg_remove_grads after reduce_ipg_grads", param.numel()) + + param_id = self.get_param_id(param) + assert self.params_already_reduced[param_id] == False, \ + f"The parameter {param_id} has already been reduced. \ + Gradient computed twice for this partition. \ + Multiple gradient reduction is currently not supported" + + if self.contiguous_gradients: + if param.numel() > self.reduce_bucket_size: + self.extra_large_param_to_reduce = param + else: + # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening + new_grad_tensor = self.ipg_buffer[self.ipg_index].narrow(0, self.elements_in_ipg_bucket, param.numel()) + new_grad_tensor.copy_(grad_reduc.view(-1)) + grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) + + self.elements_in_ipg_bucket += param.numel() + + assert grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient" + + # deal with a use-case of transient grads that will be generated in a loop for the same computation involving some model params - e.g. when performing a tiled memory calculation that shards the normal single sub-module call into a loop over a shards. + if getattr(param, "ds_grad_is_ready", True): + self.grads_in_ipg_bucket.append(grad_reduc) + self.params_in_ipg_bucket.append((i, param.param_idx_in_group, param_id)) + + #make sure the average tensor function knows how to average the gradients + if is_moe_param(param): + self.ipg_bucket_has_moe_params = True + + self.report_ipg_memory_usage("End ipg_remove_grads", 0) + + def print_rank_0(self, message): + if dist.get_rank() == 0: + logger.info(message) + + def gradient_reduction_w_predivide(self, tensor): + if tensor.size().numel() == 0: + return tensor + + dp_world_size = dist.get_world_size(group=self.dp_process_group) + + tensor_to_allreduce = tensor + + if self.communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(self.communication_data_type) + + if self.postscale_gradients: + if self.gradient_predivide_factor != 1.0: + tensor_to_allreduce.mul_(1. / self.gradient_predivide_factor) + + dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) + + if self.gradient_predivide_factor != dp_world_size: + tensor_to_allreduce.mul_(self.gradient_predivide_factor / + (dp_world_size / float(self.sequence_parallel_size))) + else: + tensor_to_allreduce.div_(dp_world_size / float(self.sequence_parallel_size)) + dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) + + if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + tensor.copy_(tensor_to_allreduce) + + return tensor + + def allreduce_and_copy_with_multiple_ranks(self, + small_bucket, + log=None, + divide=True, + process_group=None, + bucket_ranks=None): + process_group = self.dp_process_group if process_group is None else process_group + allreduced = self.allreduce_bucket(small_bucket, log=log, divide=divide, process_group=process_group) + for buf, synced, bucket_rank in zip(small_bucket, self.unflatten(allreduced, small_bucket), bucket_ranks): + if dist.get_rank(group=process_group) == bucket_rank: + buf.copy_(synced) + + def allreduce_and_scatter(self, bucket, numel_per_bucket=500000000, log=None, divide=True, process_group=None): + small_bucket = [] + small_bucket_ranks = [] + numel = 0 + allreduce_sizes = [] + + for i, bucket_elem in enumerate(bucket): + rank, tensor = bucket_elem + small_bucket.append(tensor) + small_bucket_ranks.append(rank) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy_with_multiple_ranks(small_bucket, + log=None, + divide=divide, + process_group=process_group, + bucket_ranks=small_bucket_ranks) + small_bucket = [] + small_bucket_ranks = [] + numel = 0 + + if len(small_bucket) > 0: + self.allreduce_and_copy_with_multiple_ranks(small_bucket, + log=None, + divide=divide, + process_group=process_group, + bucket_ranks=small_bucket_ranks) + + def average_tensor(self, tensor): + if self.overlap_comm: + stream = self.reduction_stream + if not get_accelerator().resolves_data_dependency(): + stream.wait_stream(get_accelerator().current_stream()) + get_accelerator().current_stream().wait_stream(stream) + else: + stream = get_accelerator().current_stream() + + with get_accelerator().stream(stream): + if not self.reduce_scatter: + self.gradient_reduction_w_predivide(tensor) + return + + # Accumulate destination ranks and bucket offsets for each gradient slice. + # Note: potential future optimization, record access pattern of parameters + # in backward pass and partition gradients w.r.t. access pattern so that our + # bucket is guaranteed to be contiguous w.r.t. ranks + rank_and_offsets = [] + real_dp_process_group = [] + curr_size = 0 + prev_id, prev_process_group = -1, None + + process_group = self.dp_process_group + # count = 0 + for i, param_idx_in_group, param_id in self.params_in_ipg_bucket: + param = self.bit16_groups[i][param_idx_in_group] + + process_group = self.dp_process_group + + if self.ipg_bucket_has_moe_params: + process_group = self.expert_dp_process_group[param.group_name] if is_moe_param( + param) else self.dp_process_group + + partition_ids = self.param_to_partition_ids[i][param_id] + assert all([p_id < dist.get_world_size(group=process_group) for p_id in partition_ids + ]), f"world size {dist.get_world_size(group=process_group)} and p_ids: {partition_ids}" + partition_size = self.partition_size[i] + # Get all partition ids + their offsets + partition_ids_w_offsets = [] + for partition_id in partition_ids: + offset = self.grad_start_offset[i][partition_id][param_id] + partition_ids_w_offsets.append((partition_id, offset)) + partition_ids_w_offsets.sort(key=lambda t: t[1]) + + # Calculate rank and offsets for grad slices + for idx in range(len(partition_ids_w_offsets)): + partition_id, offset = partition_ids_w_offsets[idx] + + # if dist.get_rank() == 0 and count < 100: + # print(f"Rank {dist.get_rank()} rank offset id {idx} calculated dp size {dist.get_world_size(group=process_group)} real dp size {dist.get_world_size(self.real_dp_process_group[i])} and dst: {partition_id}") + # count += 1 + + # Calculate numel for grad slice depending on partition location + if idx == len(partition_ids_w_offsets) - 1: + # Last partition_id uses its own offset + numel = param.numel() - offset + else: + # Set numel to next partition's offset + numel = partition_ids_w_offsets[idx + 1][1] - offset + + # Merge bucket ranges if they belong to the same rank + if partition_id == prev_id and process_group == prev_process_group: + prev_pid, prev_size, prev_numel = rank_and_offsets[-1] + rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel) + else: + rank_and_offsets.append((partition_id, curr_size, numel)) + real_dp_process_group.append(process_group) + curr_size += numel + prev_id, prev_process_group = partition_id, process_group + + tensor.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) + + buckets = {} + for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets): + grad_slice = tensor.narrow(0, int(bucket_offset), int(numel)) + bucket_key = real_dp_process_group[i] if self.use_multi_rank_bucket_allreduce else ( + dst, real_dp_process_group[i]) + if bucket_key not in buckets: + buckets[bucket_key] = [] + if self.use_multi_rank_bucket_allreduce: + buckets[bucket_key].append((dst, grad_slice)) + else: + buckets[bucket_key].append(grad_slice) + + for bucket_key in buckets: + if self.use_multi_rank_bucket_allreduce: + self.allreduce_and_scatter(buckets[bucket_key], + numel_per_bucket=self.reduce_bucket_size, + divide=False, + process_group=bucket_key) + else: + dst, process_group = bucket_key + self.allreduce_no_retain(buckets[bucket_key], + numel_per_bucket=self.reduce_bucket_size, + rank=dst, + divide=False, + process_group=process_group) + + ############################################################################## + ############################# CPU Offload Methods############################# + ############################################################################## + def get_grad_position(self, group_id, tensor_list, first_offset, partition_size): + current_offset = 0 + + for i, tensor in enumerate(tensor_list): + param_id = self.get_param_id(tensor) + param_start_offset = 0 + + num_elements = tensor.numel() + + # we need to offset to get to the right element + if i == 0 and first_offset > 0: + tensor_offset = first_offset + num_elements = num_elements - tensor_offset + param_start_offset = first_offset + + # we dont need all elements of the tensor + if num_elements > (partition_size - current_offset): + num_elements = partition_size - current_offset + + self.grad_position[param_id] = [ + int(group_id), int(param_start_offset), + int(current_offset), int(num_elements) + ] + current_offset += num_elements + + def update_overflow_tracker_for_param_grad(self, param): + grad_accum = self.get_param_gradient_attribute(param) + if grad_accum is not None and self._has_inf_or_nan(grad_accum.data): + self.local_overflow = True + + def _get_offload_gradient_dict(self): + for param_group_index, _ in enumerate(self.optimizer.param_groups): + self.offload_gradient_dict[param_group_index] = [] + for lp_param in self.params_in_partition[param_group_index]: + param_id = self.get_param_id(lp_param) + [_, _, dest_offset, num_elements] = self.grad_position[param_id] + dest_tensor = self.single_partition_of_fp32_groups[param_group_index].grad.view(-1).narrow( + 0, dest_offset, num_elements) + self.offload_gradient_dict[param_group_index].append(dest_tensor) + + def async_accumulate_grad_in_cpu_via_gpu(self, param): + param_id = self.get_param_id(param) + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + # copy to a preexisiting buffer to avoid memory allocation penalty + dest_buffer = self.temp_grad_buffer_for_gpu_offload.view(-1).narrow(0, 0, param.numel()) + + #buffer for storing gradients for this parameter in CPU + def buffer_to_accumulate_to_in_cpu(): + if not self.fp16_master_weights_and_gradients: + buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device) + return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer + else: + return self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) + + #accumulate gradients into param.grad_accum or parts of it that belongs to this partition + def accumulate_gradients(): + grad_accum = self.get_param_gradient_attribute(param) + if not self.fp16_master_weights_and_gradients: + dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True) + grad_accum.data.view(-1).add_(dest_buffer) + else: + dest_buffer.narrow(0, source_offset, + num_elements).copy_(self.accumulated_grads_in_cpu[param_id].view(-1), + non_blocking=True) + grad_accum.data.view(-1).narrow(0, source_offset, + num_elements).add_(dest_buffer.narrow(0, source_offset, num_elements)) + + #move accumulated gradients back to CPU + def copy_gradients_to_cpu(): + grad_accum = self.get_param_gradient_attribute(param) + if not self.fp16_master_weights_and_gradients: + self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1), non_blocking=True) + else: + self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1).narrow( + 0, source_offset, num_elements), + non_blocking=True) + + if param_id not in self.accumulated_grads_in_cpu: + self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu() + + if self.micro_step_id > 0: + accumulate_gradients() + else: + copy_gradients_to_cpu() + + def set_norm_for_param_grad(self, param): + param_id = self.get_param_id(param) + grad_accum = self.get_param_gradient_attribute(param) + accumulated_grad = self.accumulated_grads_in_cpu[ + param_id] if self.gradient_accumulation_steps > 1 else grad_accum + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + start = source_offset + accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) + + self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) + + def set_norm_for_param_grad_in_gpu(self, param): + param_id = self.get_param_id(param) + grad_accum = self.get_param_gradient_attribute(param) + if grad_accum is None: + accumulated_grad = param.grad + else: + accumulated_grad = grad_accum + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + start = source_offset + accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) + + self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) + + def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param): + param_id = self.get_param_id(param) + + [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] + + dest_tensor = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) + + grad_accum = self.get_param_gradient_attribute(param) + if grad_accum is None: + src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) + else: + src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) + if not self.fp16_master_weights_and_gradients: + src_tensor = src_tensor.float() + + dest_tensor.copy_(src_tensor, non_blocking=True) + param.grad = None #offload only + + def complete_grad_norm_calculation_for_cpu_offload(self, params): + total_norm = 0.0 + norm_type = 2.0 + for p in params: + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + param_id = self.get_param_id(p) + # as some model have trainable parameters but skipped in training, + # their backward hooks in self.create_gradient_handling_hooks() will not run, + # so they have no norm_for_param_grads + if param_id in self.norm_for_param_grads: + param_norm = self.norm_for_param_grads[param_id] + total_norm += param_norm.item()**2 + else: + # As unused parameters in modules may not be expected sometimes, + # add an explicit error msg when it occurred and an option to + # avoid the error + assert self.ignore_unused_parameters, """ + This assert indicates that your module has parameters that + were not used in producing loss. + You can avoid this assert by + (1) enable ignore_unused_parameters option in zero_optimization config; + (2) making sure all trainable parameters and `forward` function + outputs participate in calculating loss. + """ + + # Sum across all model parallel GPUs. + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) + + total_norm = total_norm_cuda[0].item()**(1. / norm_type) + + if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: + total_norm = -1 + + return total_norm + + ############################################################################################ + def copy_grads_in_partition(self, param): + if self.cpu_offload: + + if self.gradient_accumulation_steps > 1: + self.async_accumulate_grad_in_cpu_via_gpu(param) + + if self.is_gradient_accumulation_boundary: + self.set_norm_for_param_grad_in_gpu(param) + + self.update_overflow_tracker_for_param_grad(param) + + self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param) + + return + #print(f"ID {self.get_param_id(param)} grad norm {param.grad.norm()}") + if self.grads_in_partition is None: + self.grads_in_partition_offset = 0 + total_size = 0 + for group in self.params_in_partition: + for param_in_partition in group: + total_size += param_in_partition.numel() + + see_memory_usage(f"before copying {total_size} gradients into partition") + self.grads_in_partition = torch.empty(int(total_size), + dtype=self.dtype, + device=get_accelerator().current_device_name()) + see_memory_usage(f"after copying {total_size} gradients into partition") + + grad_reduc = self.get_gradient_for_reduction(param) + # The allreduce buffer will be rewritten. Copy the gradients in partition to a new buffer + new_grad_tensor = self.grads_in_partition.view(-1).narrow(0, self.grads_in_partition_offset, param.numel()) + new_grad_tensor.copy_(grad_reduc.view(-1)) + grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) + #print(f"Grad norm after copy to contiguous_buffer {param.grad.data.norm()}") + self.grads_in_partition_offset += param.numel() + + def reduce_ipg_grads(self): + if self.contiguous_gradients: + if self.extra_large_param_to_reduce is not None: + assert len(self.params_in_ipg_bucket) == 1, "more than 1 param in ipg bucket, this shouldn't happen" + _, _, param_id = self.params_in_ipg_bucket[0] + assert self.get_param_id(self.extra_large_param_to_reduce + ) == param_id, "param in ipg bucket does not match extra-large param" + extra_large_grad_reduc = self.get_gradient_for_reduction(self.extra_large_param_to_reduce) + self.average_tensor(extra_large_grad_reduc.view(-1)) + self.extra_large_param_to_reduce = None + else: + self.average_tensor(self.ipg_buffer[self.ipg_index].narrow(0, 0, self.elements_in_ipg_bucket)) + else: + self.buffered_reduce_fallback(None, + self.grads_in_ipg_bucket, + elements_per_buffer=self.elements_in_ipg_bucket) + + if self.overlap_comm: + stream = self.reduction_stream + elif self.cpu_offload: + # TODO: copy_grad_stream is disabled because of race with reduce. This hurts perf and should be fixed. + # get_accelerator().synchronize() + # stream = self.copy_grad_stream + stream = get_accelerator().current_stream() + else: + stream = get_accelerator().current_stream() + + with get_accelerator().stream(stream): + for group_idx, param_idx_in_group, param_id in self.params_in_ipg_bucket: + param = self.bit16_groups[group_idx][param_idx_in_group] + + assert self.params_already_reduced[param_id] == False, \ + f"The parameter {param_id} has already been reduced. \ + Gradient computed twice for this partition. \ + Multiple gradient reduction is currently not supported" + + self.params_already_reduced[param_id] = True + if self.partition_gradients: + if not self.is_param_in_current_partition[param_id]: + if self.overlap_comm and self.contiguous_gradients is False: + # Clear grads of other partitions during the next reduction + # to avoid clearing them before the reduction is complete. + if self.previous_reduced_grads is None: + self.previous_reduced_grads = [] + self.previous_reduced_grads.append(param) + else: + self.clear_grad_attribute(param) + elif self.contiguous_gradients: + self.copy_grads_in_partition(param) + else: # zero stage 1 - partition only optimizer state + if self.contiguous_gradients and self.is_param_in_current_partition[param_id]: + self.copy_grads_in_partition(param) + + self.grads_in_ipg_bucket = [] + self.params_in_ipg_bucket = [] + self.ipg_bucket_has_moe_params = False + self.elements_in_ipg_bucket = 0 + ##################################################################### + + def process_gradients(self, param, i): + self.backward_prologue() + if self.use_grad_accum_attribute: + self._fill_param_grad_accum_attribute(param) + if self.partition_gradients or self.overlap_comm: + self.reduce_ready_partitions_and_remove_grads(param, i) + + def reduce_ready_partitions_and_remove_grads(self, param, i): + if self.partition_gradients or self.is_gradient_accumulation_boundary: + self.reduce_independent_p_g_buckets_and_remove_grads(param, i) + + def zero_reduced_gradients(self, partition_id, i): + + def are_all_related_partitions_reduced(params_id): + for partition_id in self.param_to_partition_ids[i][params_id]: + if not self.is_partition_reduced[i][partition_id]: + return False + return True + + for params_id in self.is_grad_computed[i][partition_id]: + if are_all_related_partitions_reduced(params_id): + self.param_dict[params_id].grad = None # dead code + + def flatten_and_print(self, message, tensors, start=0, n=5): + flatten_tensor = self.flatten(tensors) + + def print_func(): + logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n)) + + self.sequential_execution(print_func, message) + + def get_grads_to_reduce(self, i, partition_id): + + def get_reducible_portion(key): + grad = self.param_dict[key].grad + total_elements = grad.numel() + start = self.grad_start_offset[i][partition_id][key] + num_elements = min(total_elements - start, + self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key]) + if not pg_correctness_test: + if num_elements == total_elements: + return grad + else: + return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements)) + else: + if num_elements == total_elements: + return grad.clone() + else: + return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements)) + + grads_to_reduce = [] + for key in self.is_grad_computed[i][partition_id]: + grad = get_reducible_portion(key) + grads_to_reduce.append(grad) + return grads_to_reduce + + def sequential_execution(self, function, message, group=None): + if group is None: + group = self.dp_process_group + if dist.get_rank(group=group) == 0: + logger.info(message) + for id in range(dist.get_world_size(group=group)): + if id == dist.get_rank(group=group): + function() + dist.barrier(group=group) + + def set_none_gradients_to_zero(self, i, partition_id): + for param_id in self.is_grad_computed[i][partition_id]: + param = self.param_dict[param_id] + if param.grad is None: + param.grad = torch.zeros_like(param) + + ######################Reduction Related Methods############################## + def allreduce_bucket(self, bucket, rank=None, log=None, divide=True, process_group=None): + tensor = self.flatten(bucket) + + process_group = self.dp_process_group if process_group is None else process_group + + tensor_to_allreduce = tensor + + if pg_correctness_test or self.sequence_parallel_size > 1: + communication_data_type = torch.float32 + else: + communication_data_type = self.communication_data_type + + if communication_data_type != tensor.dtype: + tensor_to_allreduce = tensor.to(communication_data_type) + + if divide: + tensor_to_allreduce.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size)) + + if rank is None: + # "All Reducing" + dist.all_reduce(tensor_to_allreduce, group=process_group) + else: + global_rank = dist.get_global_rank(process_group, rank) + dist.reduce(tensor_to_allreduce, global_rank, group=process_group) + + if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: + if rank is None or rank == dist.get_rank(group=process_group): + tensor.copy_(tensor_to_allreduce) + + return tensor + + def _clear_previous_reduced_grads(self): + if self.previous_reduced_grads is not None: + for param in self.previous_reduced_grads: + self.clear_grad_attribute(param) + self.previous_reduced_grads = None + + # if rank is specified do a reduction instead of an allreduce + def allreduce_and_copy(self, small_bucket, rank=None, log=None, divide=True, process_group=None): + process_group = self.dp_process_group if process_group is None else process_group + if self.overlap_comm: + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + # It is safe to clear the previously reduced grads of other partitions + self._clear_previous_reduced_grads() + stream = self.reduction_stream + else: + stream = get_accelerator().current_stream() + + with get_accelerator().stream(stream): + allreduced = self.allreduce_bucket( + small_bucket, + rank=rank, + log=log, + divide=divide, + process_group=process_group, + ) + if rank is None or rank == dist.get_rank(group=self.dp_process_group): + for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): + buf.copy_(synced) + + def allreduce_no_retain( + self, + bucket, + numel_per_bucket=500000000, + rank=None, + log=None, + divide=True, + process_group=None, + ): + small_bucket = [] + numel = 0 + for tensor in bucket: + small_bucket.append(tensor) + numel = numel + tensor.numel() + if numel > numel_per_bucket: + self.allreduce_and_copy(small_bucket, rank=rank, log=None, divide=divide, process_group=process_group) + small_bucket = [] + numel = 0 + + if len(small_bucket) > 0: + self.allreduce_and_copy(small_bucket, rank=rank, log=log, divide=divide, process_group=process_group) + + # allows using reduction of gradients instead of using all_reduce + + def buffered_reduce_fallback(self, rank, grads, elements_per_buffer=500000000, log=None): + split_buckets = split_half_float_double(grads) + + for i, bucket in enumerate(split_buckets): + self.allreduce_no_retain(bucket, numel_per_bucket=elements_per_buffer, rank=rank, log=log) + + ############################################################################# + ############################################################################# + ############################################################################# + + # views the tensor as multiple partitions and returns + # those partitions + def get_data_parallel_partitions(self, tensor, group_id): + partitions = [] + + dp = dist.get_world_size(group=self.real_dp_process_group[group_id]) + # dp_id = dist.get_rank(group=self.real_dp_process_group[group_id]) + + total_num_elements = tensor.numel() + + base_size = total_num_elements // dp + remaining = total_num_elements % dp + + start = 0 + for id in range(dp): + partition_size = base_size + if id < remaining: + partition_size = partition_size + 1 + partitions.append(tensor.narrow(0, start, partition_size)) + start = start + partition_size + return partitions + + def get_partition_info(self, tensor_list, partition_size, partition_id): + params_in_partition = [] + params_not_in_partition = [] + + start_index = partition_size * partition_id + end_index = partition_size * (partition_id + 1) + + current_index = 0 + first_offset = 0 + + for tensor in tensor_list: + + tensor_size = tensor.numel() + + if start_index <= current_index < end_index: + params_in_partition.append(tensor) + + elif current_index < start_index < (current_index + tensor_size): + params_in_partition.append(tensor) + + assert (first_offset == 0 + ), "This can happen either zero or only once as this must be the first tensor in the partition" + first_offset = start_index - current_index + + else: + params_not_in_partition.append(tensor) + + current_index = current_index + tensor_size + + return params_in_partition, params_not_in_partition, first_offset + + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + # FP32 grad should never exist. + # For speed, set model fp16 grad to None by default + # zero all pointers to grad tensors + for group in self.bit16_groups: + for p in group: + if set_to_none: + p.grad = None # epilogue and in step + p.grad_accum = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def _model_parallel_all_reduce(self, tensor, op): + """ Perform all reduce within model parallel group, if any. + """ + if self.model_parallel_group is None or self.model_parallel_world_size == 1: + pass + else: + dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group) + + def get_grad_norm_direct(self, gradients, params, norm_type=2): + """Clips gradient norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float or int): max norm of the gradients + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + """ + norm_type = float(norm_type) + all_norms = [] + if norm_type == inf: + for g in gradients: + all_norms.append(g.data.abs().max().float()) + total_norm = torch.stack(all_norms).max() + dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + # Take max across all GPUs. + self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.MAX) + else: + # if dist.get_rank() == 0: + # logger.info(f"Total Norm beginning {total_norm}") + for g, p in zip(gradients, params): + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): + all_norms.append( + torch.linalg.vector_norm(g.data.double().detach(), + ord=norm_type).to(get_accelerator().current_device_name())) + if len(all_norms) > 0: + total_norm = torch.stack(all_norms).square().sum().float() + else: + total_norm = torch.tensor(0.0, dtype=torch.float32).to(self.device) + # Sum across all model parallel Device. + dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=self.dp_process_group) + + self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.SUM) + + total_norm = total_norm.pow(1. / norm_type) + + mask_nan_or_inf_with_val_inplace(total_norm, device=self.device) + + return total_norm + + # creates a flat fused tensor from the tensor list starting at the first_offset + # in the first tensor of the list. If there are not enough elements in the tensor + # list then the flat tensor will be padded with zeros + def get_flat_partition(self, tensor_list, first_offset, partition_size, dtype, device, return_tensor_list=False): + flat_tensor_list = [] + current_size = 0 + + for i, tensor in enumerate(tensor_list): + grad_accum = self.get_param_gradient_attribute(tensor) + if grad_accum is None: + grad_accum = torch.zeros_like(tensor, dtype=dtype) + + tensor = grad_accum + num_elements = tensor.numel() + tensor_offset = 0 + + # we need to offset to get to the right element + if i == 0 and first_offset > 0: + tensor_offset = first_offset + num_elements = num_elements - tensor_offset + + # we dont need all elements of the tensor + if num_elements > (partition_size - current_size): + num_elements = partition_size - current_size + + # we need a narrow view of the tensor based on the tensor offset and number of elements that + # we need from this tensor + if tensor_offset > 0 or num_elements < tensor.numel(): + flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements))) + else: + flat_tensor_list.append(tensor) + + current_size = current_size + num_elements + + # this means its the last partition and does not align with the dp boundary. We need to pad before flattening + if current_size < partition_size: + flat_tensor_list.append(torch.zeros(int(partition_size - current_size), dtype=dtype, device=device)) + + if return_tensor_list: + return flat_tensor_list + + return self.flatten(flat_tensor_list) + + def free_grad_in_param_list(self, param_list): + for p in param_list: + p.grad = None # in step + p.grad_accum = None + + def reset_cpu_buffers(self): + self.norm_for_param_grads = {} + self.local_overflow = False + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + def scaled_global_norm(self, norm_type=2): + assert norm_type == 2, "only L2 norm supported" + norm_groups = [] + for i, group in enumerate(self.bit16_groups): + if self.cpu_offload: + # complete complete_grad_norm_calculation_for_cpu_offload return python float, moving back to + # torch.tensor as else statement returns tensor as well + norm = torch.tensor(self.complete_grad_norm_calculation_for_cpu_offload(self.params_in_partition[i]), + device=self.device) + norm_groups.append(norm) + else: + norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.params_in_partition[i])) + + if self.has_moe_layers: + self._average_expert_grad_norms(norm_groups) + + # calculating L2 norm + return torch.linalg.vector_norm(torch.stack(norm_groups), ord=norm_type) + + def get_bit16_param_group(self, group_no): + bit16_partitions = self.parallel_partitioned_bit16_groups[group_no] + partition_id = dist.get_rank(group=self.real_dp_process_group[group_no]) + return [bit16_partitions[dist.get_rank(group=self.real_dp_process_group[group_no])]] + + def _optimizer_step(self, group_no): + original_param_groups = self.optimizer.param_groups + self.optimizer.param_groups = [original_param_groups[group_no]] + # Disabling this as the C++ side copy & synchronize is not working correctly + #from deepspeed.ops.adam import DeepSpeedCPUAdam + #if type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half: + # self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no)]) + #else: + # self.optimizer.step() + self.optimizer.step() + self.optimizer.param_groups = original_param_groups + + # We need to link optimizer state after the first step() call + self._lazy_init_hp_params_optimizer_state() + + def step(self, closure=None): + """ + Not supporting closure. + """ + self.micro_step_id = INITIAL_MICRO_STEP_ID + + see_memory_usage(f"In step before checking overflow") + + # First compute norm for all group so we know if there is overflow + if self.dtype == torch.float16: + self.check_overflow() + + prev_scale = self.loss_scale + self._update_scale(self.overflow) + if self.overflow: + see_memory_usage('After overflow before clearing gradients') + self.zero_grad(set_to_none=True) + if self.cpu_offload: + self.reset_cpu_buffers() + else: + self.averaged_gradients = {} + + see_memory_usage('After overflow after clearing gradients') + + for timer in OPTIMIZER_TIMERS: + self.timers(timer).start() + self.timers(timer).stop() + return + + # Step 1:- Calculate gradient norm using bit-16 grads + see_memory_usage('Before norm calculation') + scaled_global_grad_norm = self.scaled_global_norm() + self._global_grad_norm = scaled_global_grad_norm / prev_scale + see_memory_usage('After norm before optimizer') + + # Step 2:- run optimizer and upscaling simultaneously + for i, group in enumerate(self.bit16_groups): + self.timers(OPTIMIZER_GRADIENTS_TIMER).start() + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + if self.cpu_offload: + single_grad_partition = self.single_partition_of_fp32_groups[i].grad + self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) + + self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() + self.timers(OPTIMIZER_STEP_TIMER).start() + self._optimizer_step(i) + + # Disabled, this is not currently working + #from deepspeed.ops.adam import DeepSpeedCPUAdam + #if not (type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half): + # bit16_partitions = self.parallel_partitioned_bit16_groups[i] + # fp32_partition = self.single_partition_of_fp32_groups[i] + # bit16_partitions[partition_id].data.copy_(fp32_partition.data) + bit16_partitions = self.parallel_partitioned_bit16_groups[i] + fp32_partition = self.single_partition_of_fp32_groups[i] + bit16_partition_buffer = self.param_buffer_of_bit16_for_cpu_offload_groups[i] + bit16_partition_buffer.data.copy_(fp32_partition.data) + bit16_partitions[partition_id].data.copy_(bit16_partition_buffer.data, non_blocking=True) + + self.timers(OPTIMIZER_STEP_TIMER).stop() + else: + # free gradients for all the parameters that are not updated by this process(ZeRO stage2) + self.free_grad_in_param_list(self.params_not_in_partition[i]) + + # create a flat gradients for parameters updated by this process + # If we are last partition, ensure we have same size grads and partition size, if not pad with zero tensors + if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1: + single_grad_partition = self.flatten_dense_tensors_aligned( + self.averaged_gradients[i], + int(self.partition_size[i])).to(self.single_partition_of_fp32_groups[i].dtype) + else: + single_grad_partition = self.flatten(self.averaged_gradients[i]).to( + self.single_partition_of_fp32_groups[i].dtype) + assert single_grad_partition.numel() == self.partition_size[i], \ + "averaged gradients have different number of elements that partition size {} {} {} {}".format( + single_grad_partition.numel(), self.partition_size[i], i, partition_id) + + self.single_partition_of_fp32_groups[i].grad = single_grad_partition + # release all the gradient since we have already created a necessary copy in dp_grad_partition(ZeRO stage2) + self.free_grad_in_param_list(self.params_in_partition[i]) + + self.averaged_gradients[i] = None + + self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) + + self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() + + # Step 3:- run the optimizer if no offloading + self.timers(OPTIMIZER_STEP_TIMER).start() + self._optimizer_step(i) + # Step 4:- get rid of the fp32 gradients. Not needed anymore + self.single_partition_of_fp32_groups[i].grad = None + del single_grad_partition + bit16_partitions = self.parallel_partitioned_bit16_groups[i] + fp32_partition = self.single_partition_of_fp32_groups[i] + bit16_partitions[partition_id].data.copy_(fp32_partition.data) + self.timers(OPTIMIZER_STEP_TIMER).stop() + + see_memory_usage('After optimizer before all-gather') + if self.cpu_offload: + self.reset_cpu_buffers() + + self.timers(OPTIMIZER_ALLGATHER_TIMER).start() + # Gather the updated weights from everyone. + # Then all partitions of the model parameters are updated and ready for next round forward. + all_gather_dp_groups(groups_flat=self.bit16_groups_flat, + partitioned_param_groups=self.parallel_partitioned_bit16_groups, + dp_process_group=self.real_dp_process_group, + start_alignment_factor=self.nccl_start_alignment_factor, + allgather_bucket_size=self.allgather_bucket_size) + self.timers(OPTIMIZER_ALLGATHER_TIMER).stop() + + # TODO: we probably don't need this? just to be safe + for i in range(len(self.bit16_groups)): + self._update_model_bit16_weights(i) + + self.timers.log(OPTIMIZER_TIMERS) + see_memory_usage('After zero_optimizer step') + + return + + @torch.no_grad() + def update_lp_params(self): + for i, (bit16_partitions, fp32_partition) in enumerate( + zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + bit16_partitions[partition_id].data.copy_(fp32_partition.data) + + all_gather_dp_groups(groups_flat=self.bit16_groups_flat, + partitioned_param_groups=self.parallel_partitioned_bit16_groups, + dp_process_group=self.real_dp_process_group, + start_alignment_factor=self.nccl_start_alignment_factor, + allgather_bucket_size=self.allgather_bucket_size) + + def _average_expert_grad_norms(self, norm_groups): + for i, norm in enumerate(norm_groups): + if self.is_moe_param_group[i]: + scaled_norm_tensor = norm * 1.0 / dist.get_world_size(group=self.real_dp_process_group[i]) + if self.device == 'cpu': + scaled_norm_tensor = scaled_norm_tensor.to(get_accelerator().current_device_name()) + dist.all_reduce(scaled_norm_tensor, group=self.real_dp_process_group[i]) + norm_groups[i] = scaled_norm_tensor.to(self.device) + + def unscale_and_clip_grads(self, grad_groups_flat, total_norm): + # compute combined scale factor for this group + combined_scale = self.loss_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad + clip = torch.clamp(clip, min=1.0) + combined_scale = clip * self.loss_scale + + for grad in grad_groups_flat: + if isinstance(grad, list): + sub_partitions = grad + for g in sub_partitions: + g.data.mul_(1. / combined_scale) + else: + grad.data.mul_(1. / combined_scale) + + def _check_overflow(self, partition_gradients=True): + self.overflow = self.has_overflow(partition_gradients) + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params): + invalid_grad_count = torch.zeros([1], dtype=torch.float, device=get_accelerator().current_device_name()) + for p in params: + if p.grad is not None: + invalid_grad_count += self._has_inf_or_nan(p.grad) + return invalid_grad_count.bool() + + def has_overflow_partitioned_grads_serial(self): + invalid_grad_count = torch.zeros([1], dtype=torch.float, device=get_accelerator().current_device_name()) + for i in range(len(self.bit16_groups)): + for j, grad in enumerate(self.averaged_gradients[i]): + if grad is not None: + invalid_grad_count += self._has_inf_or_nan(grad) + return invalid_grad_count.bool() + + def has_overflow(self, partition_gradients=True): + if partition_gradients: + overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial() + overflow_gpu = get_accelerator().ByteTensor([overflow]) if self.cpu_offload else overflow.byte().to( + get_accelerator().current_device_name()) + '''This will capture overflow across all data parallel and expert parallel process + Since expert parallel process are a subset of data parallel process''' + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group) + + else: + params = [] + for group in self.bit16_groups: + for param in group: + params.append(param) + overflow_gpu = self.has_overflow_serial(params).byte().to(get_accelerator().current_device_name()) + + # Since each model parallel GPU carries only part of the model, + # make sure overflow flag is synced across all the model parallel GPUs + self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX) + + overflow = overflow_gpu[0].item() + return bool(overflow) + + # `x` is a torch.Tensor + @staticmethod + def _has_inf_or_nan(x, j=None): + float_x = x.float() + nan = float_x.isnan() + inf = float_x.isinf() + inf_or_nan = nan.logical_or(inf) + return inf_or_nan.float().max() + + def backward_prologue(self): + if not self.ready_for_gradients: + self.micro_step_id += 1 + if self.contiguous_gradients and self.ipg_buffer is None: + self.ipg_buffer = [] + buf_0 = torch.empty(int(self.reduce_bucket_size), + dtype=self.dtype, + device=get_accelerator().current_device_name()) + self.ipg_buffer.append(buf_0) + + # Use double buffers to avoid data access conflict when overlap_comm is enabled. + if self.overlap_comm: + buf_1 = torch.empty(int(self.reduce_bucket_size), + dtype=self.dtype, + device=get_accelerator().current_device_name()) + self.ipg_buffer.append(buf_1) + self.ipg_index = 0 + self.ready_for_gradients = True + + def backward_epilogue(self): + # Only for Stage 1, Mode 2 + if self.use_grad_accum_attribute: + self.fill_grad_accum_attribute() + + def backward(self, loss, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + self.backward_prologue() + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward(retain_graph=retain_graph) + else: + self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) + self.backward_epilogue() + + def check_overflow(self, partition_gradients=True): + self._check_overflow(partition_gradients) + + def _update_scale(self, has_overflow=False): + self.loss_scaler.update_scale(has_overflow) + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + + param_groups = property(_get_param_groups, _set_param_groups) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.loss_scaler.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) + cur_scale = property(_get_loss_scale, _set_loss_scale) + + # Return group tensor after removing paddings that are added for alignment to DP world size. + # This method works on the assumption that each group contains a single flattened tensor. + def _get_groups_without_padding(self, groups_with_padding): + groups_without_padding = [] + for i, group in enumerate(groups_with_padding): + lean_length = group.numel() - self.groups_padding[i] + groups_without_padding.append(group[:lean_length]) + + return groups_without_padding + + # Return optimizer state after removing paddings that are added for alignment. + def _get_state_without_padding(self, state_with_padding, padding): + lean_state = {} + for key, value in state_with_padding.items(): + if torch.is_tensor(value): + lean_length = value.numel() - padding + lean_state[key] = value[:lean_length] + else: + lean_state[key] = value + + return lean_state + + # Return base optimizer states. + # This method assumes that each param group contains a single flattened tensor. + def _get_base_optimizer_state(self): + optimizer_groups_state = [] + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + lean_optimizer_state = self._get_state_without_padding(self.optimizer.state[p], self.groups_padding[i]) + optimizer_groups_state.append(lean_optimizer_state) + + return optimizer_groups_state + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + state_dict = {} + state_dict[LOSS_SCALER] = self.loss_scaler + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['overflow'] = self.overflow + state_dict[CLIP_GRAD] = self.clip_grad + + if self.elastic_checkpoint: + state_dict[BASE_OPTIMIZER_STATE] = self._get_base_optimizer_state() + + if "step" in self.optimizer.param_groups[0]: + # Assuming "step" is the only item that changes through training iterations + assert all(group["step"] == self.optimizer.param_groups[0]["step"] + for group in self.optimizer.param_groups), "All param groups must have the same step value" + state_dict[BASE_OPTIMIZER_STATE_STEP] = self.optimizer.param_groups[0]["step"] + else: + state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() + + # Remove paddings for DP alignment to enable loading for other alignment values + fp32_groups_without_padding = self._get_groups_without_padding(self.single_partition_of_fp32_groups) + state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding + + state_dict[ + ZERO_STAGE] = ZeroStageEnum.gradients if self.partition_gradients else ZeroStageEnum.optimizer_states + state_dict[GROUP_PADDINGS] = self.groups_padding + state_dict[PARTITION_COUNT] = self.partition_count + + state_dict[DS_VERSION] = version + state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings + + return state_dict + + # Restore base optimizer fp32 weights from elastic checkpoint by: + # 1) Merging fp32 weights from checkpoints of all partitions + # 2) Extracting fp32 weights for current partition from merged weights + # 3) Using extracted weights to update base optimizer weights directly. + def _restore_from_elastic_fp32_weights(self, all_state_dict): + merged_single_partition_of_fp32_groups = [] + + for i in range(len(self.single_partition_of_fp32_groups)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + merged_partitions = [sd[SINGLE_PARTITION_OF_FP32_GROUPS][i] for sd in all_state_dict] + if self.is_moe_group(self.optimizer.param_groups[i]): + ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) + merged_partitions = [merged_partitions[i] for i in ranks] + flat_merged_partitions = self.flatten_dense_tensors_aligned( + merged_partitions, + self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i])) + dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, i) + merged_single_partition_of_fp32_groups.append(dp_partitions[partition_id]) + + for current, saved in zip(self.single_partition_of_fp32_groups, merged_single_partition_of_fp32_groups): + current.data.copy_(saved.data) + + # Restore base optimizer fp32 weights from ZeRO fp16 or bfloat16 weights + def _restore_from_bit16_weights(self): + for group_id, (bit16_partitions, fp32_partition) in enumerate( + zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): + partition_id = dist.get_rank(group=self.real_dp_process_group[group_id]) + fp32_partition.data.copy_(bit16_partitions[partition_id].data) + + # Refresh the fp32 master params from the fp16 or bfloat16 copies. + def refresh_fp32_params(self): + self._restore_from_bit16_weights() + + # Extract optimizer state for current partition from merged states of all partitions + def _partition_base_optimizer_state(self, state_key, all_partition_states, group_id): + partition_id = dist.get_rank(group=self.real_dp_process_group[group_id]) + alignment = self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[group_id]) + if torch.is_tensor(all_partition_states[0]): + flat_merged_partitions = self.flatten_dense_tensors_aligned(all_partition_states, alignment) + dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, group_id) + return dp_partitions[partition_id] + else: + # Assume non-tensor states are not partitioned and equal across ranks, so return first one + return all_partition_states[0] + + def _restore_step_from_elastic_checkpoint(self, all_state_dict): + assert BASE_OPTIMIZER_STATE_STEP in all_state_dict[0] + assert all(sd[BASE_OPTIMIZER_STATE_STEP] == all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] + for sd in all_state_dict), "State dicts of all partitions must have the same step value" + return all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] + + def _restore_base_optimizer_state(self, base_optimizer_group_states, base_optimizer_state_step, group_paddings): + if type(base_optimizer_group_states) == dict: + base_optimizer_group_states = base_optimizer_group_states['state'] + + saved_keys = base_optimizer_group_states[0].keys() + + for i, group in enumerate(self.optimizer.param_groups): + p = group['params'][0] + padding = 0 if group_paddings is None else group_paddings[i] + for key in saved_keys: + saved = base_optimizer_group_states[i][key] + + if torch.is_tensor(saved): + if key in self.optimizer.state[p]: + dst_tensor = self.optimizer.state[p][key] + src_tensor = _get_padded_tensor(saved, dst_tensor.numel()) + self.optimizer.state[p][key].data.copy_(src_tensor.data) + else: + self.optimizer.state[p][key] = _pad_tensor_by_size( + saved, padding, torch.float32, + torch.device('cpu') if self.cpu_offload else self.device) + else: + self.optimizer.state[p][key] = saved + + for param_group in self.optimizer.param_groups: + param_group['step'] = base_optimizer_state_step + + def get_ep_ranks(self, rank=0, group_name=None): + from deepspeed.utils import groups + expert_parallel_size_ = groups._get_expert_parallel_world_size(group_name) + world_size = groups._get_data_parallel_world_size() + rank = groups._get_expert_parallel_rank(group_name) + ranks = range(rank, world_size, expert_parallel_size_) + return list(ranks) + + # Restore base optimizer state from elastic checkpoint by + # 1) Merging optimizer state from checkpoints of all partitions + # 2) Extracting optimizer state for current partition from the merged state + # 3) Using the extracted value to directly update the base optimizer. + def _restore_elastic_base_optimizer_state(self, all_state_dict): + base_optimizer_group_states = [] + for i in range(len(self.optimizer.param_groups)): + partition_states = {} + all_partition_group_states = [sd[BASE_OPTIMIZER_STATE][i] for sd in all_state_dict] + + if self.is_moe_group(self.optimizer.param_groups[i]): + ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) + all_partition_group_states = [all_partition_group_states[i] for i in ranks] + + for key in all_partition_group_states[0].keys(): + all_partition_states = [all_states[key] for all_states in all_partition_group_states] + partition_states[key] = self._partition_base_optimizer_state(key, all_partition_states, i) + base_optimizer_group_states.append(partition_states) + + self._restore_base_optimizer_state(base_optimizer_group_states, + self._restore_step_from_elastic_checkpoint(all_state_dict), None) + + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None, + param_shapes=None): + if checkpoint_folder: + self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) + else: + self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) + + def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): + self.load_hp_checkpoint_state_from_checkpoint_dir("bit16_groups", checkpoint_folder) + + def _load_global_state(self, sd): + self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler) + self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale) + self.overflow = sd.get('overflow', self.overflow) + self.clip_grad = sd.get(CLIP_GRAD, self.clip_grad) + + ckpt_version = sd.get(DS_VERSION, False) + assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" + ckpt_version = pkg_version.parse(ckpt_version) + + # zero stage 1 mode + if not self.partition_gradients: + required_version = pkg_version.parse("0.3.17") + error_str = f"ZeRO stage 1 changed in {required_version} and is not backwards compatible " \ + "with older stage 1 checkpoints. If you'd like to load an old ZeRO-1 checkpoint " \ + "please use an older version of DeepSpeed (<= 0.5.8) and set 'legacy_stage1': true in your zero config json." + assert required_version <= ckpt_version, f"Old version: {ckpt_version} {error_str}" + + def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): + r"""Loading ZeRO checkpoint + + Arguments: + state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition. + Note that the number of saved partitions may differ from number of loading partitions to support + changing GPU count, specifically DP world size, between saving and loading checkpoints. + load_optimizer_states: Boolean indicating whether or not to load base optimizer states + load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32 + copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss). + """ + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + + # I think it should actually be ok to reload the optimizer before the model. + dp_rank = dist.get_rank(group=self.dp_process_group) + current_rank_sd = state_dict_list[dp_rank] + self._load_global_state(current_rank_sd) + + ckpt_is_rigid = isinstance(current_rank_sd[BASE_OPTIMIZER_STATE], dict) + + # padding is always at the last rank/partition + # if DP=1024 and param-group elems=16 -> padding will be 1024-16 across all but one rank + # scenario-1 (shrink): saving w. 4 gpus -> loading w. 2 gpus + # scenario-2 (expand): saving w. 2 gpus -> loading w. 4 gpus + # if load_optimizer_states: + # if new_dp_size: + # self.strip_padding() + # self.add_padding_w_new_dp_size() + # self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) + + if load_optimizer_states: + if ckpt_is_rigid: + # loading rigid ckpt into either rigid or elastic exec + self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) + else: + if self.elastic_checkpoint: + # loading elastic into elastic exec + self._restore_elastic_base_optimizer_state(state_dict_list) + else: + # loading an elastic checkpoint into rigid exec + self._restore_base_optimizer_state(current_rank_sd[BASE_OPTIMIZER_STATE], + current_rank_sd[BASE_OPTIMIZER_STATE_STEP], + current_rank_sd[GROUP_PADDINGS]) + + # At this point, the optimizer's references to the model's fp32 parameters are up to date. + # The optimizer's hyperparameters and internal buffers are also up to date. + # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still + # out of date. There are two options. + # 1: Refresh the master params from the model's fp16 params. + # This requires less storage but incurs precision loss. + # 2: Save and restore the fp32 master copies separately. + # We choose option 1 if changing DP degree and option 2 otherwise. + # + # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device + # of their associated parameters, because it's possible those buffers might not exist yet in + # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been + # constructed in the same way as the one whose state_dict we are loading, the same master params + # are guaranteed to exist, so we can just copy_() from the saved master params. + + if load_from_fp32_weights: + # option 2 from above + if self.elastic_checkpoint and not ckpt_is_rigid: + self._restore_from_elastic_fp32_weights(state_dict_list) + else: + # For non-elastic checkpoint, simply copying from saved weights of current rank is sufficient. + for current, saved in zip(self.single_partition_of_fp32_groups, + current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): + src_tensor = _get_padded_tensor(saved, current.numel()) + current.data.copy_(src_tensor.data) + else: + # option 1 from above + self._restore_from_bit16_weights() + + if load_optimizer_states: + self._link_all_hp_params() + + +def _handle_overflow(cpu_sum, x, i): + import math + rank = dist.get_rank() + if rank == 0: + t_i = -1 + for v_i, v in enumerate(x.data.contiguous().view(-1)): + if not math.isfinite(float(v)): + t_i = v_i + break + logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") + + +def estimate_zero2_model_states_mem_needs(total_params, + num_gpus_per_node=1, + num_nodes=1, + cpu_offload=True, + additional_buffer_factor=1.5): + + total_gpus = num_nodes * num_gpus_per_node + + if cpu_offload: + gpu_mem = 2 * total_params + cpu_mem = total_params * max(4 * total_gpus, 16) * additional_buffer_factor + else: + # GPU's total_params multipliers: 2 = params_16bit, + # 18 = 2_grads_16bit + 4_grads_32bit + 4_params_32bit + 8_optimizer_states_32bit(momentum and variance) + gpu_mem = 2 * total_params + int(18 * total_params / total_gpus) + cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor + + return int(cpu_mem), int(gpu_mem) + + +def model_to_params(model): + # shared params calculated only once + total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) + return total_params + + +def estimate_zero2_model_states_mem_needs_all_live(model, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients + for a given ``model`` and hardware setup. + + If you have an actual model object, use this function and everything will be derived + automatically. + + If it's a hypothetical model, use ``estimate_zero2_model_states_mem_needs_all_cold`` where you have to pass + the ``total_params`` explicitly. + + Args: + - ``model``: ``nn.Module`` object + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + total_params = model_to_params(model) + + estimate_zero2_model_states_mem_needs_all_cold(total_params=total_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + additional_buffer_factor=additional_buffer_factor) + + +def estimate_zero2_model_states_mem_needs_all_cold(total_params, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients + for a given ``model`` and hardware setup. + + If it's a hypothetical model, use this function where you have to pass + the ``total_params`` and ``largest_layer_params`` explicitly. + + If you have an actual model object, use ``estimate_zero2_model_states_mem_needs_all_live`` and everything + will be derived automatically. + + Args: + - ``total_params``: total model params + - ``num_gpus_per_node``: how many gpus per node (defaults to 1) + - ``num_nodes``: how many nodes (defaults to 1), + - ``additional_buffer_factor``: estimation factor (defaults to 1.5): + + """ + + def format_options(cpu_offload): + enabled = [] + device = f'{OffloadDeviceEnum.cpu:4}' if cpu_offload else "none" + enabled.append(f"offload_optimizer={device}") + return ", ".join(enabled) + + nodes_str = "nodes" if num_nodes > 1 else "node" + gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU" + print("Estimated memory needed for params, optim states and gradients for a:\n" + f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n" + f"SW: Model with {int(total_params/1e6)}M total params.") + print(" per CPU | per GPU | Options") + for cpu_offload in [True, False]: + cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params=total_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + cpu_offload=cpu_offload, + additional_buffer_factor=additional_buffer_factor) + + options_str = format_options(cpu_offload=cpu_offload) + print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}") diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/test.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/test.py new file mode 100644 index 0000000000000000000000000000000000000000..1904a8d395be80a4322bc25bd1008a80bdd8f982 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/test.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.runtime.zero.contiguous_memory_allocator import ContiguousMemoryAllocator + + +def test1(): + mem = ContiguousMemoryAllocator(1024, torch.half, 'cpu') + mem.print_allocation(resolution=100) + a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0) + mem.print_allocation(resolution=100) + mem.release_tensor(a1) + mem.print_allocation(resolution=100) + a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0) + a3 = mem.allocate_tensor(256).mul_(0.0).add_(3.0) + a4 = mem.allocate_tensor(128).mul_(0.0).add_(4.0) + mem.print_allocation(resolution=100) + mem.release_tensor(a3) + mem.print_allocation(resolution=100) + a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0) + a6 = mem.allocate_tensor(256).mul_(0.0).add_(6.0) + a7 = mem.allocate_tensor(128).mul_(0.0).add_(7.0) + mem.print_allocation(resolution=100) + a8 = mem.allocate_tensor(256).mul_(0.0).add_(8.0) + a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0) + mem.print_allocation(resolution=100) + mem.release_tensor(a9) + mem.release_tensor(a6) + mem.release_tensor(a2) + mem.release_tensor(a5) + + a10 = mem.allocate_tensor(512).mul_(0.0).add_(10.0) + mem.print_allocation(resolution=100) + #print(f"a4:{a4}") + #print(f"a7:{a7}") + #print(f"a8:{a8}") + #print(f"a10:{a10}") + assert (a4.norm() + a7.norm() + a8.norm() + a10.norm()).item() == 474.50, "Test failed" + + +def test2(): + mem = ContiguousMemoryAllocator(512, torch.half, 'cpu') + a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0) + a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0) + a3 = mem.allocate_tensor(64).mul_(0.0).add_(3.0) + a4 = mem.allocate_tensor(64).mul_(0.0).add_(4.0) + a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0) + a6 = mem.allocate_tensor(64).mul_(0.0).add_(6.0) + a7 = mem.allocate_tensor(64).mul_(0.0).add_(7.0) + a8 = mem.allocate_tensor(64).mul_(0.0).add_(8.0) + mem.release_tensor(a2) + mem.release_tensor(a4) + mem.release_tensor(a6) + mem.release_tensor(a8) + mem.print_allocation(resolution=100) + + a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0) + a10 = mem.allocate_tensor(64).mul_(0.0).add_(10.0) + a11 = mem.allocate_tensor(64).mul_(0.0).add_(11.0) + mem.release_tensor(a1) + mem.release_tensor(a5) + mem.print_allocation(resolution=100) + a12 = mem.allocate_tensor(128).mul_(0.0).add_(12.0) + mem.print_allocation(resolution=100) + print(f"a7:{a7}") + print(f"a9:{a9}") + print(f"a10:{a10}") + print(f"a11:{a11}") + print(f"a12:{a12}") + assert (a7.norm() + a9.norm() + a10.norm() + a11.norm() + a12.norm()) == 460.75, "TestFailed" + + +test1() +test2() diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/tiling.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/tiling.py new file mode 100644 index 0000000000000000000000000000000000000000..af1c5e54737a75cb66b39736718bf62169952a80 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/tiling.py @@ -0,0 +1,296 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import deepspeed +from deepspeed.runtime.utils import partition_uniform as partition + + +def split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=False): + """Split a tensor along its last dimension. Adapted from Megatron-LM. + + Arguments: + tensor: input tensor. + partitions: list of partition sizes to supply to torch.split + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + + # Split. + tensor_list = torch.split(tensor, partitions, dim=last_dim) + # Note: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + +class TiledLinear(torch.nn.Module): + + def __init__(self, + in_features, + out_features, + bias=True, + in_splits=1, + out_splits=1, + input_is_already_split=False, + combine_out_splits=True, + linear_cls=torch.nn.Linear, + init_linear=None, + **kwargs): + """A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce + memory requirements via tiling. + + TiledLinear breaks the input and output dimensions of a linear layer + into tiles that are processed in sequence. This class enables huge + linear layers when combined with ZeRO-3 because inactive tiles can be + partitioned and offloaded. + + .. note:: + We recommend using as few tiles as necessary. Tiling + significantly reduces memory usage, but can reduce throughput + for inexpensive layers. This due to the smaller kernels having + less parallelism and lower arithmetic intensity, while + introducing more frequent synchronization and communication. + + Args: + in_features (int): See ``torch.nn.Linear`` + out_features (int): See ``torch.nn.Linear`` + bias (bool, optional): See ``torch.nn.Linear`` + in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1. + out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1. + input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in + to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``. + combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs + into a single tensor. Defaults to ``True``. + linear_cls (class, optional): The underlying class to build individual tiles. + Defaults to ``torch.nn.Linear``. + init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of + ``init_linear``. Useful for debugging. Defaults to ``None``. + kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``. + + Raises: + RuntimeError: ``in_splits`` must be within the range [1, in_features). + RuntimeError: ``out_splits`` must be within the range of [1, out_features). + """ + + super().__init__() + + if (in_splits < 1) or (in_splits > in_features): + raise RuntimeError('in splits must be in range [1, in_features].') + if (out_splits < 1) or (out_splits > out_features): + raise RuntimeError('out splits must be in range [1, out_features].') + + # global, not necessarily local + self.in_features = in_features + self.out_features = out_features + self.use_bias = bias + + self.out_splits = out_splits + self.in_splits = in_splits + self.input_is_already_split = input_is_already_split + self.combine_out_splits = combine_out_splits + + # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features] + # For example, row_parts[p] gives the start of partition p and row_parts[p+1] + # is the exclusive end. + self.in_parts = partition(num_items=in_features, num_parts=in_splits) + self.out_parts = partition(num_items=out_features, num_parts=out_splits) + + assert len(self.out_parts) == out_splits + 1 + assert len(self.in_parts) == in_splits + 1 + assert self.out_parts[0] == 0 + assert self.out_parts[out_splits] == out_features + assert self.in_parts[in_splits] == in_features + + self.linears = torch.nn.ModuleList() + for out_id in range(out_splits): + self.linears.append(torch.nn.ModuleList()) + + local_out_dim = self.out_parts[out_id + 1] - self.out_parts[out_id] + + for in_id in range(in_splits): + #if input_size is split, we only need one bias + local_bias = bias if in_id == (in_splits - 1) else False + + local_in_dim = self.in_parts[in_id + 1] - self.in_parts[in_id] + local = linear_cls(local_in_dim, local_out_dim, bias=local_bias, **kwargs) + self.linears[out_id].append(local) + + # Optionally initialize with a known tensor + if init_linear is not None: + self.copy_params_from(init_linear) + + def forward(self, input_): + if self.in_splits > 1 and not self.input_is_already_split: + input_parts = partition(input_.shape[-1], self.in_splits) + split_sizes = [input_parts[p + 1] - input_parts[p] for p in range(self.in_splits)] + inputs = self._split_global_input(input_, split_sizes) + elif self.in_splits > 1: + inputs = input_ + assert len( + inputs) == self.in_splits, f"Col splits {self.in_splits} does not match input splits {len(inputs)}" + else: + # no splits + inputs = [input_] + + outputs = [None] * self.out_splits + for out_id in range(self.out_splits): + for in_id in range(self.in_splits): + local_output = self.linears[out_id][in_id](inputs[in_id]) + + outputs[out_id] = self._reduce_local_output(in_id=in_id, + out_id=out_id, + current_out=outputs[out_id], + new_out=local_output) + + if self.combine_out_splits: + return self._combine_output_splits(outputs) + + return outputs + + def _split_global_input(self, input, split_sizes): + """Partition an input tensor along the last dimension, aligned with given splits. + + Subclasses should override this method to account for new input types. + + Args: + input (List[Tensor]): The tensor to partition along the last dimension. + split_sizes (List[int]): The size of each partition. + + Returns: + List[Any]: A list of the chunks of ``input``. + """ + return split_tensor_along_last_dim(input, split_sizes) + + def _reduce_local_output(self, in_id, out_id, current_out, new_out): + """Reduce (sum) a new local result into the existing local results. + + Subclasses should override this method. + + For a given ``out_id``, this method is called ``in_id-1`` times. The first input + split is a simple assignment. + + Args: + in_id (int): The input split that produced ``new_out``. + out_id (int): The output split that produced ``new_out``. + current_out (Any): The reduced form of all previous ``out_id`` results. + new_out (Any): The local result from forward (``in_id``, ``out_id``)e + + Returns: + Any: The combined result of ``current_out`` and ``new_out``. + """ + + if current_out is None: + #this clone is necessary to preserve auto grad + #there is some issue with inplace update for outputs that are views + return new_out.clone() + else: + return current_out + new_out + + def _combine_output_splits(self, outputs): + """Join the splits of the output into a single result. + + Args: + outputs (List[Any]): The reduced outputs for each output split. + + Returns: + Any: The combined outputs. + """ + assert len(outputs) == self.out_splits + return torch.cat(outputs, dim=-1) + + @torch.no_grad() + def copy_params_from(self, other): + """Copy the weight and bias data from ``other``. + + This is especially useful for reproducible initialization and testing. + + Equivalent to: + + .. code-block:: python + + with torch.no_grad(): + self.weight.copy_(other.weight) + if self.bias is not None: + self.bias.copy_(other.bias) + + .. note:: + If ZeRO-3 is enabled, this is a collective operation and the + updated parameters of data-parallel rank 0 will be visible on all + ranks. See :class:`deepspeed.zero.GatheredParameters` for more + information. + + + Args: + other (``torch.nn.Linear``): the linear layer to copy from. + """ + assert hasattr(other, 'weight') + assert other.weight.size() == (self.out_features, self.in_features) + if self.use_bias: + assert hasattr(other, 'bias') + assert other.bias is not None + assert other.bias.size() == (self.out_features, ) + else: + assert other.bias is None + + for row in range(self.out_splits): + rstart = self.out_parts[row] + rstop = self.out_parts[row + 1] + + for col in range(self.in_splits): + cstart = self.in_parts[col] + cstop = self.in_parts[col + 1] + + local = self.linears[row][col] + global_weight = other.weight[rstart:rstop, cstart:cstop] + with deepspeed.zero.GatheredParameters(local.weight, modifier_rank=0): + local.weight.copy_(global_weight) + + if local.bias is not None: + with deepspeed.zero.GatheredParameters(local.bias, modifier_rank=0): + local.bias.data.copy_(other.bias[rstart:rstop].data) + + +class TiledLinearReturnBias(TiledLinear): + """Wrapper for a Linear class that returns its own bias parameter, such as + used by Megatron-LM. + """ + + def _reduce_local_output(self, in_id, out_id, current_out, new_out): + """Reduces output tensors, but not the returned bias. """ + if current_out is not None: + old_tensor, old_bias = current_out + else: + old_tensor, old_bias = None, None + + assert isinstance(new_out, tuple) + assert len(new_out) == 2 + + tensor, bias = new_out + assert tensor is not None + + tensor = super()._reduce_local_output(in_id=in_id, out_id=out_id, current_out=old_tensor, new_out=tensor) + + if bias is None: + bias = old_bias + + return tensor, bias + + def _combine_output_splits(self, outputs): + # stack output tensors + tensors = [o[0] for o in outputs] + tensor = super()._combine_output_splits(tensors) + + # stack biases if applicable + biases = [o[1] for o in outputs if o[1] is not None] + if len(biases) > 0: + bias = super()._combine_output_splits(biases) + else: + bias = None + + return tensor, bias diff --git a/lib/python3.12/site-packages/deepspeed/runtime/zero/utils.py b/lib/python3.12/site-packages/deepspeed/runtime/zero/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2d1cf17962d8a59fee5cb1ac790129fcb6c53a10 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/zero/utils.py @@ -0,0 +1,175 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from typing import List, Tuple + +import torch +from deepspeed import comm as dist +from deepspeed.utils import logger +from deepspeed.ops.adam import DeepSpeedCPUAdam +from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad +from deepspeed.ops.adam import FusedAdam +from deepspeed.ops.lion import DeepSpeedCPULion, FusedLion +from deepspeed.utils.nvtx import instrument_w_nvtx +from deepspeed.accelerator import get_accelerator + +# ensure we only warn once, otherwise every iteration will trigger a warning +warned = False + + +def _initialize_parameter_parallel_groups(parameter_parallel_size=None): + data_parallel_size = int(dist.get_world_size()) + parameter_parallel_size = parameter_parallel_size or data_parallel_size + logger.info("data_parallel_size: %s, parameter_parallel_size: %s", data_parallel_size, parameter_parallel_size) + assert data_parallel_size % parameter_parallel_size == 0, \ + 'world size should be divisible by parameter parallel size' + rank = dist.get_rank() + my_group = None + for i in range(data_parallel_size // parameter_parallel_size): + ranks = range(i * parameter_parallel_size, (i + 1) * parameter_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + my_group = group + return my_group + + +class ZeRORuntimeException(Exception): + pass + + +ZERO_SUPPORTED_OPTIMIZERS = [ + torch.optim.Adam, torch.optim.AdamW, FusedAdam, DeepSpeedCPUAdam, torch.optim.Adagrad, DeepSpeedCPUAdagrad, + DeepSpeedCPULion, FusedLion +] + +# Add apex FusedAdam to supported list if apex is installed +try: + import apex + if hasattr(apex, 'optimizers') and hasattr(apex.optimizers, 'FusedAdam'): + ZERO_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam) +except ImportError: + pass + + +def is_zero_supported_optimizer(optimizer): + if dist.get_rank() == 0: + logger.info(f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}') + return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS + + +def get_lst_from_rank0(lst: List[int]) -> None: + """ + NOTE: creates both communication and synchronization overhead so should be used + sparingly + """ + lst_tensor = torch.tensor( + lst if dist.get_rank() == 0 else [-1] * len(lst), + dtype=int, + device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])), + requires_grad=False, + ) + dist.broadcast(lst_tensor, src=0, async_op=False) + + return list(lst_tensor.cpu().numpy()) + + +@instrument_w_nvtx +def assert_ints_same_as_other_ranks(ints: List[int]) -> None: + """ + NOTE: creates both communication and synchronization overhead so should be + used sparingly + + takes a list of ints from each rank and ensures that they are the same + across ranks, throwing an exception if they are not. + """ + rank0_ints = get_lst_from_rank0(ints) + if ints != rank0_ints: + raise RuntimeError(f"disagreement between rank0 and rank{dist.get_rank()}: " + f"rank0: {rank0_ints}, rank{dist.get_rank()}: {ints}") + + +def is_builtin_type(obj): + # https://stackoverflow.com/a/17795199 + return obj.__class__.__module__ == '__builtin__' or obj.__class__.__module__ == "builtins" + + +def isinstance_namedtuple(obj: object) -> bool: + """ + Is this an instance of namedtuple/NamedTuple? + From: https://stackoverflow.com/a/62692640 + + Args: + obj (object): An object. + + Returns: + bool: True if namedtuple/NamedTuple else False. + """ + return isinstance(obj, tuple) and hasattr(obj, '_asdict') and hasattr(obj, '_fields') + + +def is_zero_param(parameter): + if not torch.is_tensor(parameter): + return False + return hasattr(parameter, 'ds_id') + + +def apply_to_tensors_only(function, value, warning_msg_fn=None): + """ + Apply `function` to every Tensor in `value`. + + Args: + functional: The function class to apply. + value (Any): Target object to apply `function` to. + + Returns: + Any: Output of `function`. + """ + if isinstance(value, (tuple, list)): + touched_outputs = [] + for elem in value: + touched_output = apply_to_tensors_only(function, elem) + touched_outputs.append(touched_output) + + if isinstance_namedtuple(value): + # namedtuples require a slightly different syntax. + return value.__class__(*touched_outputs) + + return value.__class__(touched_outputs) + elif isinstance(value, dict): + # apply inplace to avoid recreating dict inherited objects + for key in value.keys(): + value[key] = apply_to_tensors_only(function, value[key]) + return value + + elif isinstance(value, torch.Tensor): + # this also applies to torch.Tensor's subclasses like torch.nn.parameter.Parameter + touched_output = function(value) + + # restore zero param attributes if those get stripped by `backward_function` + if not is_zero_param(touched_output) and is_zero_param(value): + touched_output.ds_param_alias = value + + return touched_output + else: + if not is_builtin_type(value): + global warned + if warning_msg_fn and not warned and dist.get_rank() == 0: + logger.warning(warning_msg_fn(value)) + warned = True + return value + + +def get_mapping_to_flat_buffer(tensors: List[torch.Tensor]) -> List[Tuple[torch.Tensor, int, int]]: + tensor_infos: List[Tuple[torch.Tensor, int, int]] = [] + + offset = 0 + for tensor in tensors: + tensor_numel = tensor.numel() + # record some data so we can restore the device tensor later + tensor_infos.append((tensor, offset, tensor_numel)) + offset += tensor_numel + + return tensor_infos diff --git a/lib/python3.12/site-packages/deepspeed/sequence/__init__.py b/lib/python3.12/site-packages/deepspeed/sequence/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..208299fb8c50f73468d293b6fa5dca71649d62e7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/sequence/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1adf286011ada47e15d9960dda70001b67ae0d9f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d56774beef2fcb494e6fdc911b9a9a4f67edcac Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39b22645376908090a7735aa520e2949af7429a9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..178fc80ee83d596bf0d4dd825420e11be18ecd3e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/sequence/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/sequence/cross_entropy.py b/lib/python3.12/site-packages/deepspeed/sequence/cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..baa7bc1ea7a803055f2cca20141ea6888cc3f8c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/sequence/cross_entropy.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +import deepspeed.comm as dist + + +class _VocabSequenceParallelCrossEntropy(torch.autograd.Function): + + @staticmethod + def forward(ctx, vocab_seq_parallel_logits, target, sp_group): + # vocab_seq_parallel_logits: [S/P, B, V] + # target: [S/P, B] + # return: [S, B] + + # Need softmax for backward + softmax = torch.nn.functional.softmax(vocab_seq_parallel_logits, dim=-1) + ctx.vocab_size = vocab_seq_parallel_logits.size(2) + loss = torch.nn.functional.nll_loss(softmax.log().view(-1, ctx.vocab_size), target.view(-1), reduction='none') + + sp_world_size = dist.get_world_size(sp_group) + sp_rank = dist.get_rank(sp_group) + ctx.sp_world_size = sp_world_size + ctx.sp_rank = sp_rank + ctx.seqlen = vocab_seq_parallel_logits.size(0) * sp_world_size + batch_size = vocab_seq_parallel_logits.size(1) + + loss_all = torch.empty(ctx.seqlen, + batch_size, + dtype=vocab_seq_parallel_logits.dtype, + device=vocab_seq_parallel_logits.device) + dist.all_gather_into_tensor(loss_all, loss, group=sp_group) + + ctx.save_for_backward(softmax, target) + + return loss_all + + @staticmethod + def backward(ctx, grad_output): + softmax, target = ctx.saved_tensors + + step_seqlen = ctx.seqlen // ctx.sp_world_size + sp_rank = ctx.sp_rank + grad_output_part = grad_output[step_seqlen * sp_rank:step_seqlen * (sp_rank + 1), :] + + grad_input = softmax + grad_2d = grad_input.view(-1, ctx.vocab_size) + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) + + grad_2d[arange_1d, target.view(-1)] -= 1 + grad_input.mul_(grad_output_part.unsqueeze(dim=-1)) + + return grad_input, None, None, None + + +def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, target, sp_group): + return _VocabSequenceParallelCrossEntropy.apply(vocab_parallel_logits, target, sp_group) diff --git a/lib/python3.12/site-packages/deepspeed/sequence/fpdt_layer.py b/lib/python3.12/site-packages/deepspeed/sequence/fpdt_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..4fa2cc988a19490f91c90867fe9a01d8a120e2e8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/sequence/fpdt_layer.py @@ -0,0 +1,1225 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from typing import Optional, Any, Tuple +from torch import Tensor +from packaging import version +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator + +try: + import flash_attn + from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward + flash_attn_version = version.parse(flash_attn.__version__) +except ImportError: + _flash_attn_forward = None + _flash_attn_backward = None + +from einops import rearrange +from .layer import single_all_to_all, apply_rotary_pos_emb + + +def _rotate_half_backward(x): + x = rearrange(x, '... (j d) -> ... j d', j=2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((x2, -x1), dim=-1) + + +def apply_rotary_pos_emb_backward(grad_output, freqs_cos, freqs_sin): + rot_dim = freqs_cos.shape[-1] + grad, grad_pass = grad_output[..., :rot_dim], grad_output[..., rot_dim:] + grad_t = (grad * freqs_cos) + (_rotate_half_backward(grad * freqs_sin)) + grad = grad_t if grad_pass.shape[-1] == 0 else torch.cat((grad_t, grad_pass), dim=-1) + return grad + + +def _update_out_and_lse( + out: torch.Tensor, + lse: torch.Tensor, + block_out: torch.Tensor, + block_lse: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + + block_out = block_out.to(torch.float32) + block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1) + + new_lse = lse + torch.log1p(torch.exp(block_lse - lse)) + + out = torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out + + lse = new_lse + return out, lse + + +def update_out_and_lse( + out: Optional[torch.Tensor], + lse: Optional[torch.Tensor], + block_out: torch.Tensor, + block_lse: torch.Tensor, + slice_=None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if out is None: + if slice_ is not None: + raise RuntimeError("first update_out_and_lse should not pass slice_ args") + out = block_out.to(torch.float32) + lse = block_lse.permute(0, 2, 1).contiguous().unsqueeze(dim=-1).contiguous() + elif slice_ is not None: + slice_out, slice_lse = out[slice_], lse[slice_] + slice_out, slice_lse = _update_out_and_lse(slice_out, slice_lse, block_out, block_lse) + out[slice_], lse[slice_] = slice_out, slice_lse + else: + out, lse = _update_out_and_lse(out, lse, block_out, block_lse) + return out, lse + + +class FPDT_InputConstruct(torch.nn.Module): + + def __init__(self, tokens, labels, loss_mask, attention_mask, position_ids, args, sp_size, sp_rank) -> None: + + super(FPDT_InputConstruct, self).__init__() + self.tokens = tokens + self.labels = labels + self.loss_mask = loss_mask + self.attention_mask = attention_mask + self.position_ids = position_ids + global_seq_len = tokens.shape[1] + batch_size = tokens.shape[0] + assert global_seq_len % sp_size == 0 + assert global_seq_len % args.ds_sequence_parallel_fpdt_chunk_size == 0 + num_chunk_per_gpu = global_seq_len // args.ds_sequence_parallel_fpdt_chunk_size + local_seq_len = global_seq_len // sp_size + assert local_seq_len % num_chunk_per_gpu == 0 + + self.num_chunk_per_gpu = num_chunk_per_gpu + self.chunk_size = local_seq_len // num_chunk_per_gpu + self.sp_size = sp_size + self.sp_rank = sp_rank + self.global_seq_len = global_seq_len + self.local_seq_len = local_seq_len + self.batch_size = batch_size + self.device = tokens.device + + def generate(self): + device = self.device + totalChunks = self.global_seq_len // self.chunk_size + token_chunk_idx = torch.arange(self.global_seq_len, device=device, dtype=torch.int) // self.chunk_size + chunk_to_gpu = torch.arange(totalChunks, device=device, dtype=torch.int) + chunk_to_gpu = chunk_to_gpu.reshape(self.num_chunk_per_gpu, -1).t().contiguous() + + gather_chunk = chunk_to_gpu.flatten().unsqueeze(1).contiguous() + mask = gather_chunk == token_chunk_idx + + indices = mask.nonzero(as_tuple=False) + gather_indices = indices[:, 0] + token_chunk_indices = indices[:, 1] + indices = torch.cat([token_chunk_indices[gather_indices == i] for i in range(gather_chunk.shape[0])]) + load_balanced_loss_mask = self.loss_mask[:, indices] if self.loss_mask is not None else self.loss_mask + + indices = indices.reshape(-1, self.chunk_size)[self.num_chunk_per_gpu * self.sp_rank:self.num_chunk_per_gpu * + (self.sp_rank + 1)].flatten().contiguous() + load_balanced_tokens = self.tokens[:, indices] + load_balanced_labels = self.labels[:, indices] if self.labels is not None else self.labels + + load_balanced_attention_mask = self.attention_mask if self.attention_mask is not None else self.attention_mask + load_balanced_position_ids = self.position_ids[:, + indices] if self.position_ids is not None else self.position_ids + + return load_balanced_tokens, load_balanced_labels, load_balanced_loss_mask, load_balanced_attention_mask, load_balanced_position_ids + + +class _FPDTGPUAttentionImpl_(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, + layernorm_output, + attention_mask, + inference_params, + rotary_pos_emb, + spg, + scatter_idx, + gather_idx, + hidden_size, + projection_size, + hidden_size_per_attention_head, + kv_projection_size, + qkv_linear_weight, + qkv_linear_bias, + dropout, + num_chunks=8, + cpu_offloading=True): + + do_save = layernorm_output.requires_grad + + if rotary_pos_emb is not None: + pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) + ctx.pos_emb_cos = pos_emb_cos + ctx.pos_emb_sin = pos_emb_sin + else: + ctx.pos_emb_cos = None + ctx.pos_emb_sin = None + + with torch.no_grad(): + per_gpu_seq_len = layernorm_output.shape[0] + chunk_size = per_gpu_seq_len // num_chunks + assert chunk_size * num_chunks == per_gpu_seq_len + assert attention_mask is None + ctx.num_chunks = num_chunks + ctx.cpu_offloading = cpu_offloading + ctx.spg = spg + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + device = get_accelerator().current_device_name() + ctx.device = device + ctx.dtype = layernorm_output.dtype + ctx.projection_size = projection_size + ctx.kv_projection_size = kv_projection_size + + global_q = [] + global_k = [] + global_v = [] + + ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) + + ctx.dropout_p = dropout + ctx.window_size = (-1, -1) + ctx.alibi_slopes = None + + batch_size = layernorm_output.shape[1] + + global_o = [None for _ in range(num_chunks)] + global_lse = [None for _ in range(num_chunks)] + + for i in range(num_chunks): + + st = chunk_size * i + ed = st + chunk_size + + qkv_chunk = torch.matmul(layernorm_output[st:ed], qkv_linear_weight.t()) + qkv_linear_bias + + q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) + global_q_chunk_len = q_chunk.shape[1] + if rotary_pos_emb is not None: + q_chunk = apply_rotary_pos_emb(q_chunk, + pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], + pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) + global_q.append(q_chunk) + + k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) + if rotary_pos_emb is not None: + k_chunk = apply_rotary_pos_emb(k_chunk, + pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], + pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) + global_k.append(k_chunk) + + v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) + global_v.append(v_chunk) + + for k_i in range(len(global_k)): + causal_chunk = i == k_i + if flash_attn_version >= version.parse("2.6.0"): + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], + global_k[k_i], + global_v[k_i], + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + softcap=0.0, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + else: + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], + global_k[k_i], + global_v[k_i], + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + + global_o[i], global_lse[i] = update_out_and_lse(global_o[i], global_lse[i], block_out, block_lse) + + global_o[i] = global_o[i].to(q_chunk.dtype) + + output = [None for i in range(num_chunks)] + + for i in range(num_chunks): + global_lse[i] = global_lse[i][:, :, :, 0].permute(0, 2, 1).contiguous() + output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) + output = torch.cat(output, dim=1) + + head_dim = output.shape[-1] + + if do_save: + ctx.save_for_backward(layernorm_output) + ctx.global_q = global_q + ctx.global_k = global_k + ctx.global_v = global_v + ctx.attn_output = global_o + ctx.attn_lse = global_lse + ctx.head_dim = head_dim + ctx.batch_size = batch_size + + ctx.qkv_linear_weight = qkv_linear_weight + ctx.qkv_linear_bias = qkv_linear_bias + + return output + + @staticmethod + def backward(ctx, grad_output): + + num_chunks = ctx.num_chunks + device = ctx.device + dtype = ctx.dtype + spg = ctx.spg + scatter_idx = ctx.scatter_idx + gather_idx = ctx.gather_idx + softmax_scale = ctx.softmax_scale + dropout_p = ctx.dropout_p + window_size = ctx.window_size + alibi_slopes = ctx.alibi_slopes + + projection_size = ctx.projection_size + kv_projection_size = ctx.kv_projection_size + + layernorm_output = ctx.saved_tensors[0] + + global_q = ctx.global_q + global_k = ctx.global_k + global_v = ctx.global_v + attn_output = ctx.attn_output + lse = ctx.attn_lse + + qkv_linear_weight = ctx.qkv_linear_weight + qkv_linear_bias = ctx.qkv_linear_bias + + input_chunk_size = layernorm_output.shape[0] // num_chunks + grad_layernorm_output = [ + torch.zeros((input_chunk_size, layernorm_output.shape[1], layernorm_output.shape[2]), + device=device, + dtype=dtype) for _ in range(num_chunks) + ] + + grad_global_attn_output = [] + chunk_size = grad_output.shape[1] // num_chunks + + for i in range(num_chunks): + st = chunk_size * i + ed = st + chunk_size + grad_global_attn_output.append( + single_all_to_all(grad_output[:, st:ed].contiguous(), scatter_idx, gather_idx, 0, spg)) + + del grad_output + + dq = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] + dk = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] + dv = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] + + grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, + device=qkv_linear_weight.device, + dtype=torch.float) + grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, device=qkv_linear_weight.device, dtype=torch.float) + + for i in range(num_chunks): + k_chunk = global_k[i] + v_chunk = global_v[i] + + for q_i in range(num_chunks): + no_computation = q_i < i + if no_computation: + continue + + causal_chunk = q_i == i + + q_chunk = global_q[q_i] + attn_output_chunk = attn_output[q_i] + lse_chunk = lse[q_i] + d_out = grad_global_attn_output[q_i] + + dq_this = torch.zeros(global_q[0].shape, dtype=dtype, device=device) + dk_this = torch.zeros(global_k[0].shape, dtype=dtype, device=device) + dv_this = torch.zeros(global_v[0].shape, dtype=dtype, device=device) + + if flash_attn_version >= version.parse("2.6.0"): + _flash_attn_backward(d_out, + q_chunk, + k_chunk, + v_chunk, + attn_output_chunk, + lse_chunk, + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + softcap=0.0, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + else: + _flash_attn_backward(d_out, + q_chunk, + k_chunk, + v_chunk, + attn_output_chunk, + lse_chunk, + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + + dq[q_i].add_(dq_this.to(torch.float)) + dk[i].add_(dk_this.to(torch.float)) + dv[i].add_(dv_this.to(torch.float)) + + dk_seq_len = dk[i].shape[1] + + if ctx.pos_emb_cos is not None: + dk[i] = apply_rotary_pos_emb_backward(dk[i].to(dtype), + ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) + else: + dk[i] = dk[i].to(dtype) + dv[i] = dv[i].to(dtype) + dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg) + dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg) + + input_st = i * input_chunk_size + input_ed = input_st + input_chunk_size + + input_chunk = layernorm_output[input_st:input_ed].reshape(-1, layernorm_output.shape[-1]) + + dk[i] = dk[i].flatten(2).permute(1, 0, 2) + dv[i] = dv[i].flatten(2).permute(1, 0, 2) + l, b = dk[i].shape[0], dk[i].shape[1] + grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( + torch.matmul(dk[i].reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( + torch.matmul(dv[i].reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk[i].sum(0).sum(0)) + grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv[i].sum(0).sum(0)) + + grad_layernorm_output[i].add_( + torch.matmul(dk[i], qkv_linear_weight[projection_size:projection_size + kv_projection_size])) + grad_layernorm_output[i].add_(torch.matmul(dv[i], + qkv_linear_weight[projection_size + kv_projection_size:])) + + dk[i] = None + dv[i] = None + + for i in range(num_chunks): + dq_seq_len = dq[i].shape[1] + if ctx.pos_emb_cos is not None: + dq[i] = apply_rotary_pos_emb_backward(dq[i].to(dtype), + ctx.pos_emb_cos[:, dq_seq_len * i:dq_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dq_seq_len * i:dq_seq_len * (i + 1)]) + else: + dq[i] = dq[i].to(dtype) + dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(), gather_idx, scatter_idx, 0, spg) + + input_chunk = layernorm_output[:input_chunk_size].reshape(-1, layernorm_output.shape[-1]) + layernorm_output = layernorm_output[input_chunk_size:] + + dq[i] = dq[i].flatten(2).permute(1, 0, 2) + l, b = dq[i].shape[0], dq[i].shape[1] + grad_qkv_linear_weight[:projection_size].add_(torch.matmul(dq[i].reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_bias[:projection_size].add_(dq[i].sum(0).sum(0)) + + grad_layernorm_output[i].add_(torch.matmul(dq[i], qkv_linear_weight[:projection_size])) + + dq[i] = None + + return torch.cat( + grad_layernorm_output, + dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( + dtype), grad_qkv_linear_bias.to(dtype), None, None, None + + +class SequenceChunk: + + def __init__(self, chunk: torch.Tensor, device=None, is_in_use=False): + + self.chunk_shape = chunk.shape + self.chunk_dtype = chunk.dtype + self.device = chunk.device if device is None else device + + cpu_chunk = torch.empty(chunk.shape, dtype=chunk.dtype, device='cpu', pin_memory=True) + + if get_accelerator().on_accelerator(chunk): + cpu_chunk.copy_(chunk, non_blocking=True) + else: + cpu_chunk = chunk + + self.cpu_chunk = cpu_chunk + + self.gpu_chunk = chunk if is_in_use else None + + def load_to_gpu(self): + assert self.gpu_chunk is None + if self.gpu_chunk is not None: + pass + else: + gpu_chunk = torch.empty(self.chunk_shape, device=self.device, dtype=self.chunk_dtype) + gpu_chunk.copy_(self.cpu_chunk, non_blocking=True) + self.gpu_chunk = gpu_chunk + + def get_gpu_chunk(self): + assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device + return self.gpu_chunk + + def check_gpu_chunk(self, ): + assert (self.gpu_chunk is not None) and ( + self.gpu_chunk.device == self.device + ), f"gpu_chunk {self.gpu_chunk is not None} shound be on {self.device}, but it is now on {self.gpu_chunk.device}" + return True + + def offload(self): + assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device + del self.gpu_chunk + self.gpu_chunk = None + + def overwrite_to_cpu(self): + assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device + self.cpu_chunk.copy_(self.gpu_chunk, non_blocking=True) + + +class _FPDTGPUOffloadingAttentionImpl_(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, + layernorm_output, + attention_mask, + inference_params, + rotary_pos_emb, + spg, + scatter_idx, + gather_idx, + hidden_size, + projection_size, + hidden_size_per_attention_head, + kv_projection_size, + qkv_linear_weight, + qkv_linear_bias, + dropout, + num_chunks=8, + cpu_offloading=True): + + do_save = layernorm_output.requires_grad + + if rotary_pos_emb is not None: + pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) + ctx.pos_emb_cos = pos_emb_cos + ctx.pos_emb_sin = pos_emb_sin + else: + ctx.pos_emb_cos = None + ctx.pos_emb_sin = None + with torch.no_grad(): + per_gpu_seq_len = layernorm_output.shape[0] + chunk_size = per_gpu_seq_len // num_chunks + assert chunk_size * num_chunks == per_gpu_seq_len + assert attention_mask is None + ctx.num_chunks = num_chunks + ctx.cpu_offloading = cpu_offloading + ctx.spg = spg + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + ctx.chunk_size = chunk_size + device = get_accelerator().current_device_name() + ctx.device = device + ctx.dtype = layernorm_output.dtype + ctx.projection_size = projection_size + ctx.kv_projection_size = kv_projection_size + + global_q = [] + global_k = [] + global_v = [] + + ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) + + ctx.dropout_p = dropout + ctx.window_size = (-1, -1) + ctx.alibi_slopes = None + + batch_size = layernorm_output.shape[1] + + global_o = [] + global_lse = [] + + layernorm_output_cpu = [] + final_output = [] + + offload_stream = get_accelerator().Stream() + general_offload_stream = get_accelerator().Stream() + compute_stream = get_accelerator().default_stream() + + q_compute_chunk_idx = 0 + kv_compute_chunk_idx = 0 + for i in range(num_chunks): + + qkv_chunk = torch.matmul(layernorm_output[:chunk_size], + qkv_linear_weight.t()) + qkv_linear_bias # torch.Size([18126, 1, 12288]) + + with get_accelerator().stream(general_offload_stream): + layernorm_output_cpu.append(SequenceChunk(layernorm_output[:chunk_size])) + + layernorm_output = layernorm_output[chunk_size:] + + q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) + global_q_chunk_len = q_chunk.shape[1] + + k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) + + v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( + qkv_chunk.shape[0], qkv_chunk.shape[1], -1, + hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd + v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) + + dist.barrier() + + if ctx.pos_emb_cos is not None: + pos_emb_cos_chunk = pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] + pos_emb_sin_chunk = pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] + + q_chunk = apply_rotary_pos_emb(q_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) + k_chunk = apply_rotary_pos_emb(k_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + with get_accelerator().stream(offload_stream): + global_q.append(SequenceChunk(q_chunk, is_in_use=True)) + global_k.append(SequenceChunk(k_chunk, is_in_use=True)) + global_v.append(SequenceChunk(v_chunk, is_in_use=True)) + + del qkv_chunk + + cur_attn_output = None + cur_attn_lse = None + for k_i in range(len(global_k)): + causal_chunk = i == k_i + with get_accelerator().stream(compute_stream): + if flash_attn_version >= version.parse("2.6.0"): + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + softcap=0.0, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + else: + block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + ctx.dropout_p, + ctx.softmax_scale, + causal=causal_chunk, + window_size=ctx.window_size, + alibi_slopes=ctx.alibi_slopes, + return_softmax=False) + cur_attn_output, cur_attn_lse = update_out_and_lse(cur_attn_output, cur_attn_lse, block_out, + block_lse) + + can_offload_kv = True + if k_i != (len(global_k) - 1) or i != (num_chunks - 1): + if k_i != (len(global_k) - 1): + next_kv_compute_chunk_idx = k_i + 1 + else: + next_kv_compute_chunk_idx = 0 + + if next_kv_compute_chunk_idx == kv_compute_chunk_idx: + can_offload_kv = False + else: + if next_kv_compute_chunk_idx != (len(global_k) - 1): + with get_accelerator().stream(offload_stream): + global_k[next_kv_compute_chunk_idx].load_to_gpu() + global_v[next_kv_compute_chunk_idx].load_to_gpu() + + if i == num_chunks - 1 and k_i == num_chunks - 1: + with get_accelerator().stream(offload_stream): + global_q[0].load_to_gpu() + global_k[0].load_to_gpu() + global_v[0].load_to_gpu() + global_o[0].load_to_gpu() + global_lse[0].load_to_gpu() + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + if can_offload_kv: + global_k[kv_compute_chunk_idx].offload() + global_v[kv_compute_chunk_idx].offload() + kv_compute_chunk_idx = next_kv_compute_chunk_idx + + global_q[q_compute_chunk_idx].offload() + q_compute_chunk_idx += 1 + + all2all_output = single_all_to_all( + cur_attn_output.to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) + final_output.append(all2all_output) + with get_accelerator().stream(general_offload_stream): + global_o.append(SequenceChunk(cur_attn_output.to(ctx.dtype))) + global_lse.append(SequenceChunk(cur_attn_lse[:, :, :, 0].permute(0, 2, 1).contiguous())) + + compute_stream.wait_stream(general_offload_stream) + compute_stream.synchronize() + + final_output = torch.cat(final_output, dim=1) + + head_dim = final_output.shape[-1] + + if do_save: + ctx.layernorm_output = layernorm_output_cpu + ctx.global_q = global_q + ctx.global_k = global_k + ctx.global_v = global_v + ctx.attn_output = global_o + ctx.attn_lse = global_lse + ctx.head_dim = head_dim + ctx.batch_size = batch_size + + ctx.qkv_linear_weight = qkv_linear_weight + ctx.qkv_linear_bias = qkv_linear_bias + + return final_output + + @staticmethod + def backward(ctx, grad_output): + num_chunks = ctx.num_chunks + device = grad_output.device + dtype = ctx.dtype + spg = ctx.spg + scatter_idx = ctx.scatter_idx + gather_idx = ctx.gather_idx + softmax_scale = ctx.softmax_scale + dropout_p = ctx.dropout_p + window_size = ctx.window_size + alibi_slopes = ctx.alibi_slopes + + projection_size = ctx.projection_size + kv_projection_size = ctx.kv_projection_size + + layernorm_output = ctx.layernorm_output + + global_q = ctx.global_q + global_k = ctx.global_k + global_v = ctx.global_v + attn_output = ctx.attn_output + lse = ctx.attn_lse + + qkv_linear_weight = ctx.qkv_linear_weight + qkv_linear_bias = ctx.qkv_linear_bias + + offload_stream = get_accelerator().Stream() + general_offload_stream = get_accelerator().Stream() + compute_stream = get_accelerator().default_stream() + + chunk_size = grad_output.shape[1] // num_chunks + assert chunk_size == layernorm_output[0].cpu_chunk.shape[0] + + grad_layernorm_output = [ + torch.zeros(layernorm_output[0].chunk_shape, device=device, dtype=dtype) for _ in range(num_chunks) + ] + + grad_global_attn_output = [None for _ in range(num_chunks)] + + q_compute_chunk_idx = 0 + kv_compute_chunk_idx = 0 + last_q_accum_idx = 0 + + with get_accelerator().stream(general_offload_stream): + layernorm_output[0].load_to_gpu() + grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, + device=qkv_linear_weight.device, + dtype=torch.float) + grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, + device=qkv_linear_weight.device, + dtype=torch.float) + + grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), scatter_idx, + gather_idx, 0, spg) + get_accelerator().synchronize() + grad_output = grad_output[:, chunk_size:] + + with get_accelerator().stream(offload_stream): + grad_global_attn_output[0] = SequenceChunk(grad_global_attn_output_chunk, is_in_use=True) + dq = [ + SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device=device), is_in_use=True) + ] + [ + SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device='cpu', pin_memory=True), + device) for _ in range(num_chunks - 1) + ] + dk_accum = torch.zeros(global_k[0].chunk_shape, dtype=torch.float, device=device) + dv_accum = torch.zeros(global_v[0].chunk_shape, dtype=torch.float, device=device) + + for i in range(num_chunks): + for q_i in range(num_chunks): + no_computation = q_i < i + if no_computation: + continue + + causal_chunk = q_i == i + + dq_this = torch.zeros(global_q[0].chunk_shape, dtype=dtype, device=device) + dk_this = torch.zeros(global_k[0].chunk_shape, dtype=dtype, device=device) + dv_this = torch.zeros(global_v[0].chunk_shape, dtype=dtype, device=device) + + with get_accelerator().stream(compute_stream): + if flash_attn_version >= version.parse("2.6.0"): + _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + attn_output[q_compute_chunk_idx].get_gpu_chunk(), + lse[q_compute_chunk_idx].get_gpu_chunk(), + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + softcap=0.0, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + else: + _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), + global_q[q_compute_chunk_idx].get_gpu_chunk(), + global_k[kv_compute_chunk_idx].get_gpu_chunk(), + global_v[kv_compute_chunk_idx].get_gpu_chunk(), + attn_output[q_compute_chunk_idx].get_gpu_chunk(), + lse[q_compute_chunk_idx].get_gpu_chunk(), + dq_this, + dk_this, + dv_this, + dropout_p, + softmax_scale, + causal_chunk, + window_size, + alibi_slopes=alibi_slopes, + deterministic=False, + rng_state=None) + + if i != (len(global_k) - 1): + if q_i != (len(global_q) - 1): + next_q_compute_chunk_idx = q_i + 1 + else: + next_q_compute_chunk_idx = i + 1 + + can_offload_q = True + + if next_q_compute_chunk_idx == q_compute_chunk_idx: + can_offload_q = False + else: + with get_accelerator().stream(offload_stream): + if i > 0 or q_i > 0: + if can_offload_q and last_q_accum_idx != i: # the first q chunk calculate in the loop will be sent out, therefore we do not offload it + dq[last_q_accum_idx].offload() + dq[next_q_compute_chunk_idx].load_to_gpu() + global_q[next_q_compute_chunk_idx].load_to_gpu() + attn_output[next_q_compute_chunk_idx].load_to_gpu() + lse[next_q_compute_chunk_idx].load_to_gpu() + if grad_global_attn_output[next_q_compute_chunk_idx] is not None: + grad_global_attn_output[next_q_compute_chunk_idx].load_to_gpu() + + if grad_global_attn_output[next_q_compute_chunk_idx] is None: + grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), + scatter_idx, gather_idx, 0, spg) + dist.barrier() + grad_output = grad_output[:, chunk_size:] + grad_global_attn_output[next_q_compute_chunk_idx] = SequenceChunk( + grad_global_attn_output_chunk, is_in_use=True) + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + with get_accelerator().stream(compute_stream): + dq[q_compute_chunk_idx].check_gpu_chunk() + dq[q_compute_chunk_idx].gpu_chunk.add_(dq_this) + dk_accum.add_(dk_this) + dv_accum.add_(dv_this) + + offload_stream.wait_stream(compute_stream) + with get_accelerator().stream(offload_stream): + dq[q_compute_chunk_idx].overwrite_to_cpu() + + if can_offload_q: + global_q[q_compute_chunk_idx].offload() + attn_output[q_compute_chunk_idx].offload() + lse[q_compute_chunk_idx].offload() + grad_global_attn_output[q_compute_chunk_idx].offload() + + last_q_accum_idx = q_compute_chunk_idx + q_compute_chunk_idx = next_q_compute_chunk_idx + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + dk_seq_len = dk_accum.shape[1] + + if ctx.pos_emb_cos is not None: + dq_accum = apply_rotary_pos_emb_backward(dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype), + ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) + dk_accum = apply_rotary_pos_emb_backward(dk_accum.to(dtype), + ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], + ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) + else: + dq_accum = dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype) + dk_accum = dk_accum.to(dtype) + dv_accum = dv_accum.to(dtype) + + dq_accum = single_all_to_all(dq_accum.contiguous(), gather_idx, scatter_idx, 0, spg) + dk_accum = single_all_to_all(dk_accum.contiguous(), gather_idx, scatter_idx, 0, spg) + dv_accum = single_all_to_all(dv_accum.contiguous(), gather_idx, scatter_idx, 0, spg) + + general_offload_stream.synchronize() + compute_stream.wait_stream(general_offload_stream) + dist.barrier() + + with get_accelerator().stream(compute_stream): + input_chunk = layernorm_output[i].get_gpu_chunk().reshape(-1, layernorm_output[i].chunk_shape[-1]) + + dq_accum = dq_accum.flatten(2).permute(1, 0, 2) + dk_accum = dk_accum.flatten(2).permute(1, 0, 2) + dv_accum = dv_accum.flatten(2).permute(1, 0, 2) + + l, b = dk_accum.shape[0], dk_accum.shape[1] + + grad_qkv_linear_weight[:projection_size].add_( + torch.matmul(dq_accum.reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( + torch.matmul(dk_accum.reshape(l * b, -1).t(), input_chunk)) + grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( + torch.matmul(dv_accum.reshape(l * b, -1).t(), input_chunk)) + + grad_qkv_linear_bias[:projection_size].add_(dq_accum.sum(0).sum(0)) + grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk_accum.sum(0).sum(0)) + grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv_accum.sum(0).sum(0)) + + grad_layernorm_output[i].add_(torch.matmul(dq_accum, qkv_linear_weight[:projection_size])) + grad_layernorm_output[i].add_( + torch.matmul(dk_accum, qkv_linear_weight[projection_size:projection_size + kv_projection_size])) + grad_layernorm_output[i].add_( + torch.matmul(dv_accum, qkv_linear_weight[projection_size + kv_projection_size:])) + + del dq_accum, dk_accum, dv_accum + dk_accum = torch.zeros(global_k[i].chunk_shape, dtype=torch.float, device=device) + dv_accum = torch.zeros(global_v[i].chunk_shape, dtype=torch.float, device=device) + dq[kv_compute_chunk_idx].offload() + dq[kv_compute_chunk_idx] = None + + if i != (len(global_k) - 1): + next_kv_compute_chunk_idx = kv_compute_chunk_idx + 1 + with get_accelerator().stream(offload_stream): + global_k[next_kv_compute_chunk_idx].load_to_gpu() + global_v[next_kv_compute_chunk_idx].load_to_gpu() + + with get_accelerator().stream(general_offload_stream): + layernorm_output[next_kv_compute_chunk_idx].load_to_gpu() + + compute_stream.wait_stream(offload_stream) + compute_stream.synchronize() + + layernorm_output[kv_compute_chunk_idx].offload() + global_k[kv_compute_chunk_idx].offload() + global_v[kv_compute_chunk_idx].offload() + kv_compute_chunk_idx = next_kv_compute_chunk_idx + + return torch.cat( + grad_layernorm_output, + dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( + dtype), grad_qkv_linear_bias.to(dtype), None, None, None + + +class FPDT_Attention(torch.nn.Module): + + def __init__(self, + config, + first_weight, + first_bias, + second_weight, + second_bias, + sequence_process_group, + gather_idx: int = 0, + scatter_idx: int = 2, + return_bias=True, + chunk_size=65536, + enable_offloading=True) -> None: + + super(FPDT_Attention, self).__init__() + if _flash_attn_forward is None or _flash_attn_backward is None: + raise ImportError( + "DeepSpeed FPDT requires flash-attn 2.6.3. Please install it with `pip install flash-attn --no-build-isolation`." + ) + + self.spg = sequence_process_group + self.scatter_idx = scatter_idx + self.gather_idx = gather_idx + self.config = config + + self.projection_size = config.kv_channels * config.num_attention_heads + self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads + self.kv_projection_size = config.kv_channels * config.num_key_value_heads + self.hidden_size = config.hidden_size + + self.qkv_linear_weight = first_weight + self.qkv_linear_bias = first_bias + self.qkv_dense_weight = second_weight + self.qkv_dense_bias = second_bias + + self.reture_bias = return_bias + self.dropout = config.attention_dropout + + self.chunk_size = chunk_size + self.double_buffer = enable_offloading + + def forward(self, + layernorm_output, + attention_mask, + inference_params, + rotary_pos_emb, + cpu_offloading=True) -> Tensor: + self.num_chunks_attn = layernorm_output.shape[0] * dist.get_world_size(self.spg) // self.chunk_size + + if not cpu_offloading or self.num_chunks_attn == 1: + output = _FPDTGPUAttentionImpl_.apply(layernorm_output, attention_mask, inference_params, rotary_pos_emb, + self.spg, self.scatter_idx, self.gather_idx, self.hidden_size, + self.projection_size, self.hidden_size_per_attention_head, + self.kv_projection_size, self.qkv_linear_weight, + self.qkv_linear_bias, self.dropout, self.num_chunks_attn, + cpu_offloading) + else: + output = _FPDTGPUOffloadingAttentionImpl_.apply( + layernorm_output, attention_mask, inference_params, rotary_pos_emb, self.spg, self.scatter_idx, + self.gather_idx, self.hidden_size, self.projection_size, self.hidden_size_per_attention_head, + self.kv_projection_size, self.qkv_linear_weight, self.qkv_linear_bias, self.dropout, + self.num_chunks_attn, cpu_offloading) + + output = output.flatten(2).permute(1, 0, 2).contiguous() + + output = torch.matmul(output, self.qkv_dense_weight.t()) + if not self.reture_bias: + output += self.qkv_dense_bias + return output, self.qkv_dense_bias if self.reture_bias else None + + +@torch.jit.script +def bias_gelu(x): + return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) + + +@torch.jit.script +def bias_gelu_back(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out) + return ff * g + + +class FPDT_FFN(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, x, w1, b1, w2, b2, add_bias, chunk_size): + do_save = x.requires_grad + ctx.add_bias = add_bias + device = x.device + + with torch.no_grad(): + num_chunk = x.shape[0] // chunk_size + ctx.num_chunk = num_chunk + result = torch.empty(x.shape, device=device, dtype=x.dtype) + assert chunk_size * num_chunk == x.shape[0] + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + x_ = torch.matmul(x[st:ed], w1.t()) + b1 + x_ = bias_gelu(x_) + if add_bias: + result[st:ed] = torch.matmul(x_, w2.t()) + b2 + else: + result[st:ed] = torch.matmul(x_, w2.t()) + + del x_ + + if do_save: + ctx.device = device + ctx.dtype = x.dtype + ctx.save_for_backward(x, w1, b1, w2, b2) + ctx.grad_x_shape = x.shape + return result.to(x.dtype), b2 if not add_bias else None + + @staticmethod + def backward(ctx, grad_output, grad_bias): + x, w1, b1, w2, b2 = ctx.saved_tensors + device = ctx.device + dtype = ctx.dtype + add_bias = ctx.add_bias + + num_chunk = ctx.num_chunk + chunk_size = x.shape[0] // num_chunk + assert chunk_size * num_chunk == grad_output.shape[0] + + grad_w2 = torch.zeros(w2.shape, device=device, dtype=torch.float) + grad_b2 = torch.zeros(b2.shape, device=device, dtype=torch.float) + grad_w1 = torch.zeros(w1.shape, device=device, dtype=torch.float) + grad_b1 = torch.zeros(b1.shape, device=device, dtype=torch.float) + + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + x_chunk = x[st:ed] + + before_act = (torch.matmul(x_chunk, w1.t()) + b1) + before_act_2 = before_act**2 + tanh_out = torch.tanh(0.79788456 * before_act * (1 + 0.044715 * before_act_2)) + ff = 0.5 * before_act * ((1 - tanh_out * tanh_out) * + (0.79788456 + 0.1070322243 * before_act_2)) + 0.5 * (1 + tanh_out) + grad_w2.add_( + torch.matmul(grad_output[st:ed].reshape(-1, grad_output.shape[2]).t(), + (before_act * 0.5 * (1 + tanh_out)).reshape(-1, before_act.shape[2]))) + del before_act, before_act_2, tanh_out + + grad_inter = torch.matmul(grad_output[st:ed], w2) * ff + del ff + + grad_w1.add_(torch.matmul( + grad_inter.reshape(-1, grad_inter.shape[2]).t(), x_chunk.reshape(-1, x.shape[2]))) + grad_b1.add_(grad_inter.sum(0).sum(0)) + + x[st:ed].copy_(torch.matmul(grad_inter, w1)) + + del grad_inter + + if add_bias: + grad_b2.add_(grad_output[st:ed].sum(0).sum(0)) + + return x, grad_w1.to(dtype), grad_b1.to(dtype), grad_w2.to(dtype), grad_b2.to(dtype), None, None + + +class FPDT_LogitsLoss(torch.autograd.Function): + generate_vmap_rule = False + + @staticmethod + def forward(ctx: Any, lm_output, labels, logit_weights, rank, spg_size, spg, num_chunk): + labels = labels.t() + chunk_size = lm_output.shape[0] // num_chunk + assert chunk_size * num_chunk == lm_output.shape[0] + batch_size, local_seq_len = lm_output.shape[1], lm_output.shape[0] + loss = torch.empty((batch_size, local_seq_len), dtype=torch.float, device=lm_output.device) + + ctx.num_chunk = num_chunk + ctx.chunk_size = chunk_size + ctx.device = lm_output.device + ctx.dtype = lm_output.dtype + + ctx.rank = rank + ctx.local_seq_len = local_seq_len + with torch.no_grad(): + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + logits_chunk = torch.matmul(lm_output[st:ed], logit_weights.t()).float() + + vocab_size = logits_chunk.size(2) + # nll + softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) + loss_chunk = torch.nn.functional.nll_loss(softmax.log().reshape(-1, vocab_size).contiguous(), + labels[st:ed, :].reshape(-1).contiguous(), + reduction='none') + loss[:, st:ed] = loss_chunk.reshape(chunk_size, batch_size).t() + + del logits_chunk + ctx.save_for_backward(lm_output.to('cpu'), labels) + ctx.logit_weights = logit_weights + + seqlen = local_seq_len * spg_size + batch_size = loss.size(0) + loss = loss.t().contiguous() + loss_all = torch.empty(seqlen, batch_size, dtype=loss.dtype, device=loss.device).contiguous() + + dist.allgather_fn(loss_all, loss, group=spg) + + return loss_all + + @staticmethod + def backward(ctx, grad_output): + lm_output, labels = ctx.saved_tensors + logit_weights = ctx.logit_weights + device = ctx.device + dtype = ctx.dtype + num_chunk = ctx.num_chunk + chunk_size = ctx.chunk_size + + rank = ctx.rank + local_seq_len = ctx.local_seq_len + + grad_output = grad_output[rank * local_seq_len:(rank + 1) * local_seq_len] + grad_lm_output = [None for _ in range(num_chunk)] + grad_logit_weights = torch.zeros(logit_weights.shape, device=grad_output.device, dtype=torch.float) + for i in range(num_chunk): + st = i * chunk_size + ed = st + chunk_size + lm_output_chunk = lm_output[st:ed].to(device) + logits_chunk = torch.matmul(lm_output_chunk, logit_weights.t()).float() + + # nll + softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) + vocab_size = logits_chunk.size(2) + + grad_input = softmax + grad_2d = grad_input.reshape(-1, vocab_size).contiguous() + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=device) + + grad_2d[arange_1d, labels[st:ed, :].reshape(-1).contiguous()] -= 1 + grad_input.mul_(grad_output[:chunk_size, :].unsqueeze(dim=-1)) + grad_input = grad_input.to(dtype) + + grad_output = grad_output[chunk_size:].contiguous() + + grad_lm_output_chunk = torch.matmul(grad_input, logit_weights) + grad_lm_output[i] = grad_lm_output_chunk + + grad_logit_weights.add_( + torch.matmul( + grad_input.reshape(-1, grad_input.shape[2]).t(), + lm_output_chunk.reshape(-1, lm_output_chunk.shape[2]))) + + return torch.cat(grad_lm_output, dim=0).to(dtype), None, grad_logit_weights.to(dtype), None, None, None, None diff --git a/lib/python3.12/site-packages/deepspeed/sequence/layer.py b/lib/python3.12/site-packages/deepspeed/sequence/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..ecbe0d94120e2312d82ce65bf6c12fcd5dbd48d9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/sequence/layer.py @@ -0,0 +1,440 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +import torch + +from typing import Any, Tuple +from torch import Tensor +from torch.nn import Module + +from einops import rearrange + +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.tp_shard import get_shard_size_list, set_num_kv_heads, get_num_kv_heads +from deepspeed.utils import groups + + +def _generate_layout_params(scatter_idx, batch_dim_idx, seq_world_size, input): + """ + This function generates the parameters required for `permute` and `reshape` operations, + which are used to process data before and after `all2all` communication. + """ + if batch_dim_idx == 0: + if scatter_idx < 2: + bs, global_seq_len, num_local_head, head_dim = input.shape + pre_all2all_inp_shape = [bs, seq_world_size, global_seq_len // seq_world_size, num_local_head, head_dim] + pre_all2all_permute_idx = (1, 0, 2, 3, 4) + + post_all2all_permute_idx = (1, 2, 0, 3, 4) + post_all2all_res_shape = [bs, global_seq_len // seq_world_size, seq_world_size * num_local_head, head_dim] + else: + bs, local_seq_len, num_total_head, head_dim = input.shape + assert num_total_head % seq_world_size == 0, f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + pre_all2all_inp_shape = [bs, local_seq_len, seq_world_size, num_total_head // seq_world_size, head_dim] + pre_all2all_permute_idx = (2, 0, 1, 3, 4) + + post_all2all_permute_idx = (1, 0, 2, 3, 4) + post_all2all_res_shape = [bs, seq_world_size * local_seq_len, num_total_head // seq_world_size, head_dim] + else: + if scatter_idx < 2: + global_seq_len, bs, num_local_head, head_dim = input.shape + pre_all2all_inp_shape = [seq_world_size, global_seq_len // seq_world_size, bs, num_local_head, head_dim] + pre_all2all_permute_idx = None + + post_all2all_permute_idx = (1, 2, 0, 3, 4) + post_all2all_res_shape = [bs, seq_world_size * global_seq_len, num_local_head // seq_world_size, head_dim] + else: + local_seq_len, bs, num_total_head, head_dim = input.shape + assert num_total_head % seq_world_size == 0, f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + pre_all2all_inp_shape = [local_seq_len, bs, seq_world_size, num_total_head // seq_world_size, head_dim] + pre_all2all_permute_idx = (2, 0, 1, 3, 4) + post_all2all_permute_idx = None + post_all2all_res_shape = [local_seq_len * seq_world_size, bs, num_total_head // seq_world_size, head_dim] + + return pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape + + +def post_all2all(permute_idx, res_shape): + """ + Post-processing function for `all2all` communication. + """ + + def post_func(input): + if permute_idx is not None: + input = input.permute(permute_idx).contiguous() + output = input.reshape(res_shape).contiguous() + + return output + + return post_func + + +def pre_all2all_fun(permute_idx, inp_shape, input): + """ + Pre-processing function for `all2all` communication. + """ + input_t = input.reshape(inp_shape).contiguous() + if permute_idx is not None: + input_t = input_t.permute(permute_idx).contiguous() + return input_t + + +def _rotate_half(x): + """ + change sign so the last dimension becomes [-odd, +even] + """ + x = rearrange(x, '... (j d) -> ... j d', j=2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(t, freqs_cos, freqs_sin): + """ + input tensor t is of shape [seq_length, ..., dim] + rotary positional embeding tensor freqs is of shape [seq_length, ..., dim] + check https://kexue.fm/archives/8265 for detailed formulas + """ + rot_dim = freqs_cos.shape[-1] + # ideally t_pass is empty so rotary pos embedding is applied to all tensor t + t, t_pass = t[..., :rot_dim], t[..., rot_dim:] + + # first part is cosine component + # second part is sine component, need to change signs with _rotate_half method + t = (t * freqs_cos) + (_rotate_half(t) * freqs_sin) + + res = t if t_pass.shape[-1] == 0 else torch.cat((t, t_pass), dim=-1) + return res + + +def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group): + seq_world_size = dist.get_world_size(group) + inp_shape = list(input.shape) + assert batch_dim_idx in [0, 1], "batch_dim_idx must be either 0 or 1" + + if not (scatter_idx < 2): + input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size) + input = input.transpose(0, scatter_idx).contiguous() + local_heads = input_splits[groups._get_sequence_parallel_rank()] + output_splits = [local_heads] * seq_world_size + + output_buffer_shape = [seq_world_size * local_heads] + list(input.shape[1:]) + output = torch.empty(output_buffer_shape, device=input.device, dtype=input.dtype) + dist.all_to_all_single(output,input,output_split_sizes=output_splits,\ + input_split_sizes=input_splits,group=group) + ###[seq_ws*local_heads, ...] to [seq_ws, local_heads, ...] + output = output.view(seq_world_size, local_heads, *output.shape[1:]) + ###[seq_ws,local_heads,b,seq_len,...] to [seq_ws,seq_len,b,local_heads,...] + + ### batch_dim_idx=0 [seq_ws,local_heads,seq_len,b,...] to [b, seq_ws, seq_len, local_heads ...] + ### batch_dim_idx=1 [seq_ws,local_heads,b,seq_len,...] to [seq_ws,seq_len,b,local_heads,...] + if batch_dim_idx == 0: + order = [3, 0, 2, 1] + list(range(4, len(output.shape))) + output = output.permute(order).contiguous() + ###[b, seq_ws*local_seq_len, local_heads,...] + output = output.view(output.shape[0], inp_shape[gather_idx] * seq_world_size, + *output.shape[3:]).contiguous() + elif batch_dim_idx == 1: + output = output.transpose(1, 3).contiguous() + ###[seq_ws*local_seq_len, b, local_heads,...] + output = output.view(inp_shape[gather_idx] * seq_world_size, *output.shape[2:]).contiguous() + else: + # The compatibility handling of 4D and 3D tensors, standardizing to 3D. + input = input.reshape(input.shape[0], input.shape[1], -1) + + if batch_dim_idx == 0: #b,s,h + input = input.permute(1, 2, 0).contiguous() #s,h,b + elif batch_dim_idx == 1: #s,b,h + input = input.transpose(1, 2).contiguous() #s,h,b + seq_len, h, batch_size = input.shape + num_local_heads_list = get_shard_size_list(get_num_kv_heads(), seq_world_size) + local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()] + h_dim = h // local_heads + local_seq_len = seq_len // seq_world_size + + input = input.view(seq_len * h, batch_size) + local_seq_len_with_heads = int(input.shape[0] / seq_world_size) # dim size of local_seq_len*local_heads*hdim + input_splits = [local_seq_len_with_heads] * seq_world_size + coeff = local_seq_len_with_heads // local_heads #per head: dim size of local_seq_len*hdim + + #uneven seq_world_size coeff, total_heads/local_heads. + heads_scale_coeff = get_num_kv_heads() / local_heads + + output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list] + output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads) + total_h = int(inp_shape[gather_idx] * heads_scale_coeff) + output = torch.empty(output_buff_d1_size, input.shape[1], device=input.device, dtype=input.dtype) + dist.all_to_all_single(output,input,output_split_sizes=output_splits, \ + input_split_sizes=input_splits,group=group) + ################## + #suppose 7 heads divide into 4 ranks [2,2,2,1] + #chunk_num_heads_small=floor(7/4)=1 + #chunk_num_heads_large=ceil(7/4)=2 + #num_chunk_heads_large=len([2,2,2])=3, all2all_buffer_counts + #num_chunk_heads_small=len([1])=1, all2all_buffer_counts + #total_num_large_heads=sum([2,2,2])=7 + #total_num_small_heads=sum([1])=1 + + chunk_num_heads_small = get_num_kv_heads() // seq_world_size # even heads compatible + chunk_num_heads_large = chunk_num_heads_small + 1 + num_chunk_heads_large = get_num_kv_heads() % seq_world_size + num_chunk_heads_small = seq_world_size - num_chunk_heads_large + total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large + total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small + + heads_large_combine_size = coeff * total_num_large_heads + heads_small_combine_size = coeff * total_num_small_heads + heads_large_chunk, heads_small_chunk = output.split([heads_large_combine_size, heads_small_combine_size], + dim=0) + heads_large_chunk = heads_large_chunk.view(num_chunk_heads_large, local_seq_len, chunk_num_heads_large, h_dim, + batch_size) + heads_small_chunk = heads_small_chunk.view(num_chunk_heads_small, local_seq_len, chunk_num_heads_small, h_dim, + batch_size) + if batch_dim_idx == 0: + #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[batch,local_seq_len,all2all_buffer_counts*n_heads,dim] + order = [4, 1, 0, 2, 3] + heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(batch_size, local_seq_len, + total_num_large_heads, h_dim) + heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(batch_size, local_seq_len, + total_num_small_heads, h_dim) + elif batch_dim_idx == 1: + #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[local_seq_len,batch,all2all_buffer_counts*n_heads,dim] + order = [1, 4, 0, 2, 3] + heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(local_seq_len, batch_size, + total_num_large_heads, h_dim) + heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(local_seq_len, batch_size, + total_num_small_heads, h_dim) + + output = torch.cat([heads_large_chunk, heads_small_chunk], dim=2).contiguous() + + inp_shape[scatter_idx] = inp_shape[scatter_idx] // seq_world_size + output_shape= inp_shape[: gather_idx] + \ + [total_h,] + \ + inp_shape[gather_idx + 1:] + + output = output.view(output_shape) + + return output + + +def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, async_op=False, handle=None, type=None): + seq_world_size = dist.get_world_size(group) + # we only need num_heads once + num_heads = input.shape[2] + + if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2): + # Assuming here that the number of heads for q is consistent with kv + # If not, additional logic is required for cases like GQA + if get_num_kv_heads() is None: + assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})" + # set heads at first call by num_total_heads. + # then use ``get_num_kv_heads() is not None`` to re-entry uneven path. + set_num_kv_heads(num_heads) + assert async_op == False, "uneven head sp does not support async op" + return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group) + + pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = _generate_layout_params( + scatter_idx, batch_dim_idx, seq_world_size, input) + + input_t = pre_all2all_fun(pre_all2all_permute_idx, pre_all2all_inp_shape, input) + + post_all2all_fun = post_all2all(post_all2all_permute_idx, post_all2all_res_shape) + output = torch.empty_like(input_t) + work = dist.all_to_all_single(output, input_t, group=group, async_op=async_op) + + if async_op: + if type in ('dq', 'dk'): + handle[type + '_work'] = work + handle[type + '_grad'] = output + handle[type + '_post_all2all_func'] = post_all2all_fun + return output.view(post_all2all_res_shape) + + res = post_all2all_fun(output) + return res + + +class _DimZeroAllToAll(torch.autograd.Function): + """Differentiable All2All across dimension 0.""" + + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, input: Tensor) -> Tensor: + world_size = dist.get_world_size(group) + assert input.shape[0] == world_size, f"Dim 0 {input.shape[0]} is not world size" + + ctx.group = group + + output = torch.empty_like(input).contiguous() + # torch.distributed.nn.functional.all_to_all_single(output, input.contiguous(), group=group) + dist.all_to_all_single(output, input.contiguous(), group=group) + return output + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor]: + return (None, _DimZeroAllToAll.apply(ctx.group, *grad_output)) + + +class _SeqAllToAll(torch.autograd.Function): + + @staticmethod + def forward(ctx: Any, + group: dist.ProcessGroup, + input: Tensor, + scatter_idx: int, + gather_idx: int, + batch_dim_idx: int, + stream=None, + handle=None, + type=None, + is_fwd=True) -> Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + ctx.stream = stream + ctx.handle = handle + ctx.type = type + ctx.batch_dim_idx = batch_dim_idx + if ctx.handle is None: + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) + + else: + # overlap communication path + if not is_fwd and type == 'o': + assert ctx.stream != None + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) + get_accelerator().current_stream().wait_stream(ctx.stream) + # The computation of d o_weight can overlap with the communication of d o_input + + elif not is_fwd and type in ('q', 'k'): + # Achieve communication overlap by pipelining the matrix computation and communication of dq, dk, and dv + type = 'd' + type + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, True, handle, type) + + elif is_fwd and type in ('q', 'k'): + # Achieve communication overlap by pipelining the matrix computation and communication of q, k, and v + type = 'fwd_' + type + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False, handle, type) + + else: + res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) + + return res + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: + + return (None, + _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx, ctx.batch_dim_idx, + ctx.stream, ctx.handle, ctx.type, False), None, None, None, None, None, None, None) + + +class DistributedAttention(torch.nn.Module): + """Initialization. + + Arguments: + local_attention (Module): local attention with q,k,v + sequence_process_group (ProcessGroup): sequence parallel process group + scatter_idx (int): scatter_idx for all2all comm + gather_idx (int): gather_idx for all2all comm + """ + + def __init__( + self, + local_attention: Module, + sequence_process_group: dist.ProcessGroup, + scatter_idx: int = 2, + gather_idx: int = 0, + sp_stream=None, + ) -> None: + + super(DistributedAttention, self).__init__() + self.local_attn = local_attention + self.spg = sequence_process_group + self.scatter_idx = scatter_idx + self.gather_idx = gather_idx + self.sp_overlap_comm = False + self.overlap_handles = None + self.sp_stream = sp_stream + if sp_stream is not None: + self.overlap_handles = {} + self.sp_overlap_comm = True + self.default_stream = get_accelerator().default_stream() + + def layer_sync(self, layer): + if self.sp_overlap_comm and hasattr(layer, 'done_event'): + self.default_stream.wait_event(layer.done_event) + + def forward(self, + query: Tensor, + key: Tensor, + value: Tensor, + batch_dim_idx: int, + rotary_pos_emb=None, + *args: Any, + **kwargs) -> Tensor: + """ forward + + Arguments: + query (Tensor): query input to the layer + key (Tensor): key input to the layer + value (Tensor): value input to the layer + batch_dim_idx (int): indicating which dim is batch + args: other args + + Returns: + * output (Tensor): context output + """ + + # TODO Merge three alltoall calls into one + # TODO (Reza): change the api on the megatron-deepspeed side so that we only receive all data (q,k, and v) together! + #in shape : e.g., [s/p:h:] + + def bwd_hook(layer_type): + + def pre_hook_fun(grad): + type = 'd' + layer_type + self.overlap_handles[type + '_work'].wait() + self.sp_stream.wait_stream(self.default_stream) + all2all_output = self.overlap_handles[type + '_grad'] + grad = list(grad) + grad[0] = self.overlap_handles[type + '_post_all2all_func'](all2all_output) + grad = tuple(grad) + + return pre_hook_fun + + self.layer_sync(query) + query_layer = _SeqAllToAll.apply(self.spg, query, self.scatter_idx, self.gather_idx, batch_dim_idx, None, + self.overlap_handles, 'q') + self.layer_sync(key) + key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx, batch_dim_idx, None, + self.overlap_handles, 'k') + if self.sp_overlap_comm: + self.default_stream.wait_stream(self.sp_stream) + + value_layer = _SeqAllToAll.apply(self.spg, value, self.scatter_idx, self.gather_idx, batch_dim_idx, None, + self.overlap_handles, 'v') + + if self.sp_overlap_comm: + # Register a hook to synchronize dq and dk after the all-to-all + # operation when the gradient data is used. + # Place this logic after the q, k, v all-to-all operation to + # improve interpreter speed to + # call and launch of the forward all-to-all communication. + grad_fn_q = query.grad_fn.next_functions[0][0] + grad_fn_q.register_prehook(bwd_hook(layer_type='q')) + grad_fn_k = key.grad_fn.next_functions[0][0] + grad_fn_k.register_prehook(bwd_hook(layer_type='k')) + + #out shape : e.g., [s:h/p:] + if rotary_pos_emb is not None: + pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) + query_layer = apply_rotary_pos_emb(query_layer, pos_emb_cos, pos_emb_sin) + key_layer = apply_rotary_pos_emb(key_layer, pos_emb_cos, pos_emb_sin) + + context_layer = self.local_attn(query_layer, key_layer, value_layer, *args, **kwargs) + + output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx, batch_dim_idx, + self.sp_stream, self.overlap_handles, 'o') + + #out e.g., [s/p::h] + return output diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d1ce6f6dbdaf0ddabeda914ee3de430fa83aba3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/debug.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/debug.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..684bf4bbd9c0df91221fe604707d415f82e6caab Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/debug.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f56a251f258dde4b6663d20dfb787e843bd3ba9b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77d887ca019797b8ce48219fe834a2c03d84535d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/timer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/timer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3432fb165d90b8d1d9131139cce5ab69801402d2 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/timer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ef9dc89dcf022daf10c1e47fe1b9e7a98c3cc02 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/types.cpython-312.pyc differ