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","."]
code: 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_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
              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_sizs 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:
              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
  
imports: import shutil
  import subprocess
  import torch
  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 *
label_desc: Use image node_id 26 for calling the Autotuner obj's underlying member method code with example usage: obj.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) and returns: tuning_micro_batch_sizes, tuning_micro_batch_sizes