mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 15:33:51 +08:00
100 lines
4.5 KiB
Python
100 lines
4.5 KiB
Python
'''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
|