Add internal determinism flag configuration for ORTModule (#9074)

This commit is contained in:
Thiago Crepaldi 2021-09-21 15:11:41 -04:00 committed by GitHub
parent b175f98dcc
commit 153767bab4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 213 additions and 117 deletions

View file

@ -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

View file

@ -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.

View file

@ -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)

View file

@ -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

View file

@ -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"

View file

@ -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()

View file

@ -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