diff --git a/orttraining/orttraining/python/training/orttrainer.py b/orttraining/orttraining/python/training/orttrainer.py index f7b3e5e68b..3965c32523 100644 --- a/orttraining/orttraining/python/training/orttrainer.py +++ b/orttraining/orttraining/python/training/orttrainer.py @@ -95,7 +95,6 @@ class ORTTrainer(object): Inputs to the combined PyTorch model are concatenation of the :py:attr:`model`'s input and :py:attr:`loss_fn`'s label input. Outputs of the combined PyTorch model are concatenation of :py:attr:`loss_fn`'s loss output and :py:attr:`model`'s outputs. options (ORTTrainerOptions, default is None): options for additional features. - Example: .. code-block:: python @@ -121,7 +120,9 @@ class ORTTrainer(object): ort_trainer = ORTTrainer(model, model_desc, optim_config, loss_fn) """ - def __init__(self, model, model_desc, optim_config, loss_fn=None, options=None): + def __init__(self, model, model_desc, optim_config, + loss_fn=None, + options=None): assert model is not None, "'model' is required and must be either a 'torch.nn.Module' or ONNX model" assert isinstance(model_desc, dict), "'model_desc' must be a 'dict'" assert isinstance(optim_config, optim._OptimizerConfig),\ @@ -205,7 +206,8 @@ class ORTTrainer(object): self._train_step_info = TrainStepInfo(self.optim_config) self._training_session = None self._load_state_dict = None - self._init_session() + self._init_session(provider_options=self.options._validated_opts['provider_options'], + session_options=self.options.session_options) def eval_step(self, *args, **kwargs): r"""Evaluation step method @@ -564,7 +566,10 @@ class ORTTrainer(object): return onnx_model - def _create_ort_training_session(self, optimizer_state_dict={}): + def _create_ort_training_session(self, + optimizer_state_dict={}, + session_options=None, + provider_options=None): # Validating frozen_weights names unused_frozen_weights = [n for n in self.options.utils.frozen_weights\ if n not in [i.name for i in self._onnx_model.graph.initializer]] @@ -662,7 +667,7 @@ class ORTTrainer(object): ort_parameters.model_with_training_graph_path = self.options.debug.graph_save_paths.model_with_training_graph_path # SessionOptions - session_options = ort.SessionOptions() + session_options = ort.SessionOptions() if session_options is None else session_options session_options.use_deterministic_compute = self.options.debug.deterministic_compute if (self.options.graph_transformer.attn_dropout_recompute or self.options.graph_transformer.gelu_recompute or @@ -676,10 +681,14 @@ class ORTTrainer(object): del self._training_session # Set provider-specific options if needed - def get_providers(): + def get_providers(provider_options): providers = ort.get_available_providers() - - if 'cuda' in self.options.device.id.lower(): + + if provider_options: + for provider_name in provider_options: + providers[providers.index(provider_name)] = (provider_name, provider_options[provider_name]) + #default: using cuda + elif 'cuda' in self.options.device.id.lower(): cuda_ep_options = {"device_id": _utils.get_device_index(self.options.device.id)} cuda_ep_name = ("ROCMExecutionProvider" if self.is_rocm_pytorch else "CUDAExecutionProvider") @@ -700,7 +709,7 @@ class ORTTrainer(object): # TrainingSession self._training_session = ort.TrainingSession(self._onnx_model.SerializeToString(), ort_parameters, - session_options, get_providers()) + session_options, get_providers(provider_options)) # I/O bindings self._train_io_binding = self._training_session.io_binding() @@ -731,9 +740,13 @@ class ORTTrainer(object): if self._load_state_dict: optimizer_state_dict = self._load_state_dict() - self._init_session(optimizer_state_dict) + self._init_session(optimizer_state_dict, + session_options=self.options.session_options, + provider_options=self.options._validated_opts['provider_options']) - def _init_session(self, optimizer_state_dict={}): + def _init_session(self, optimizer_state_dict={}, + session_options=None, + provider_options=None): if self._onnx_model is None: return @@ -742,7 +755,9 @@ class ORTTrainer(object): # Create training session used by train_step # pass all optimizer states to the backend - self._create_ort_training_session(optimizer_state_dict) + self._create_ort_training_session(optimizer_state_dict, + session_options=session_options, + provider_options=provider_options) # Update model description to update dtype when mixed precision is enabled # C++ backend modifies model's output dtype from float32 to float16 for mixed precision @@ -867,9 +882,20 @@ class ORTTrainer(object): # Keep all finite flag on CPU to match backend implementation # This prevents CPU -> GPU -> CPU copies between frontend and backend target_device = 'cpu' + # the self.options.device may be a device that pytorch does not recognize. + # in that case, we temporary prefer to leave the input/output on CPU and let ORT session + # to move the data between device and host. + # so output will be on the same device as input. + try: + test_pt_device = torch.device(target_device) + except: + #in this case, input/output must on CPU + assert(input.device == 'cpu') + target_device = 'cpu' + torch_tensor = torch.zeros(output_desc.shape, device=target_device, dtype=output_desc.dtype_amp if output_desc.dtype_amp else output_desc.dtype) - iobinding.bind_output(output_desc.name, torch_tensor.device.type, _utils.get_device_index(self.options.device.id), + iobinding.bind_output(output_desc.name, torch_tensor.device.type, _utils.get_device_index(target_device), _utils.dtype_torch_to_numpy(torch_tensor.dtype), list(torch_tensor.size()), torch_tensor.data_ptr()) result[output_desc.name] = torch_tensor @@ -1300,7 +1326,9 @@ class ORTTrainer(object): # create a new training session after loading initializer states onto the onnx graph # pass the populated states to the training session to populate the backend graph - self._init_session(optimizer_state_dict) + self._init_session(optimizer_state_dict, + session_options=self.options.session_options, + provider_options=self.options._validated_opts['provider_options']) def save_checkpoint(self, path, user_dict={}, include_optimizer_states=True): """Persists ORTTrainer state dictionary on disk along with user_dict. diff --git a/orttraining/orttraining/python/training/orttrainer_options.py b/orttraining/orttraining/python/training/orttrainer_options.py index 8208ea22be..fab4cba259 100644 --- a/orttraining/orttraining/python/training/orttrainer_options.py +++ b/orttraining/orttraining/python/training/orttrainer_options.py @@ -3,7 +3,7 @@ import torch from .optim import lr_scheduler from .amp import loss_scaler - +import onnxruntime as ort class ORTTrainerOptions(object): r"""Settings used by ONNX Runtime training backend @@ -277,7 +277,18 @@ class ORTTrainerOptions(object): 'default' : True } } - } + }, + 'provider_options':{ + 'type': 'dict', + 'default': {}, + 'required': False, + 'schema': {} + }, + 'session_options': { + 'type': 'SessionOptions', + 'nullable': True, + 'default': None + }, } Keyword arguments: @@ -383,6 +394,11 @@ class ORTTrainerOptions(object): _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 @@ -459,9 +475,13 @@ class ORTTrainerOptionsValidator(cerberus.Validator): _LOSS_SCALER = cerberus.TypeDefinition( 'loss_scaler', (loss_scaler.LossScaler,), ()) + _SESSION_OPTIONS = cerberus.TypeDefinition( + 'session_options', (ort.SessionOptions,),()) + 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 def _check_is_callable(field, value, error): @@ -731,5 +751,17 @@ _ORTTRAINER_OPTIONS_SCHEMA = { 'default': True } } - } + }, + 'provider_options':{ + 'type': 'dict', + 'default_setter': lambda _: {}, + 'required': False, + 'allow_unknown': True, + 'schema': {} + }, + 'session_options': { + 'type': 'session_options', + 'nullable': True, + 'default': None + }, } diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index f6d347f860..dcd16f1a4a 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -18,6 +18,7 @@ from onnxruntime.training import _utils, amp, checkpoint, optim, orttrainer, Tra model_desc_validation as md_val,\ orttrainer_options as orttrainer_options import _test_commons,_test_helpers +from onnxruntime import SessionOptions ############################################################################### @@ -99,7 +100,9 @@ def testORTTrainerOptionsDefaultValues(test_input): 'extra_postprocess': None, 'onnx_opset_version' : 12, 'enable_onnx_contrib_ops': True, - } + }, + 'provider_options':{}, + 'session_options': None, } actual_values = orttrainer_options.ORTTrainerOptions(test_input) @@ -532,6 +535,15 @@ def testLRSchedulerUpdateImpl(lr_scheduler, expected_values): assert_allclose(lr_list[0], expected_values[optimization_step], rtol=rtol, err_msg="lr mismatch") +def testInstantiateORTTrainerOptions(): + session_options = SessionOptions() + session_options.enable_mem_pattern = False + provider_options = {'EP1': {'key':'val'}} + opts = {'session_options' : session_options, + 'provider_options' : provider_options} + opts = orttrainer.ORTTrainerOptions(opts) + assert(opts.session_options.enable_mem_pattern is False) + assert(opts._validated_opts['provider_options']['EP1']['key'] == 'val') @pytest.mark.parametrize("step_fn, lr_scheduler, expected_lr_values, device", [ ('train_step', None, None, 'cuda'),