123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135 |
- '''Copyright The Microsoft DeepSpeed Team'''
- """
- Copyright (c) Microsoft Corporation
- Licensed under the MIT license.
- """
- 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
|