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: Tunes Zero stages, micro batch size per GPU, and other Zero configurations. Performance metrics of different tuning spaces are recorded in self.records.
          
docstr_tok: ["Tunes","Zero","stages",",","micro","batch","size","per","GPU",",","and","other","Zero","configurations",".","Performance","metrics","of","different","tuning","spaces","are","recorded","in","self.records","."]
code: 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."""
      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."
      )
  
      # TODO: FIX THIS
      stage = self.user_config.get(ZERO_OPTIMIZATION, {}).get(
          ZERO_OPTIMIZATION_STAGE, "all"
      )
      stage = "all"
      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
      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
      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
      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"
              )
              _, _, _ = self.tune_space(
                  DEFAULT_TUNING_SPACE_ZERO_3,
                  prev_max_mbs=max_mbs,
                  prev_best_mbs=mbs,
                  prev_best_metric_val=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
  
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 18 for calling the Autotuner obj's underlying member method code with example usage: obj.tune() without return types