|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
|
