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 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
  
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 19 for calling the Autotuner obj's underlying member method code with example usage: obj.tune_space(tuning_space, prev_max_mbs, prev_best_mbs, prev_best_metric_val) and returns: max_micro_batch_size,  best_mbs,  best_metric_val, int,  int,  int, int,  int,  int, max_micro_batch_size,  fast_best_mbs,  fast_best_metric_val, int,  int,  int, max_micro_batch_size,  fast_best_mbs,  fast_best_metric_val