diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 8d5af41297..04d1f064e2 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -125,6 +125,10 @@ class GraphExecutionManager(GraphExecutionInterface): # To be instantiated in the concrete implementation of GraphExecutionManager self._export_mode = None + # Exporter can take extra arguments for ORTModule extensions + # It cannot overlap with required/immutable arguments (validated in runtime) + self._export_extra_kwargs = {} + # Related to training graph shape inference self._current_input_shape = None # default execution order is priority-based for both dynamic/static shape input for now @@ -353,18 +357,23 @@ class GraphExecutionManager(GraphExecutionInterface): try: with torch.set_grad_enabled(self._enable_custom_autograd_function), \ _logger.suppress_os_stream_output(log_level=self._debug_options.logging.log_level): + required_export_kwargs = {'input_names': self._input_info.names, + 'output_names': output_names, + 'opset_version': ONNX_OPSET_VERSION, + 'do_constant_folding': False, + 'training': self._export_mode, + 'dynamic_axes': self._input_info.dynamic_axes, + 'verbose': self._debug_options.logging.log_level < LogLevel.WARNING, + 'export_params': False, + 'keep_initializers_as_inputs': True} + invalid_args = self._export_extra_kwargs.keys() & required_export_kwargs.keys() + assert len(invalid_args) == 0,\ + f"The following PyTorch exporter arguments cannot be specified: '{invalid_args}'." torch.onnx.export(self._flattened_module, sample_inputs_as_tuple, f, - input_names=self._input_info.names, - output_names=output_names, - opset_version=ONNX_OPSET_VERSION, - do_constant_folding=False, - training=self._export_mode, - dynamic_axes=self._input_info.dynamic_axes, - verbose=self._debug_options.logging.log_level < LogLevel.WARNING, - export_params=False, - keep_initializers_as_inputs=True) + **required_export_kwargs, + **self._export_extra_kwargs) except Exception as e: raise wrap_exception(ORTModuleONNXModelException, RuntimeError(f'There was an error while exporting the PyTorch model to ONNX: ' diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 40a9933cbc..b9cb2c6af4 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -23,7 +23,13 @@ from distutils.version import LooseVersion from onnxruntime.training.ortmodule._custom_gradient_registry import register_gradient -from onnxruntime.training.ortmodule import ORTModule, _utils, _io, DebugOptions, LogLevel, _fallback, _graph_execution_manager +from onnxruntime.training.ortmodule import (ORTModule, + _utils, + _io, + DebugOptions, + LogLevel, + _fallback, + _graph_execution_manager) import _test_helpers # Import autocasting libs @@ -3862,7 +3868,7 @@ def test_ortmodule_determinism_flag(is_training,deterministic): model = ORTModule(model) model.train(is_training) - for i in range(5): + for _ in range(5): x = torch.randn(N, D_in) _ = model(x) @@ -3900,3 +3906,74 @@ def test_ortmodule_gradient_builder(): ort_prediction = run_step(ort_model, ort_x) _test_helpers.assert_values_are_close(ort_prediction, pt_prediction) _test_helpers.assert_values_are_close(ort_x.grad, pt_x.grad ) + + +def test_override_pytorch_exporter_kwargs(): + device = 'cuda' + + N, D_in, H, D_out = 64, 784, 500, 10 + x = torch.randn(N, D_in, device=device) + model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) + + ort_model = ORTModule(model) + ort_model._torch_module._execution_manager(True)._export_extra_kwargs = {'custom_opsets': None} + + # Make sure model runs without any exception + prediction = ort_model(x) + assert prediction is not None + prediction = prediction.sum() + prediction.backward() + + +def test_override_pytorch_exporter_kwargs__invalid(): + device = 'cuda' + + N, D_in, H, D_out = 64, 784, 500, 10 + x = torch.randn(N, D_in, device=device) + model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) + + ort_model = ORTModule(model) + ort_model._torch_module._execution_manager(True)._export_extra_kwargs = {'verbose': False} + with pytest.raises(_fallback.ORTModuleONNXModelException) as type_error: + _ = ort_model(x) + assert "The following PyTorch exporter arguments cannot be specified: '{'verbose'}'." in str(type_error.value) + + +def test_override_pytorch_exporter_kwargs_using_ortmodule_extension__invalid(): + device = 'cuda' + + class ORTModuleExtension(ORTModule): + def __init__(self, module, debug_options=None): + super().__init__(module, debug_options) + for training_mode in [False, True]: + self._torch_module._execution_manager(training_mode)._export_extra_kwargs = {'verbose': None} + + N, D_in, H, D_out = 64, 784, 500, 10 + x = torch.randn(N, D_in, device=device) + model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) + ort_model = ORTModuleExtension(model) + + with pytest.raises(_fallback.ORTModuleONNXModelException) as type_error: + _ = ort_model(x) + assert "The following PyTorch exporter arguments cannot be specified: '{'verbose'}'." in str(type_error.value) + +def test_override_pytorch_exporter_kwargs_using_ortmodule_extension(): + device = 'cuda' + + class ORTModuleExtension(ORTModule): + def __init__(self, module, debug_options=None): + super().__init__(module, debug_options) + # modify GraphExecutionManager internally + for training_mode in [False, True]: + self._torch_module._execution_manager(training_mode)._export_extra_kwargs = {'custom_opsets': None} + + N, D_in, H, D_out = 64, 784, 500, 10 + x = torch.randn(N, D_in, device=device) + model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) + ort_model = ORTModuleExtension(model) + + # Make sure model runs without any exception + prediction = ort_model(x) + assert prediction is not None + prediction = prediction.sum() + prediction.backward()