From 153767bab42cb5b51296c8fb269be6ccdeacf99d Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Tue, 21 Sep 2021 15:11:41 -0400 Subject: [PATCH] Add internal determinism flag configuration for ORTModule (#9074) --- .../python/training/ortmodule/__init__.py | 47 +++++--- .../ortmodule/_graph_execution_manager.py | 25 +++-- .../training/ortmodule/_inference_manager.py | 34 +++--- .../training/ortmodule/_training_manager.py | 80 ++++++++----- .../python/training/ortmodule/ortmodule.py | 20 ++-- .../python/orttraining_test_ortmodule_api.py | 18 +++ .../orttraining_test_ortmodule_fallback.py | 106 +++++++++++------- 7 files changed, 213 insertions(+), 117 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/__init__.py b/orttraining/orttraining/python/training/ortmodule/__init__.py index 742c92295a..660a4b26e5 100644 --- a/orttraining/orttraining/python/training/ortmodule/__init__.py +++ b/orttraining/orttraining/python/training/ortmodule/__init__.py @@ -5,11 +5,11 @@ import os import sys +import torch +from onnxruntime import set_seed from packaging import version - -from ._fallback import (_FallbackManager, - _FallbackPolicy, +from ._fallback import (_FallbackPolicy, ORTModuleFallbackException, ORTModuleInitException, wrap_exception) @@ -30,6 +30,7 @@ ORTMODULE_FALLBACK_POLICY = _FallbackPolicy.FALLBACK_UNSUPPORTED_DEVICE |\ _FallbackPolicy.FALLBACK_UNSUPPORTED_ONNX_MODEL |\ _FallbackPolicy.FALLBACK_BAD_INITIALIZATION ORTMODULE_FALLBACK_RETRY = False +ORTMODULE_IS_DETERMINISTIC = torch.are_deterministic_algorithms_enabled() # Verify minimum PyTorch version is installed before proceding to ONNX Runtime initialization try: @@ -39,41 +40,59 @@ try: if runtime_pytorch_version < minimum_runtime_pytorch_version: raise wrap_exception(ORTModuleInitException, RuntimeError( - f'ONNX Runtime ORTModule frontend requires PyTorch version greater or equal to {MINIMUM_RUNTIME_PYTORCH_VERSION_STR}, ' - f'but version {torch.__version__} was found instead.')) + 'ONNX Runtime ORTModule frontend requires PyTorch version greater' + f' or equal to {MINIMUM_RUNTIME_PYTORCH_VERSION_STR},' + f' but version {torch.__version__} was found instead.')) except ORTModuleFallbackException as e: # Initialization fallback is handled at ORTModule.__init__ _FALLBACK_INIT_EXCEPTION = e except ImportError as e: - raise RuntimeError(f'PyTorch {MINIMUM_RUNTIME_PYTORCH_VERSION_STR} must be installed in order to run ONNX Runtime ORTModule frontend!') from e + raise RuntimeError(f'PyTorch {MINIMUM_RUNTIME_PYTORCH_VERSION_STR} must be ' + 'installed in order to run ONNX Runtime ORTModule frontend!') from e # Verify whether PyTorch C++ extensions are already compiled if not is_torch_cpp_extensions_installed(TORCH_CPP_DIR) and '-m' not in sys.argv: - _FALLBACK_INIT_EXCEPTION = wrap_exception(ORTModuleInitException, - EnvironmentError( - f"ORTModule's extensions were not detected at '{TORCH_CPP_DIR}' folder. " - "Run `python -m torch_ort.configure` before using `ORTModule` frontend.")) + _FALLBACK_INIT_EXCEPTION = wrap_exception( + ORTModuleInitException, + EnvironmentError( + f"ORTModule's extensions were not detected at '{TORCH_CPP_DIR}' folder. " + "Run `python -m torch_ort.configure` before using `ORTModule` frontend.")) # Initalized ORT's random seed with pytorch's initial seed # in case user has set pytorch seed before importing ORTModule -import sys -from onnxruntime import set_seed set_seed((torch.initial_seed() % sys.maxsize)) + # Override torch.manual_seed and torch.cuda.manual_seed def override_torch_manual_seed(seed): set_seed(int(seed % sys.maxsize)) return torch_manual_seed(seed) + + torch_manual_seed = torch.manual_seed torch.manual_seed = override_torch_manual_seed + def override_torch_cuda_manual_seed(seed): set_seed(int(seed % sys.maxsize)) return torch_cuda_manual_seed(seed) + + torch_cuda_manual_seed = torch.cuda.manual_seed torch.cuda.manual_seed = override_torch_cuda_manual_seed + +def _use_deterministic_algorithms(enabled): + global ORTMODULE_IS_DETERMINISTIC + ORTMODULE_IS_DETERMINISTIC = enabled + + +def _are_deterministic_algorithms_enabled(): + global ORTMODULE_IS_DETERMINISTIC + return ORTMODULE_IS_DETERMINISTIC + + # ORTModule must be loaded only after all validation passes -from .ortmodule import ORTModule -from .debug_options import DebugOptions, LogLevel +from .ortmodule import ORTModule # noqa: E402 +from .debug_options import DebugOptions, LogLevel # noqa: E402 diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 8fe54e1784..beb063bef6 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -4,17 +4,20 @@ # -------------------------------------------------------------------------- from .debug_options import DebugOptions, LogLevel -from . import _utils, _io, _logger, torch_cpp_extensions as _cpp_ext, _onnx_models +from . import (_utils, + _io, + _logger, + torch_cpp_extensions as _cpp_ext, + _onnx_models, + _are_deterministic_algorithms_enabled) from ._custom_autograd_function import custom_autograd_function_enabler from ._custom_autograd_function_exporter import _post_process_after_export from ._graph_execution_interface import GraphExecutionInterface from ._fallback import (_FallbackManager, - _FallbackPolicy, - ORTModuleFallbackException, - ORTModuleDeviceException, - ORTModuleONNXModelException, - ORTModuleTorchModelException, - wrap_exception) + ORTModuleDeviceException, + ORTModuleONNXModelException, + ORTModuleTorchModelException, + wrap_exception) from ._gradient_accumulation_manager import GradientAccumulationManager from onnxruntime.training.ortmodule import ONNX_OPSET_VERSION @@ -232,6 +235,12 @@ class GraphExecutionManager(GraphExecutionInterface): def _get_session_config(self): """Creates and returns the session configuration to be used for the ExecutionAgent""" + + if _are_deterministic_algorithms_enabled(): + if self._debug_options.logging.log_level <= _logger.LogLevel.INFO: + warnings.warn("ORTModule's determinism will be enabled because PyTorch's determinism is enabled.", + UserWarning) + providers = None provider_options = None if self._device.type == 'cuda': @@ -256,7 +265,7 @@ class GraphExecutionManager(GraphExecutionInterface): session_options = onnxruntime.SessionOptions() session_options.enable_mem_pattern = False session_options.enable_mem_reuse = False - session_options.use_deterministic_compute = False + session_options.use_deterministic_compute = _are_deterministic_algorithms_enabled() # default to PRIORITY_BASED execution order session_options.execution_order = onnxruntime.ExecutionOrder.PRIORITY_BASED # 0:Verbose, 1:Info, 2:Warning. 3:Error, 4:Fatal. Default is 2. diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index 8a754891f4..d4e2e59caa 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -3,14 +3,19 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from . import _utils, _io, _logger -from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck +from . import (_utils, + _io, + _logger, + _are_deterministic_algorithms_enabled, + _use_deterministic_algorithms) +from ._graph_execution_manager import (GraphExecutionManager, + _RunStateInfo, + _SkipCheck) from ._execution_agent import InferenceAgent from .debug_options import DebugOptions from ._fallback import ORTModuleFallbackException, _FallbackPolicy, _FallbackManager from onnxruntime.capi import _pybind_state as C -import onnx import torch import warnings @@ -66,9 +71,9 @@ class InferenceManager(GraphExecutionManager): return self._fallback_manager.fallback(self._original_module, self._debug_options.logging.log_level, *inputs, **kwargs) try: - if self._first_skip_check_warning == True and self._skip_check.is_disabled() == False \ - and self._debug_options.logging.log_level <= _logger.LogLevel.WARNING: - # Only change this after the firs time a warning is issued. + # Issue at most one warning message about fast path + if self._first_skip_check_warning is True and self._skip_check.is_disabled() is False \ + and self._debug_options.logging.log_level <= _logger.LogLevel.WARNING: self._first_skip_check_warning = False warnings.warn(f"Fast path enabled - skipping checks." f"rebuild gradient graph: {self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT)}," @@ -78,7 +83,7 @@ class InferenceManager(GraphExecutionManager): # If exporting module to ONNX for the first time, this skip check will not take effect. # It will only take effect on subsequent forward calls. build_graph = False - if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT) == False or \ + if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT) is False or \ not self._onnx_models.exported_model: # Exporting module to ONNX for the first time build_graph = self._export_model(*inputs, **kwargs) @@ -93,13 +98,16 @@ class InferenceManager(GraphExecutionManager): # If creating the execution agent for the first time, this skip check will not take effect. # It will only take effect on subsequent forward calls. create_execution_session = False - if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT) == False or \ - not self._execution_agent: + if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT) is False or \ + not self._execution_agent: module_device = _utils.get_device_from_module( self._original_module) - # The inference session should be created every time - # the graph was built or if the device changed between calls to forward - create_execution_session = build_graph or self._device != module_device + + create_execution_session = (build_graph or self._device != module_device or + torch.are_deterministic_algorithms_enabled() is not + _are_deterministic_algorithms_enabled()) + _use_deterministic_algorithms(torch.are_deterministic_algorithms_enabled()) + if self._device != module_device: self._device = module_device @@ -107,7 +115,7 @@ class InferenceManager(GraphExecutionManager): # Create execution session creates the inference_session self._create_execution_agent() - if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) == False: + if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) is False: # Assert that the input and model device match _utils._check_same_device(self._device, "Input argument to forward", *inputs) diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index f2d2726812..b5ef430dbf 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -3,11 +3,19 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from . import _utils, _io, _logger -from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck +from . import (_utils, + _io, + _logger, + _are_deterministic_algorithms_enabled, + _use_deterministic_algorithms) +from ._graph_execution_manager import (GraphExecutionManager, + _RunStateInfo, + _SkipCheck) from ._execution_agent import TrainingAgent from .debug_options import DebugOptions -from ._fallback import ORTModuleFallbackException, _FallbackPolicy, _FallbackManager +from ._fallback import (ORTModuleFallbackException, + _FallbackPolicy, + _FallbackManager) from onnxruntime.capi import _pybind_state as C from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_type @@ -15,6 +23,7 @@ from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_typ import torch import warnings + class TrainingManager(GraphExecutionManager): """Concrete instance of GraphExecutionManager that is able to manage the training model @@ -61,23 +70,24 @@ class TrainingManager(GraphExecutionManager): # Fallback to PyTorch due to failures *external* to forward(), # typically from initialization if self._fallback_manager.is_pending(): - return self._fallback_manager.fallback(self._original_module, self._debug_options.logging.log_level, *inputs, **kwargs) + return self._fallback_manager.fallback(self._original_module, self._debug_options.logging.log_level, + *inputs, **kwargs) try: - if self._first_skip_check_warning == True and self._skip_check.is_disabled() == False \ - and self._debug_options.logging.log_level <= _logger.LogLevel.WARNING: + if self._first_skip_check_warning is True and self._skip_check.is_disabled() is False \ + and self._debug_options.logging.log_level <= _logger.LogLevel.WARNING: # Only change this after the firs time a warning is issued. self._first_skip_check_warning = False warnings.warn(f"Fast path enabled - skipping checks." - f"rebuild gradient graph: {self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT)}," - f"execution agent recreation: {self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT)}," - f"device check: {self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE)}", UserWarning) + f" Rebuild graph: {self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT)}," + f" Execution agent: {self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT)}," + f" Device check: {self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE)}", UserWarning) # If exporting module to ONNX for the first time, this skip check will not take effect. # It will only take effect on subsequent forward calls. build_gradient_graph = False - if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT) == False or \ - not self._onnx_models.exported_model: + if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT) is False or \ + not self._onnx_models.exported_model: build_gradient_graph = self._export_model(*inputs, **kwargs) if build_gradient_graph: # If model was exported, then initialize the graph builder @@ -105,13 +115,14 @@ class TrainingManager(GraphExecutionManager): # If creating the execution agent for the first time, this skip check will not take effect. # It will only take effect on subsequent forward calls. create_execution_session = False - if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT) == False or \ - not self._execution_agent: + if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT) is False or \ + not self._execution_agent: device = _utils.get_device_from_module(self._original_module) or \ _utils.get_device_from_inputs(inputs, kwargs) - # The _training_session/_inference_session should be created every time - # the graph was built or if the device changed between calls to forward - create_execution_session = build_gradient_graph or self._device != device + create_execution_session = (build_gradient_graph or self._device != device or + torch.are_deterministic_algorithms_enabled() is not + _are_deterministic_algorithms_enabled()) + _use_deterministic_algorithms(torch.are_deterministic_algorithms_enabled()) if self._device != device: self._device = device @@ -119,8 +130,10 @@ class TrainingManager(GraphExecutionManager): # Create execution session creates the training_session self._create_execution_agent() - self._gradient_accumulation_manager.initialize(self._enable_grad_acc_optimization, self._flattened_module, self._graph_info) - + self._gradient_accumulation_manager.initialize(self._enable_grad_acc_optimization, + self._flattened_module, + self._graph_info) + self._gradient_accumulation_manager.maybe_update_cache_before_run() class _ORTModuleFunction(torch.autograd.Function): @@ -138,19 +151,20 @@ class TrainingManager(GraphExecutionManager): Module outputs are returned to the user ''' - if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) == False: + if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) is False: # Assert that the input and model device match _utils._check_same_device(self._device, "Input argument to forward", *inputs) - user_outputs, ctx.run_info = TrainingManager.execution_session_run_forward(self._execution_agent, - self._onnx_models.optimized_model, - self._device, - self._gradient_accumulation_manager, - *inputs) + user_outputs, ctx.run_info = TrainingManager.execution_session_run_forward( + self._execution_agent, + self._onnx_models.optimized_model, + self._device, + self._gradient_accumulation_manager, + *inputs) # Disable materializing grads then None object will not be # converted to a tensor filled with zeros prior to calling backward. - # Save shape, device and type info to ctx for materializing tensor in backward if output grad is None. + # Save shape/device/type info to ctx for materializing tensor in backward if output grad is None. ctx.set_materialize_grads(False) # Mark the outputs tensors needed in backward computation @@ -158,7 +172,7 @@ class TrainingManager(GraphExecutionManager): # as this tensor is also kept in ORT's PartialGraphState # This call is to invoke pytorch's version check to detect the potential inplace corruption # If ORT is caching tensors, the module_output_indices_requires_save_for_backward field - # might also have indices of cached tensors that are not passed over to pytorch, and they don't + # might also have indices of cached tensors that are not passed over to pytorch, and they don't # need marking with save_for_backward() for idx in self._graph_info.module_output_indices_requires_save_for_backward: if idx < len(self._graph_info.user_output_names): @@ -176,7 +190,7 @@ class TrainingManager(GraphExecutionManager): '''Performs backward pass based on grad wrt module output''' assert ctx.run_info is not None, 'forward() or __call__() methods must be called before backward()' - if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) == False: + if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) is False: _utils._check_same_device(self._device, "Input argument to backward", *grad_outputs) # Unpack saved_tensor to trigger version detection that catches inplace corruption @@ -204,7 +218,8 @@ class TrainingManager(GraphExecutionManager): grad_output = torch.tensor(0., device=device, dtype=dtype) elif not grad_output.is_contiguous(): grad_output = grad_output.contiguous() - backward_inputs.push_back(_utils._torch_tensor_to_dlpack(grad_output), grad_output.dtype == torch.bool) + backward_inputs.push_back(_utils._torch_tensor_to_dlpack(grad_output), + grad_output.dtype is torch.bool) backward_inputs.shrink_to_fit() # Run and get results @@ -267,7 +282,10 @@ class TrainingManager(GraphExecutionManager): # Fallback to PyTorch due to failures *during* forward(), # (e.g. export, model/input post-processing, forward, output processing, etc) if self._fallback_manager.is_pending(): - return self._fallback_manager.fallback(self._original_module, self._debug_options.logging.log_level, *inputs, **kwargs) + return self._fallback_manager.fallback(self._original_module, + self._debug_options.logging.log_level, + *inputs, + **kwargs) def _build_graph(self): """Build an optimized gradient graph using the module_graph_builder""" @@ -288,7 +306,7 @@ class TrainingManager(GraphExecutionManager): C.OrtDevice(get_ort_device_type(self._device), C.OrtDevice.default_memory(), _utils.get_device_index(self._device) - )] * (len(self._graph_info.user_output_names) + + )] * (len(self._graph_info.user_output_names) + len(self._graph_info.frontier_node_arg_map)) bw_fetches_names = [output.name for output in self._onnx_models.optimized_model.graph.output] @@ -310,7 +328,7 @@ class TrainingManager(GraphExecutionManager): def _reinitialize_graph_builder(self, input_info): """Return true if the module graph builder was reinitialized""" - # Model could have unused parameters which are dropped after export and so not a part of self._graph_initializer_names_to_train. + # Model may have unused params dropped after export and not part of self._graph_initializer_names_to_train # To see if any trainable initializers changed, compare self._graph_initializer_names_to_train # with initializers in module named_parameters that are known to the onnx graph. initializer_names_to_train_set_user_model = {name for name, param in diff --git a/orttraining/orttraining/python/training/ortmodule/ortmodule.py b/orttraining/orttraining/python/training/ortmodule/ortmodule.py index e2cd8fdec5..a642261600 100644 --- a/orttraining/orttraining/python/training/ortmodule/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule/ortmodule.py @@ -10,14 +10,19 @@ from ._custom_op_symbolic_registry import CustomOpSymbolicRegistry from ._custom_gradient_registry import CustomGradientRegistry from . import _utils from .debug_options import DebugOptions -from ._fallback import _FallbackManager, _FallbackPolicy, ORTModuleFallbackException, ORTModuleTorchModelException, wrap_exception -from . import _FALLBACK_INIT_EXCEPTION, MINIMUM_RUNTIME_PYTORCH_VERSION_STR, ORTMODULE_FALLBACK_POLICY, ORTMODULE_FALLBACK_RETRY +from ._fallback import (_FallbackManager, + _FallbackPolicy, + ORTModuleFallbackException) +from . import (_FALLBACK_INIT_EXCEPTION, + ORTMODULE_FALLBACK_POLICY, + ORTMODULE_FALLBACK_RETRY) + from onnxruntime.tools import pytorch_export_contrib_ops import functools import torch -from typing import Iterator, Optional, Tuple, TypeVar, Set, Callable -import warnings +from typing import Iterator, Optional, Tuple, TypeVar, Callable + # Needed to override PyTorch methods T = TypeVar('T', bound='Module') @@ -59,7 +64,6 @@ class ORTModule(torch.nn.Module): try: # Read ORTModule module initialization status - global _FALLBACK_INIT_EXCEPTION if _FALLBACK_INIT_EXCEPTION: raise _FALLBACK_INIT_EXCEPTION @@ -296,8 +300,8 @@ class ORTModule(torch.nn.Module): yield from self._torch_module.named_modules(*args, **kwargs) def __getattr__(self, name: str): - if '_is_initialized' in self.__dict__ and self.__dict__['_is_initialized'] == True: - # If ORTModule is intitialized and attribute is not found in ORTModule, + if '_is_initialized' in self.__dict__ and self.__dict__['_is_initialized'] is True: + # If ORTModule is initialized and attribute is not found in ORTModule, # it must be present in the user's torch.nn.Module. Forward the call to # the user's model. assert '_torch_module' in self.__dict__, "ORTModule does not have a reference to the user's model" @@ -311,7 +315,7 @@ class ORTModule(torch.nn.Module): # If the name is an attribute of ORTModule, update only ORTModule self.__dict__[name] = value - elif '_is_initialized' in self.__dict__ and self.__dict__['_is_initialized'] == True: + elif '_is_initialized' in self.__dict__ and self.__dict__['_is_initialized'] is True: assert '_torch_module' in self.__dict__, "ORTModule does not have a reference to the user's model" diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 4256f58a04..ef4d894f96 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -3846,3 +3846,21 @@ def test_ortmodule_skip_check_load_from_os_env(policy_str, policy): assert ort_model._torch_module._execution_manager(training_mode)._skip_check == policy del os.environ['ORTMODULE_SKIPCHECK_POLICY'] + +@pytest.mark.parametrize("is_training,deterministic", + list(itertools.product([True,False],repeat=2))) +def test_ortmodule_determinism_flag(is_training,deterministic): + + torch.use_deterministic_algorithms(deterministic) + + N, D_in, H, D_out = 64, 784, 500, 10 + model = NeuralNetSinglePositionalArgument(D_in, H, D_out) + model = ORTModule(model) + model.train(is_training) + + for i in range(5): + x = torch.randn(N, D_in) + _ = model(x) + + from onnxruntime.training.ortmodule import _are_deterministic_algorithms_enabled + assert _are_deterministic_algorithms_enabled() is torch.are_deterministic_algorithms_enabled() diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py index ee6ce75ccc..d4cab55f74 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py @@ -9,7 +9,7 @@ import torch import pytest import warnings -from onnxruntime.training.ortmodule import ORTModule, _utils, _io, LogLevel, _fallback, TORCH_CPP_DIR +from onnxruntime.training.ortmodule import ORTModule, _fallback, TORCH_CPP_DIR from onnxruntime.training.ortmodule.torch_cpp_extensions import is_installed as is_torch_cpp_extensions_installed import _test_helpers from _orttraining_ortmodule_models import (NeuralNetSinglePositionalArgument, @@ -19,12 +19,13 @@ from _orttraining_ortmodule_models import (NeuralNetSinglePositionalArgument, # PyTorch model definitions for tests + @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_forward(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_FORCE_TORCH_FORWARD policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_FORCE_TORCH_FORWARD policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen if fallback_enabled: @@ -38,6 +39,7 @@ def test_ortmodule_fallback_forward(is_training, fallback_enabled, matching_poli os.environ['ORTMODULE_FALLBACK_RETRY'] = str(not persist_fallback) from dataclasses import dataclass + @dataclass class Point: x: int @@ -46,6 +48,7 @@ def test_ortmodule_fallback_forward(is_training, fallback_enabled, matching_poli class UnsupportedInputModel(torch.nn.Module): def __init__(self): super(UnsupportedInputModel, self).__init__() + def forward(self, point): return point.x * point.y @@ -60,7 +63,8 @@ def test_ortmodule_fallback_forward(is_training, fallback_enabled, matching_poli if fallback_enabled: if matching_policy: if i > 0 and persist_fallback: - assert ort_model._torch_module._execution_manager(is_training=is_training)._fallback_manager._exception is not None + assert ort_model._torch_module._execution_manager( + is_training=is_training)._fallback_manager._exception is not None ort_out = ort_model(inputs) pt_out = pt_model(inputs) assert ort_out == pt_out @@ -75,11 +79,11 @@ def test_ortmodule_fallback_forward(is_training, fallback_enabled, matching_poli @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_device__multiple(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_DEVICE policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_DEVICE policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DATA) is used to verify that the fallback does not happen if fallback_enabled: @@ -128,11 +132,11 @@ def test_ortmodule_fallback_device__multiple(is_training, fallback_enabled, matc @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_device__mismatch(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_DEVICE policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_DEVICE policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DATA) is used to verify that the fallback does not happen if fallback_enabled: @@ -165,22 +169,26 @@ def test_ortmodule_fallback_device__mismatch(is_training, fallback_enabled, matc if matching_policy: with pytest.raises(RuntimeError) as e: ort_model(inputs) - assert "Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!" in str(e.value) + assert ("Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!" + in str(e.value)) else: with pytest.raises(_fallback.ORTModuleDeviceException) as e: ort_model(inputs) - assert f"Input argument to forward found on device {input_device}, but expected it to be on module device {ort_model_device}." in str(e.value) + assert (f"Input argument to forward found on device {input_device}, " + f"but expected it to be on module device {ort_model_device}." in str(e.value)) else: with pytest.raises(_fallback.ORTModuleDeviceException) as e: ort_model(inputs) - assert f"Input argument to forward found on device {input_device}, but expected it to be on module device {ort_model_device}." in str(e.value) + assert (f"Input argument to forward found on device {input_device}, " + f"but expected it to be on module device {ort_model_device}." in str(e.value)) + @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_output(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_DATA policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_DATA policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen if fallback_enabled: @@ -208,7 +216,8 @@ def test_ortmodule_fallback_output(is_training, fallback_enabled, matching_polic if fallback_enabled: if matching_policy: if i > 0 and persist_fallback: - assert ort_model._torch_module._execution_manager(is_training=is_training)._fallback_manager._exception is not None + assert ort_model._torch_module._execution_manager( + is_training=is_training)._fallback_manager._exception is not None ort_out = ort_model(x, y, z) pt_out = pt_model(x, y, z) _test_helpers.assert_values_are_close(ort_out.out1, pt_out.out1, rtol=0, atol=0) @@ -223,12 +232,13 @@ def test_ortmodule_fallback_output(is_training, fallback_enabled, matching_polic ort_model(x, y, z) assert 'ORTModule does not support the following model output type' in str(runtime_error.value) + @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_input(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_DATA policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_DATA policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen if fallback_enabled: @@ -252,7 +262,8 @@ def test_ortmodule_fallback_input(is_training, fallback_enabled, matching_policy if fallback_enabled: if matching_policy: if i > 0 and persist_fallback: - assert ort_model._torch_module._execution_manager(is_training=is_training)._fallback_manager._exception is not None + assert ort_model._torch_module._execution_manager( + is_training=is_training)._fallback_manager._exception is not None ort_out = ort_model(inputs, 'hello') pt_out = pt_model(inputs, 'hello') _test_helpers.assert_values_are_close(ort_out, pt_out, rtol=0, atol=0) @@ -267,11 +278,11 @@ def test_ortmodule_fallback_input(is_training, fallback_enabled, matching_policy @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_torch_model(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_TORCH_MODEL policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_TORCH_MODEL policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen if fallback_enabled: @@ -311,12 +322,13 @@ def test_ortmodule_fallback_torch_model(is_training, fallback_enabled, matching_ ort_model = ORTModule(pt_model) assert "ORTModule is not compatible with torch.nn.DataParallel" in str(ex_info.value) + @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_init__torch_version(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_TORCH_MODEL policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_TORCH_MODEL policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen from packaging import version @@ -354,7 +366,8 @@ def test_ortmodule_fallback_init__torch_version(is_training, fallback_enabled, m else: with pytest.raises(_fallback.ORTModuleInitException) as ex_info: ort_model = ORTModule(pt_model) - assert "ONNX Runtime ORTModule frontend requires PyTorch version greater or equal to" in str(ex_info.value) + assert "ONNX Runtime ORTModule frontend requires PyTorch version greater or equal to" in str( + ex_info.value) else: with pytest.raises(_fallback.ORTModuleInitException) as ex_info: # Initialize with fallback policy because Exception will happen during __init__ @@ -366,11 +379,12 @@ def test_ortmodule_fallback_init__torch_version(is_training, fallback_enabled, m @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) -def test_ortmodule_fallback_init__missing_cpp_extensions(is_training, fallback_enabled, matching_policy, persist_fallback): + list(itertools.product([True, False], repeat=4))) +def test_ortmodule_fallback_init__missing_cpp_extensions(is_training, fallback_enabled, matching_policy, + persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_TORCH_MODEL policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_TORCH_MODEL policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen if is_torch_cpp_extensions_installed(TORCH_CPP_DIR): @@ -414,12 +428,14 @@ def test_ortmodule_fallback_init__missing_cpp_extensions(is_training, fallback_e ort_model = ORTModule(pt_model) assert "ORTModule's extensions were not detected" in str(ex_info.value) + @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) -def test_ortmodule_fallback_onnx_model__custom_autograd(is_training, fallback_enabled, matching_policy, persist_fallback): + list(itertools.product([True, False], repeat=4))) +def test_ortmodule_fallback_onnx_model__custom_autograd(is_training, fallback_enabled, matching_policy, + persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_ONNX_MODEL policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_ONNX_MODEL policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen if fallback_enabled: @@ -450,7 +466,8 @@ def test_ortmodule_fallback_onnx_model__custom_autograd(is_training, fallback_en if fallback_enabled: if matching_policy: if i > 0 and persist_fallback: - assert ort_model._torch_module._execution_manager(is_training=is_training)._fallback_manager._exception is not None + assert ort_model._torch_module._execution_manager( + is_training=is_training)._fallback_manager._exception is not None pt_out = pt_model(x.mm(w1)).mm(w2) ort_out = ort_model(x.mm(w1)).mm(w2) _test_helpers.assert_values_are_close(ort_out, pt_out, rtol=0, atol=0) @@ -464,12 +481,13 @@ def test_ortmodule_fallback_onnx_model__custom_autograd(is_training, fallback_en _ = ort_model(x.mm(w1)).mm(w2) assert "There was an error while exporting the PyTorch model to ONNX" in str(ex_info.value) + @pytest.mark.parametrize("is_training,fallback_enabled,matching_policy,persist_fallback", - list(itertools.product([True,False],repeat=4))) + list(itertools.product([True, False], repeat=4))) def test_ortmodule_fallback_onnx_model__missing_op(is_training, fallback_enabled, matching_policy, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise - # fallback_enabled: True results in PyTorch executing the forward graph instead of ORT backend - # matching_policy: True results in properly matching FALLBACK_UNSUPPORTED_ONNX_MODEL policy to ORTModuleDeviceException exception. + # fallback_enabled: True PyTorch executes the forward graph instead of ORT backend + # matching_policy: True matches FALLBACK_UNSUPPORTED_ONNX_MODEL policy to ORTModuleDeviceException exception. # Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen if fallback_enabled: @@ -497,7 +515,8 @@ def test_ortmodule_fallback_onnx_model__missing_op(is_training, fallback_enabled if fallback_enabled: if matching_policy: if i > 0 and persist_fallback: - assert ort_model._torch_module._execution_manager(is_training=is_training)._fallback_manager._exception is not None + assert ort_model._torch_module._execution_manager( + is_training=is_training)._fallback_manager._exception is not None pt_out = pt_model(x, y) ort_out = ort_model(x, y) _test_helpers.assert_values_are_close(ort_out, pt_out, rtol=0, atol=0) @@ -511,8 +530,9 @@ def test_ortmodule_fallback_onnx_model__missing_op(is_training, fallback_enabled _ = ort_model(x, y) assert "There was an error while exporting the PyTorch model to ONNX" in str(ex_info.value) + @pytest.mark.parametrize("is_training,persist_fallback", - list(itertools.product([True,False],repeat=2))) + list(itertools.product([True, False], repeat=2))) def test_ortmodule_fallback_warn_message(is_training, persist_fallback): # is_training: True for torch.nn.Module training model, eval mode otherwise