class_docstr: 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.
class_docstr_tok: ["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","."]
docstr: 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.
docstr_tok: ["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","."]
code: 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
  
imports: 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
label_desc: Use image node_id 17 for calling the Autotuner obj's underlying member method code with example usage: obj._generate_experiments(tuning_space, max_train_batch_size_per_gpu) and returns: exps, exps