wehub-resource-sync
78 lines
2.9 KiB
Python
78 lines
2.9 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from deepspeed.runtime.config_utils import get_scalar_param
|
|
from .constants import *
|
|
|
|
VALID_VALUES = {
|
|
CHECKPOINT_TAG_VALIDATION: CHECKPOINT_TAG_VALIDATION_MODES,
|
|
CHECKPOINT_WRITER_TYPE: CHECKPOINT_WRITER_TYPES,
|
|
CHECKPOINT_DATA_PARALLEL: CHECKPOINT_DATA_PARALLEL_UNITS
|
|
}
|
|
|
|
CHECKPOINT_DEFAULT_DICT = {
|
|
CHECKPOINT_TAG_VALIDATION: CHECKPOINT_TAG_VALIDATION_DEFAULT,
|
|
CHECKPOINT_SERIALIZATION: CHECKPOINT_SERIALIZATION_DEFAULT,
|
|
CHECKPOINT_WRITER: CHECKPOINT_WRITER_DEFAULT
|
|
}
|
|
|
|
|
|
def _validate_config_values(config_name, config_dict, valid_values):
|
|
for key, value in config_dict.items():
|
|
if value is None:
|
|
continue
|
|
if key in valid_values.keys():
|
|
assert value in valid_values[key], \
|
|
f"{config_name} contains invalid value {value} for {key}, expecting one of {valid_values[key]}"
|
|
|
|
|
|
def _make_upper_case(value):
|
|
return value if value is None else value.upper()
|
|
|
|
|
|
def get_checkpoint_writer_config(param_dict):
|
|
writer_dict = param_dict.get(CHECKPOINT_WRITER, None)
|
|
if writer_dict is None:
|
|
return CHECKPOINT_WRITER_DEFAULT
|
|
|
|
writer_config = {
|
|
CHECKPOINT_WRITER_TYPE:
|
|
_make_upper_case(get_scalar_param(writer_dict, CHECKPOINT_WRITER_TYPE, CHECKPOINT_WRITER_TYPE_DEFAULT)),
|
|
CHECKPOINT_IO_BUFFER_SIZE:
|
|
get_scalar_param(writer_dict, CHECKPOINT_IO_BUFFER_SIZE, CHECKPOINT_IO_BUFFER_SIZE_DEFAULT),
|
|
CHECKPOINT_IO_BUFFER_DOUBLE:
|
|
get_scalar_param(writer_dict, CHECKPOINT_IO_BUFFER_DOUBLE, CHECKPOINT_IO_BUFFER_DOUBLE_DEFAULT),
|
|
CHECKPOINT_IO_STATISTICS:
|
|
get_scalar_param(writer_dict, CHECKPOINT_IO_STATISTICS, CHECKPOINT_IO_STATISTICS_DEFAULT),
|
|
CHECKPOINT_DATA_PARALLEL:
|
|
_make_upper_case(get_scalar_param(writer_dict, CHECKPOINT_DATA_PARALLEL, CHECKPOINT_DATA_PARALLEL_DEFAULT)),
|
|
CHECKPOINT_WRITER_DECOUPLED:
|
|
get_scalar_param(writer_dict, CHECKPOINT_WRITER_DECOUPLED, CHECKPOINT_WRITER_DECOUPLED_DEFAULT),
|
|
CHECKPOINT_IO_MULTIPLIER:
|
|
get_scalar_param(writer_dict, CHECKPOINT_IO_MULTIPLIER, CHECKPOINT_IO_MULTIPLIER_DEFAULT),
|
|
}
|
|
_validate_config_values(CHECKPOINT_WRITER, writer_config, VALID_VALUES)
|
|
|
|
return writer_config
|
|
|
|
|
|
def get_checkpoint_config(param_dict):
|
|
checkpoint_dict = param_dict.get(CHECKPOINT, None)
|
|
if checkpoint_dict is None:
|
|
return CHECKPOINT_DEFAULT_DICT
|
|
|
|
checkpoint_config = {
|
|
CHECKPOINT_TAG_VALIDATION:
|
|
get_scalar_param(checkpoint_dict, CHECKPOINT_TAG_VALIDATION, CHECKPOINT_TAG_VALIDATION_DEFAULT).upper(),
|
|
CHECKPOINT_SERIALIZATION:
|
|
get_scalar_param(checkpoint_dict, CHECKPOINT_SERIALIZATION, CHECKPOINT_SERIALIZATION_DEFAULT),
|
|
CHECKPOINT_WRITER:
|
|
get_checkpoint_writer_config(checkpoint_dict)
|
|
}
|
|
|
|
_validate_config_values(CHECKPOINT, checkpoint_config, VALID_VALUES)
|
|
|
|
return checkpoint_config
|