Skip to content

Refactor/Pydantify autotuning config #2769

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mrwyattii wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Refactor/Pydantify autotuning config #2769

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4b145f0
initial refactor of autotuning config
mrwyattii Jan 31, 2023
f8ac247
add back previously remove constants
mrwyattii Jan 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions deepspeed/autotuning/autotuner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ def __init__(self, args, active_resources):
if not os.path.exists(self.results_dir):
try:
os.makedirs(self.results_dir, exist_ok=True)
logger.info(f"Created autotuning resutls directory: {self.exps_dir}")
logger.info(f"Created autotuning results directory: {self.results_dir}")
except:
logger.error(
f"Failed to create {self.results_dir}, please check `results_dir` in the autotuning config file is accessible by all the nodes in the job."
Expand Down Expand Up @@ -645,9 +645,9 @@ def tune_space(self,
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:
if self.autotuning_config.tuner_type == AutotuningTunerEnum.model_based:
t = ModelBasedTuner(exps, self.rm, self.metric(), tuning_space)
elif self.autotuning_config.tuner_type == AUTOTUNING_TUNER_RANDOM:
elif self.autotuning_config.tuner_type == AutotuningTunerEnum.random:
t = RandomTuner(exps, self.rm, self.metric())
else:
t = GridSearchTuner(exps, self.rm, self.metric())
Expand Down Expand Up @@ -710,7 +710,7 @@ def model_info_profile_run(self):
"""
logger.info("Starting model info profile run.")
model_info = self.autotuning_config.model_info
if model_info and MODEL_INFO_NUM_PARAMS in model_info:
if model_info.num_params != None:
return model_info

ds_config = copy.deepcopy(self.user_config)
Expand Down
205 changes: 76 additions & 129 deletions deepspeed/autotuning/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,132 +3,79 @@
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
from typing import Dict
from pydantic import Field, validator, root_validator
from enum import Enum
from deepspeed.runtime.config_utils import DeepSpeedConfigModel


def get_autotuning_config(param_dict):
return DeepSpeedAutotuningConfig(**param_dict.get("autotuning", {}))


class AutotuningMetricEnum(str, Enum):
latency = "latency"
throughput = "throughput"
flops = "flops"
forward = "forward"
steps = "steps"


class AutotuningTunerEnum(str, Enum):
gridsearch = "gridsearch"
random = "random"
model_based = "model_based"


class ModelInfoConfig(DeepSpeedConfigModel):
profile: bool = False
num_params: int = Field(None, ge=0)
hidden_size: int = Field(None, ge=0)
num_layers: int = Field(None, ge=0)


class DeepSpeedAutotuningConfig(DeepSpeedConfigModel):
enabled: bool = False
fast: bool = True
results_dir: str = "autotuning_results"
exps_dir: str = "autotuning_exps"
overwrite: bool = True
start_profile_step: int = Field(3, ge=0)
end_profile_step: int = Field(5, ge=0)
metric: AutotuningMetricEnum = "throughput"
metric_path: str = None
tuner_type: AutotuningTunerEnum = "gridsearch"
tuner_num_trials: int = Field(50, ge=0)
tuner_early_stopping: int = Field(5, ge=0)
arg_mappings: Dict[str, str] = None
model_info: ModelInfoConfig = {}
model_info_path: str = None
mp_size: int = Field(1, ge=1)
max_train_batch_size: int = Field(None, ge=1)
min_train_batch_size: int = Field(1, ge=1)
max_train_micro_batch_size_per_gpu: int = Field(1024, ge=1)
min_train_micro_batch_size_per_gpu: int = Field(1, ge=1)
num_tuning_micro_batch_sizes: int = Field(3, ge=1)

@validator("results_dir", "exps_dir")
def assert_non_empty_str(cls, field_value, values):
assert field_value != "", "field cannot by empty"
return field_value

@root_validator
def check_profile_start_end(cls, values):
start_step = values.get("start_profile_step")
end_step = values.get("end_profile_step")
assert start_step <= end_step, f"start_profiling_step ({start_step}) cannot be greater than end_profiling_step ({end_step})"
return values

@root_validator
def check_min_max_batch_sizes(cls, values):
max_batch = values.get("max_train_batch_size")
min_batch = values.get("min_train_batch_size")
assert min_batch <= max_batch, f"min_train_batch_size ({min_batch}) cannot be greater than max_train_batch_size ({max_batch})"

max_micro_batch = values.get("max_train_micro_batch_size")
min_micro_batch = values.get("min_train_micro_batch_size")
assert min_micro_batch <= max_micro_batch, f"min_train_micro_batch_size ({min_micro_batch}) cannot be greater than max_train_micro_batch_size ({max_micro_batch})"
return values
100 changes: 0 additions & 100 deletions deepspeed/autotuning/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,109 +29,9 @@
#########################################
# autotuner configuration constants
#########################################
# Autotuner. By default, this feature is not enabled.
# Users can configure in ds_config.json as below example:
AUTOTUNING_FORMAT = """
autotuner should be enabled as:
"session_params": {
"autotuning": {
"enabled": true,
"start_step": 5,
"end_step": 15
}
}
"""

AUTOTUNING = "autotuning"

AUTOTUNING_ENABLED = "enabled"
AUTOTUNING_ENABLED_DEFAULT = False

AUTOTUNING_FAST = "fast"
AUTOTUNING_FAST_DEFAULT = True

AUTOTUNING_RESULTS_DIR = "results_dir"
AUTOTUNING_RESULTS_DIR_DEFAULT = "autotuning_results"

AUTOTUNING_EXPS_DIR = "exps_dir"
AUTOTUNING_EXPS_DIR_DEFAULT = "autotuning_exps"

AUTOTUNING_OVERWRITE = "overwrite"
AUTOTUNING_OVERWRITE_DEFAULT = True

AUTOTUNING_START_PROFILE_STEP = "start_profile_step"
AUTOTUNING_START_PROFILE_STEP_DEFAULT = 3

AUTOTUNING_END_PROFILE_STEP = "end_profile_step"
AUTOTUNING_END_PROFILE_STEP_DEFAULT = 5
AUTOTUNING_METRIC_PATH = "metric_path"
AUTOTUNING_METRIC_PATH_DEFAULT = None

AUTOTUNING_TUNER_TYPE = "tuner_type"
AUTOTUNING_TUNER_GRIDSEARCH = "gridsearch"
AUTOTUNING_TUNER_RANDOM = "random"
AUTOTUNING_TUNER_MODELBASED = "model_based"
AUTOTUNING_TUNER_TYPE_DEFAULT = AUTOTUNING_TUNER_GRIDSEARCH
AUTOTUNING_TUNER_EARLY_STOPPING = "tuner_early_stopping"
AUTOTUNING_TUNER_EARLY_STOPPING_DEFAULT = 5
AUTOTUNING_TUNER_NUM_TRIALS = "tuner_num_trials"
AUTOTUNING_TUNER_NUM_TRIALS_DEFAULT = 50

AUTOTUNING_ARG_MAPPINGS = "arg_mappings"
AUTOTUNING_ARG_MAPPINGS_DEFAULT = None

AUTOTUNING_MAX_TRAIN_BATCH_SIZE = "max_train_batch_size"
AUTOTUNING_MAX_TRAIN_BATCH_SIZE_DEFAULT = None
AUTOTUNING_MIN_TRAIN_BATCH_SIZE = "min_train_batch_size"
AUTOTUNING_MIN_TRAIN_BATCH_SIZE_DEFAULT = 1
AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU = "max_train_micro_batch_size_per_gpu"
AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT = 1024
AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU = "min_train_micro_batch_size_per_gpu"
AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT = 1
AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES = "num_tuning_micro_batch_sizes"
AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES_DEFAULT = 3

AUTOTUNING_MP_SIZE = "mp_size"
AUTOTUNING_MP_SIZE_DEFAULT = 1

AUTOTUNING_METRIC = "metric"
AUTOTUNING_METRIC_LATENCY = "latency"
AUTOTUNING_METRIC_THROUGHPUT = "throughput"
AUTOTUNING_METRIC_FLOPS = "flops"
AUTOTUNING_METRIC_FORWARD = "forward"
AUTOTUNING_METRIC_BACKWRAD = "flops"
AUTOTUNING_METRIC_STEPS = "step"
AUTOTUNING_METRIC_DEFAULT = AUTOTUNING_METRIC_THROUGHPUT

#########################################
# MODEL INFO
#########################################
AUTOTUNING_MODEL_INFO_PATH = "model_info_path"
AUTOTUNING_MODEL_INFO_PATH_DEFAULT = None

MODEL_INFO_FORMAT = '''
"model_info": {
"num_params": 1000000000,
"hidden_size": 10,
"num_layers": 12,
}
'''
MODEL_INFO = "model_info"
MODEL_INFO_PROFILE = "profile"
MODEL_INFO_PROFILE_DEFAULT = False
MODEL_INFO_NUM_PARAMS = "num_params"
MODEL_INFO_NUM_PARAMS_DEFAULT = None
MODEL_INFO_HIDDEN_SIZE = "hideen_size"
MODEL_INFO_HIDDEN_SIZE_DEFAULT = None
MODEL_INFO_NUM_LAYERS = "num_layers"
MODEL_INFO_NUM_LAYERS_DEFAULT = None

MODEL_INFO_KEY_DEFAULT_DICT = {
MODEL_INFO_PROFILE: MODEL_INFO_PROFILE_DEFAULT,
MODEL_INFO_NUM_PARAMS: MODEL_INFO_NUM_PARAMS_DEFAULT,
MODEL_INFO_HIDDEN_SIZE: MODEL_INFO_HIDDEN_SIZE_DEFAULT,
MODEL_INFO_NUM_LAYERS: MODEL_INFO_NUM_LAYERS_DEFAULT
}

#########################################
# autotunner search space constants
Expand Down
4 changes: 2 additions & 2 deletions deepspeed/runtime/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@
)

from ..profiling.config import DeepSpeedFlopsProfilerConfig
from ..autotuning.config import DeepSpeedAutotuningConfig
from ..autotuning.config import get_autotuning_config
from ..nebula.config import DeepSpeedNebulaConfig

from ..compression.config import get_compression_config, get_quantize_enabled
Expand Down Expand Up @@ -864,7 +864,7 @@ def _initialize_params(self, param_dict):
self.wall_clock_breakdown = (get_wall_clock_breakdown(param_dict)
| self.flops_profiler_config.enabled)
self.memory_breakdown = get_memory_breakdown(param_dict)
self.autotuning_config = DeepSpeedAutotuningConfig(param_dict)
self.autotuning_config = get_autotuning_config(param_dict)

(
self.eigenvalue_enabled,
Expand Down
4 changes: 1 addition & 3 deletions deepspeed/runtime/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -662,9 +662,7 @@ def autotuning_metric(self):

def autotuning_profile_model_info(self):
return self.autotuning_enabled(
) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.get(
"profile",
False)
) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.profile

def sparse_gradients_enabled(self):
return self._config.sparse_gradients_enabled
Expand Down