import cerberus
import torch
import onnxruntime as ort
from . import PropagateCastOpsStrategy
from .amp import loss_scaler
from .optim import lr_scheduler
[docs]class ORTTrainerOptions(object):
r"""Settings used by ONNX Runtime training backend
The parameters are hierarchically organized to facilitate configuration through semantic groups
that encompasses features, such as distributed training, etc.
Input validation is performed on the input dict during instantiation to ensure
that supported parameters and values are passed in. Invalid input results
in :py:obj:`ValueError` exception with details on it.
Args:
options (dict): contains all training options
_validate (bool, default is True): for internal use only
Supported schema for kwargs:
.. code-block:: python
schema = {
'batch' : {
'type' : 'dict',
'required': False,
'default' : {},
'schema' : {
'gradient_accumulation_steps' : {
'type' : 'integer',
'min' : 1,
'default' : 1
}
},
},
'device' : {
'type' : 'dict',
'required': False,
'default' : {},
'schema' : {
'id' : {
'type' : 'string',
'default' : 'cuda'
},
'mem_limit' : {
'type' : 'integer',
'min' : 0,
'default' : 0
}
}
},
'distributed': {
'type': 'dict',
'default': {},
'required': False,
'schema': {
'world_rank': {
'type': 'integer',
'min': 0,
'default': 0
},
'world_size': {
'type': 'integer',
'min': 1,
'default': 1
},
'local_rank': {
'type': 'integer',
'min': 0,
'default': 0
},
'data_parallel_size': {
'type': 'integer',
'min': 1,
'default': 1
},
'horizontal_parallel_size': {
'type': 'integer',
'min': 1,
'default': 1
},
'pipeline_parallel' : {
'type': 'dict',
'default': {},
'required': False,
'schema': {
'pipeline_parallel_size': {
'type': 'integer',
'min': 1,
'default': 1
},
'num_pipeline_micro_batches': {
'type': 'integer',
'min': 1,
'default': 1
},
'pipeline_cut_info_string': {
'type': 'string',
'default': ''
},
'sliced_schema': {
'type': 'dict',
'default': {},
'keysrules': {'type': 'string'},
'valuesrules': {
'type': 'list',
'schema': {'type': 'integer'}
}
},
'sliced_axes': {
'type': 'dict',
'default': {},
'keysrules': {'type': 'string'},
'valuesrules': {'type': 'integer'}
},
'sliced_tensor_names': {
'type': 'list',
'schema': {'type': 'string'},
'default': []
}
}
},
'allreduce_post_accumulation': {
'type': 'boolean',
'default': False
},
'deepspeed_zero_optimization': {
'type': 'dict',
'default': {},
'required': False,
'schema': {
'stage': {
'type': 'integer',
'min': 0,
'max': 1,
'default': 0
},
}
},
'enable_adasum': {
'type': 'boolean',
'default': False
}
}
},
'lr_scheduler' : {
'type' : 'optim.lr_scheduler',
'nullable' : True,
'default' : None
},
'mixed_precision' : {
'type' : 'dict',
'required': False,
'default' : {},
'schema' : {
'enabled' : {
'type' : 'boolean',
'default' : False
},
'loss_scaler' : {
'type' : 'amp.loss_scaler',
'nullable' : True,
'default' : None
}
}
}
},
'graph_transformer': {
'type': 'dict',
'required': False,
'default': {},
'schema': {
'attn_dropout_recompute': {
'type': 'boolean',
'default': False
},
'gelu_recompute': {
'type': 'boolean',
'default': False
},
'transformer_layer_recompute': {
'type': 'boolean',
'default': False
},
'number_recompute_layers': {
'type': 'integer',
'min': 0,
'default': 0
},
'propagate_cast_ops_config': {
'type': 'dict',
'required': False,
'default': {},
'schema': {
'propagate_cast_ops_strategy': {
'type': 'onnxruntime.training.PropagateCastOpsStrategy',
'default': PropagateCastOpsStrategy.FLOOD_FILL
},
'propagate_cast_ops_level': {
'type': 'integer',
'default': 1
},
'propagate_cast_ops_allow': {
'type': 'list',
'schema': {'type': 'string'},
'default': []
}
},
'allow_layer_norm_mod_precision': {
'type': 'boolean',
'default': False
}
}
},
'utils' : {
'type' : 'dict',
'required': False,
'default' : {},
'schema' : {
'frozen_weights' : {
'type' : 'list',
'default' : []
},
'grad_norm_clip' : {
'type' : 'boolean',
'default' : True
},
'memory_efficient_gradient' : {
'type' : 'boolean',
'default' : False
},
'run_symbolic_shape_infer' : {
'type' : 'boolean',
'default' : False
}
}
},
'debug' : {
'type' : 'dict',
'required': False,
'default' : {},
'schema' : {
'deterministic_compute' : {
'type' : 'boolean',
'default' : False
},
'check_model_export' : {
'type' : 'boolean',
'default' : False
},
'graph_save_paths' : {
'type' : 'dict',
'default': {},
'required': False,
'schema': {
'model_after_graph_transforms_path': {
'type': 'string',
'default': ''
},
'model_with_gradient_graph_path':{
'type': 'string',
'default': ''
},
'model_with_training_graph_path': {
'type': 'string',
'default': ''
},
'model_with_training_graph_after_optimization_path': {
'type': 'string',
'default': ''
},
}
},
}
},
'_internal_use' : {
'type' : 'dict',
'required': False,
'default' : {},
'schema' : {
'enable_internal_postprocess' : {
'type' : 'boolean',
'default' : True
},
'extra_postprocess' : {
'type' : 'callable',
'nullable' : True,
'default' : None
},
'onnx_opset_version': {
'type': 'integer',
'min' : 12,
'max' : 13,
'default': 12
},
'enable_onnx_contrib_ops' : {
'type' : 'boolean',
'default' : True
}
},
}
},
'_internal_use' : {
'type' : 'dict',
'required': False,
'default' : {},
'schema' : {
'enable_internal_postprocess' : {
'type' : 'boolean',
'default' : True
},
'extra_postprocess' : {
'type' : 'callable',
'nullable' : True,
'default' : None
},
'onnx_opset_version': {
'type': 'integer',
'min' : 12,
'max' :14,
'default': 14
},
'enable_onnx_contrib_ops' : {
'type' : 'boolean',
'default' : True
}
},
'provider_options':{
'type': 'dict',
'default': {},
'required': False,
'schema': {}
},
'session_options': {
'type': 'SessionOptions',
'nullable': True,
'default': None
},
}
Keyword arguments:
batch (dict):
batch related settings
batch.gradient_accumulation_steps (int, default is 1):
number of steps to accumulate before do collective gradient reduction
device (dict):
compute device related settings
device.id (string, default is 'cuda'):
device to run training
device.mem_limit (int):
maximum memory size (in bytes) used by device.id
distributed (dict):
distributed training options.
distributed.world_rank (int, default is 0):
rank ID used for data/horizontal parallelism
distributed.world_size (int, default is 1):
number of ranks participating in parallelism
distributed.data_parallel_size (int, default is 1):
number of ranks participating in data parallelism
distributed.horizontal_parallel_size (int, default is 1):
number of ranks participating in horizontal parallelism
distributed.pipeline_parallel (dict):
Options which are only useful to pipeline parallel.
distributed.pipeline_parallel.pipeline_parallel_size (int, default is 1):
number of ranks participating in pipeline parallelism
distributed.pipeline_parallel.num_pipeline_micro_batches (int, default is 1):
number of micro-batches. We divide input batch into micro-batches and run the graph.
distributed.pipeline_parallel.pipeline_cut_info_string (string, default is ''):
string of cutting ids for pipeline partition.
distributed.allreduce_post_accumulation (bool, default is False):
True enables overlap of AllReduce with computation, while False,
postpone AllReduce until all gradients are ready
distributed.deepspeed_zero_optimization:
DeepSpeed ZeRO options.
distributed.deepspeed_zero_optimization.stage (int, default is 0):
select which stage of DeepSpeed ZeRO to use. Stage 0 means disabled.
distributed.enable_adasum (bool, default is False):
enable `Adasum <https://arxiv.org/abs/2006.02924>`_
algorithm for AllReduce
lr_scheduler (optim._LRScheduler, default is None):
specifies learning rate scheduler
mixed_precision (dict):
mixed precision training options
mixed_precision.enabled (bool, default is False):
enable mixed precision (fp16)
mixed_precision.loss_scaler (amp.LossScaler, default is None):
specifies a loss scaler to be used for fp16. If not specified,
:py:class:`.DynamicLossScaler` is used with default values.
Users can also instantiate :py:class:`.DynamicLossScaler` and
override its parameters. Lastly, a completely new implementation
can be specified by extending :py:class:`.LossScaler` class from scratch
graph_transformer (dict):
graph transformer related configurations
graph_transformer.attn_dropout_recompute(bool, default False)
graph_transformer.gelu_recompute(bool, default False)
graph_transformer.transformer_layer_recompute(bool, default False)
graph_transformer.number_recompute_layers(bool, default False)
graph_transformer.propagate_cast_ops_config (dict):
graph_transformer.propagate_cast_ops_config.strategy(PropagateCastOpsStrategy, default FLOOD_FILL)
Specify the choice of the cast propagation optimization strategy, either, NONE, INSERT_AND_REDUCE or FLOOD_FILL.
NONE strategy does not perform any cast propagation transformation on the graph, although other optimizations
locally change cast operations, for example, in order to fuse Transpose and MatMul nodes, the TransposeMatMulFunsion optimization could
interchange Transpose and Cast if the Cast node exists between Transpose and MatMul.
INSERT_AND_REDUCE strategy inserts and reduces cast operations around the nodes with allowed opcodes.
FLOOD_FILL strategy expands float16 regions in the graph using the allowed opcodes, and unlike
INSERT_AND_REDUCE does not touch opcodes outside expanded float16 region.
graph_transformer.propagate_cast_ops_config.level(integer, default 1)
Optimize by moving Cast operations if propagate_cast_ops_level is non-negative.
Use predetermined list of opcodes considered safe to move before/after cast operation
if propagate_cast_ops_level is positive and use propagate_cast_ops_allow otherwise.
graph_transformer.propagate_cast_ops_config.allow(list of str, [])
List of opcodes to be considered safe to move before/after cast operation if propagate_cast_ops_level is zero.
attn_dropout_recompute (bool, default is False):
enable recomputing attention dropout to save memory
gelu_recompute (bool, default is False):
enable recomputing Gelu activation output to save memory
transformer_layer_recompute (bool, default is False):
enable recomputing transformer layerwise to save memory
number_recompute_layers (int, default is 0)
number of layers to apply transformer_layer_recompute, by default system will
apply recompute to all the layers, except for the last one
utils (dict):
miscellaneous options
utils.frozen_weights (list of str, []):
list of model parameter names to skip training (weights don't change)
utils.grad_norm_clip (bool, default is True):
enables gradient norm clipping for 'AdamOptimizer' and 'LambOptimizer'
utils.memory_efficient_gradient (bool, default is False):
enables use of memory aware gradient builder.
utils.run_symbolic_shape_infer (bool, default is False):
runs symbolic shape inference on the model
debug (dict):
debug options
debug.deterministic_compute (bool, default is False)
forces compute to be deterministic accross runs
debug.check_model_export (bool, default is False)
compares PyTorch model outputs with ONNX model outputs in inference before the first
train step to ensure successful model export
debug.graph_save_paths (dict):
paths used for dumping ONNX graphs for debugging purposes
debug.graph_save_paths.model_after_graph_transforms_path (str, default is "")
path to export the ONNX graph after training-related graph transforms have been applied.
No output when it is empty.
debug.graph_save_paths.model_with_gradient_graph_path (str, default is "")
path to export the ONNX graph with the gradient graph added. No output when it is empty.
debug.graph_save_paths.model_with_training_graph_path (str, default is "")
path to export the training ONNX graph with forward, gradient and optimizer nodes.
No output when it is empty.
debug.graph_save_paths.model_with_training_graph_after_optimization_path (str, default is "")
outputs the optimized training graph to the path if nonempty.
_internal_use (dict):
internal options, possibly undocumented, that might be removed without notice
_internal_use.enable_internal_postprocess (bool, default is True):
enable internal internal post processing of the ONNX model
_internal_use.extra_postprocess (callable, default is None)
a functor to postprocess the ONNX model and return a new ONNX model.
It does not override :py:attr:`._internal_use.enable_internal_postprocess`, but complement it
_internal_use.onnx_opset_version (int, default is 14):
ONNX opset version used during model exporting.
_internal_use.enable_onnx_contrib_ops (bool, default is True)
enable PyTorch to export nodes as contrib ops in ONNX.
This flag may be removed anytime in the future.
session_options (onnxruntime.SessionOptions):
The SessionOptions instance that TrainingSession will use.
provider_options (dict):
The provider_options for customized execution providers. it is dict map from EP name to
a key-value pairs, like {'EP1' : {'key1' : 'val1'}, ....}
Example:
.. code-block:: python
opts = ORTTrainerOptions({
'batch' : {
'gradient_accumulation_steps' : 128
},
'device' : {
'id' : 'cuda:0',
'mem_limit' : 2*1024*1024*1024,
},
'lr_scheduler' : optim.lr_scheduler.LinearWarmupLRScheduler(),
'mixed_precision' : {
'enabled': True,
'loss_scaler': amp.LossScaler(loss_scale=float(1 << 16))
}
})
fp16_enabled = opts.mixed_precision.enabled
"""
def __init__(self, options={}):
# Keep a copy of original input for debug
self._original_opts = dict(options)
# Used for logging purposes
self._main_class_name = self.__class__.__name__
# Validates user input
self._validated_opts = dict(self._original_opts)
validator = ORTTrainerOptionsValidator(_ORTTRAINER_OPTIONS_SCHEMA)
self._validated_opts = validator.validated(self._validated_opts)
if self._validated_opts is None:
raise ValueError(f"Invalid options: {validator.errors}")
# Convert dict in object
for k, v in self._validated_opts.items():
setattr(self, k, self._wrap(v))
def __repr__(self):
return "{%s}" % str(
", ".join(
"'%s': %s" % (k, repr(v))
for (k, v) in self.__dict__.items()
if k not in ["_original_opts", "_validated_opts", "_main_class_name"]
)
)
def _wrap(self, v):
if isinstance(v, (tuple, list, set, frozenset)):
return type(v)([self._wrap(i) for i in v])
else:
return _ORTTrainerOptionsInternal(self._main_class_name, v) if isinstance(v, dict) else v
class _ORTTrainerOptionsInternal(ORTTrainerOptions):
r"""Internal class used by ONNX Runtime training backend for input validation
NOTE: Users MUST NOT use this class in any way!
"""
def __init__(self, main_class_name, options):
# Used for logging purposes
self._main_class_name = main_class_name
# We don't call super().__init__(options) here but still called it "_validated_opts"
# instead of "_original_opts" because it has been validated in the top-level
# ORTTrainerOptions's constructor.
self._validated_opts = dict(options)
# Convert dict in object
for k, v in dict(options).items():
setattr(self, k, self._wrap(v))
class ORTTrainerOptionsValidator(cerberus.Validator):
_LR_SCHEDULER = cerberus.TypeDefinition("lr_scheduler", (lr_scheduler._LRScheduler,), ())
_LOSS_SCALER = cerberus.TypeDefinition("loss_scaler", (loss_scaler.LossScaler,), ())
_SESSION_OPTIONS = cerberus.TypeDefinition("session_options", (ort.SessionOptions,), ())
_PROPAGATE_CAST_OPS_STRATEGY = cerberus.TypeDefinition(
"propagate_cast_ops_strategy", (PropagateCastOpsStrategy,), ()
)
types_mapping = cerberus.Validator.types_mapping.copy()
types_mapping["lr_scheduler"] = _LR_SCHEDULER
types_mapping["loss_scaler"] = _LOSS_SCALER
types_mapping["session_options"] = _SESSION_OPTIONS
types_mapping["propagate_cast_ops_strategy"] = _PROPAGATE_CAST_OPS_STRATEGY
def _check_is_callable(field, value, error):
result = False
try:
# Python 3
result = value is None or callable(value)
except:
# Python 3 but < 3.2
if hasattr(value, "__call__"):
result = True
if not result:
error(field, "Must be callable or None")
_ORTTRAINER_OPTIONS_SCHEMA = {
"batch": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {"gradient_accumulation_steps": {"type": "integer", "min": 1, "default": 1}},
},
"device": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"id": {"type": "string", "default": "cuda"},
"mem_limit": {"type": "integer", "min": 0, "default": 0},
},
},
"distributed": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"world_rank": {"type": "integer", "min": 0, "default": 0},
"world_size": {"type": "integer", "min": 1, "default": 1},
"local_rank": {"type": "integer", "min": 0, "default": 0},
"data_parallel_size": {"type": "integer", "min": 1, "default": 1},
"horizontal_parallel_size": {"type": "integer", "min": 1, "default": 1},
"pipeline_parallel": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"pipeline_parallel_size": {"type": "integer", "min": 1, "default": 1},
"num_pipeline_micro_batches": {"type": "integer", "min": 1, "default": 1},
"pipeline_cut_info_string": {"type": "string", "default": ""},
"sliced_schema": {
"type": "dict",
"default_setter": lambda _: {},
"keysrules": {"type": "string"},
"valuesrules": {"type": "list", "schema": {"type": "integer"}},
},
"sliced_axes": {
"type": "dict",
"default_setter": lambda _: {},
"keysrules": {"type": "string"},
"valuesrules": {"type": "integer"},
},
"sliced_tensor_names": {"type": "list", "schema": {"type": "string"}, "default": []},
},
},
"allreduce_post_accumulation": {"type": "boolean", "default": False},
"deepspeed_zero_optimization": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"stage": {"type": "integer", "min": 0, "max": 1, "default": 0},
},
},
"enable_adasum": {"type": "boolean", "default": False},
},
},
"lr_scheduler": {"type": "lr_scheduler", "nullable": True, "default": None},
"mixed_precision": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"enabled": {"type": "boolean", "default": False},
"loss_scaler": {"type": "loss_scaler", "nullable": True, "default": None},
},
},
"graph_transformer": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"attn_dropout_recompute": {"type": "boolean", "default": False},
"gelu_recompute": {"type": "boolean", "default": False},
"transformer_layer_recompute": {"type": "boolean", "default": False},
"number_recompute_layers": {"type": "integer", "min": 0, "default": 0},
"propagate_cast_ops_config": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"strategy": {
"type": "propagate_cast_ops_strategy",
"nullable": True,
"default": PropagateCastOpsStrategy.FLOOD_FILL,
},
"level": {"type": "integer", "min": -1, "default": 1},
"allow": {"type": "list", "schema": {"type": "string"}, "default": []},
},
},
},
},
"utils": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"frozen_weights": {"type": "list", "default": []},
"grad_norm_clip": {"type": "boolean", "default": True},
"memory_efficient_gradient": {"type": "boolean", "default": False},
"run_symbolic_shape_infer": {"type": "boolean", "default": False},
},
},
"debug": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"deterministic_compute": {"type": "boolean", "default": False},
"check_model_export": {"type": "boolean", "default": False},
"graph_save_paths": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"model_after_graph_transforms_path": {"type": "string", "default": ""},
"model_with_gradient_graph_path": {"type": "string", "default": ""},
"model_with_training_graph_path": {"type": "string", "default": ""},
"model_with_training_graph_after_optimization_path": {"type": "string", "default": ""},
},
},
},
},
"_internal_use": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"schema": {
"enable_internal_postprocess": {"type": "boolean", "default": True},
"extra_postprocess": {"check_with": _check_is_callable, "nullable": True, "default": None},
"onnx_opset_version": {"type": "integer", "min": 12, "max": 14, "default": 14},
"enable_onnx_contrib_ops": {"type": "boolean", "default": True},
},
},
"provider_options": {
"type": "dict",
"default_setter": lambda _: {},
"required": False,
"allow_unknown": True,
"schema": {},
},
"session_options": {"type": "session_options", "nullable": True, "default": None},
}