Add files using upload-large-folder tool

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

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.")

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

lib/python3.12/site-packages/deepspeed/autotuning/config_templates/template_zero0.json

@@ -0,0 +1,5 @@
+{
+  "zero_optimization": {
+    "stage": 0
+  }
+}

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
+  }
+}

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
+  }
+}

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
+  }
+}

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"

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}")

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

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

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)

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

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

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

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) + " "

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 *

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'

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.'

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)

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

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)

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