86 lines
2.2 KiB
Python
86 lines
2.2 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
|
|
#########################################
|
|
# Validation modes
|
|
#########################################
|
|
class ValidationMode:
|
|
WARN = "WARN"
|
|
IGNORE = "IGNORE"
|
|
FAIL = "FAIL"
|
|
|
|
|
|
#########################################
|
|
# Checkpoint config params
|
|
#########################################
|
|
# "checkpoint": {tag_validation=["Ignore"|"Warn"|"Fail"]}
|
|
CHECKPOINT_FORMAT = '''
|
|
"checkpoint": {
|
|
"tag_validation": [Ignore|Warn|Fail],
|
|
"checkpoint_serialization": False,
|
|
"writer": {
|
|
"type": [mock|python|fast],
|
|
"decoupled": [True|False]
|
|
"io_buffer_size": 64e6,
|
|
"io_buffer_double": True,
|
|
"show_statistics": False,
|
|
"data_parallel": [replica|socket|machine],
|
|
"io_multiplier": 1,
|
|
}
|
|
}
|
|
'''
|
|
CHECKPOINT = "checkpoint"
|
|
CHECKPOINT_TAG_VALIDATION = "tag_validation"
|
|
CHECKPOINT_TAG_VALIDATION_DEFAULT = ValidationMode.WARN
|
|
CHECKPOINT_TAG_VALIDATION_MODES = [ValidationMode.WARN, ValidationMode.IGNORE, ValidationMode.FAIL]
|
|
|
|
CHECKPOINT_SERIALIZATION = "checkpoint_serialization"
|
|
CHECKPOINT_SERIALIZATION_DEFAULT = True
|
|
|
|
CHECKPOINT_WRITER = "writer"
|
|
CHECKPOINT_WRITER_DEFAULT = None
|
|
|
|
CHECKPOINT_WRITER_TYPE = "type"
|
|
|
|
|
|
class CheckpointWriterType:
|
|
MOCK = "MOCK"
|
|
PYTHON = "PYTHON"
|
|
FAST = "FAST"
|
|
|
|
|
|
CHECKPOINT_WRITER_TYPE_DEFAULT = CheckpointWriterType.FAST
|
|
CHECKPOINT_WRITER_TYPES = [CheckpointWriterType.MOCK, CheckpointWriterType.PYTHON, CheckpointWriterType.FAST]
|
|
|
|
CHECKPOINT_IO_BUFFER_SIZE = "io_buffer_size"
|
|
CHECKPOINT_IO_BUFFER_SIZE_DEFAULT = 64 * (1024**2)
|
|
|
|
CHECKPOINT_IO_BUFFER_DOUBLE = "io_buffer_double"
|
|
CHECKPOINT_IO_BUFFER_DOUBLE_DEFAULT = True
|
|
|
|
CHECKPOINT_IO_MULTIPLIER = "io_multiplier"
|
|
CHECKPOINT_IO_MULTIPLIER_DEFAULT = 1
|
|
|
|
CHECKPOINT_IO_STATISTICS = "show_statistics"
|
|
CHECKPOINT_IO_STATISTICS_DEFAULT = False
|
|
|
|
CHECKPOINT_DATA_PARALLEL = "data_parallel"
|
|
CHECKPOINT_DATA_PARALLEL_DEFAULT = None
|
|
|
|
|
|
class CheckpointDataParallel:
|
|
REPLICA = "REPLICA"
|
|
SOCKET = "SOCKET"
|
|
MACHINE = "MACHINE"
|
|
|
|
|
|
CHECKPOINT_DATA_PARALLEL_UNITS = [
|
|
CheckpointDataParallel.REPLICA, CheckpointDataParallel.SOCKET, CheckpointDataParallel.MACHINE
|
|
]
|
|
|
|
CHECKPOINT_WRITER_DECOUPLED = "decoupled"
|
|
CHECKPOINT_WRITER_DECOUPLED_DEFAULT = False
|