saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix all-or-nothing fallback for bad ORTModule init (#9277)

Browse source 
* Fix all-or-nothing fallback for bad ORTModule init

* Address comments
This commit is contained in:
Thiago Crepaldi 2021-10-06 15:12:27 -04:00 committed by GitHub
parent 510b58c877
commit 52d067402a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 47 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
orttraining/orttraining/python/training/ortmodule/__init__.py
View file

@ -26,8 +26,7 @@ _FALLBACK_INIT_EXCEPTION = None
ORTMODULE_FALLBACK_POLICY = _FallbackPolicy.FALLBACK_UNSUPPORTED_DEVICE |\
_FallbackPolicy.FALLBACK_UNSUPPORTED_DATA |\
_FallbackPolicy.FALLBACK_UNSUPPORTED_TORCH_MODEL |\
_FallbackPolicy.FALLBACK_UNSUPPORTED_ONNX_MODEL |\
_FallbackPolicy.FALLBACK_BAD_INITIALIZATION
_FallbackPolicy.FALLBACK_UNSUPPORTED_ONNX_MODEL
ORTMODULE_FALLBACK_RETRY = False
ORTMODULE_IS_DETERMINISTIC = torch.are_deterministic_algorithms_enabled()
@ -57,7 +56,7 @@ except ImportError as e:
if not is_torch_cpp_extensions_installed(ORTMODULE_TORCH_CPP_DIR) and '-m' not in sys.argv:
_FALLBACK_INIT_EXCEPTION = wrap_exception(
ORTModuleInitException,
EnvironmentError(
RuntimeError(
f"ORTModule's extensions were not detected at '{ORTMODULE_TORCH_CPP_DIR}' folder. "
"Run `python -m torch_ort.configure` before using `ORTModule` frontend."))

17
orttraining/orttraining/python/training/ortmodule/_fallback.py
View file

@ -67,11 +67,13 @@ class _FallbackManager(object):
'''
def __init__(self,
pytorch_module: torch.nn.Module,
policy: _FallbackPolicy,
retry: bool):
# Read policy from environment variable for testing purposes
self._original_module = pytorch_module
# Read policy from environment variable for testing purposes
policy = os.getenv('ORTMODULE_FALLBACK_POLICY', policy)
if isinstance(policy, str):
policy = _FallbackPolicy[policy]
@ -127,6 +129,15 @@ class _FallbackManager(object):
if log_level <= _logger.LogLevel.INFO:
warnings.warn(
f'Fallback for policy {policy.name} is pending.', UserWarning)
# ORTModuleInitException exceptions do not call `fallback()` through `GraphExecutionManager`,
# Instead, it fallbacks to PyTorch implicitly through `ORTModule._torch_module = TorchModulePytorch(module)`
if log_level <= _logger.LogLevel.WARNING and policy == _FallbackPolicy.FALLBACK_BAD_INITIALIZATION:
warnings.warn(
(f'Fallback to PyTorch due to exception {type(exception)} was triggered. '
'Report this issue with a minimal repro at https://www.github.com/microsoft/onnxruntime. '
f'See details below:\n\n{_utils.get_exception_as_string(exception)}'), UserWarning)
self._exception = exception
if override_policy is None:
@ -147,7 +158,7 @@ class _FallbackManager(object):
return self._exception is not None
def fallback(self, model: torch.nn.Module, log_level: _logger.LogLevel, *inputs, **kwargs):
def fallback(self, log_level: _logger.LogLevel, *inputs, **kwargs):
'''Executes user PyTorch `model` using the provided inputs and return the result'''
assert self.is_pending(), '`fallback` can only be called when there is a pending fallback'
@ -161,4 +172,4 @@ class _FallbackManager(object):
# Pending fallbacks are resetted to enforce retries
if self.retry:
self._exception = None
return model(*inputs, **kwargs)
return self._original_module(*inputs, **kwargs)

4
orttraining/orttraining/python/training/ortmodule/_inference_manager.py
View file

@ -68,7 +68,7 @@ class InferenceManager(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._debug_options.logging.log_level, *inputs, **kwargs)
try:
# Issue at most one warning message about fast path
@ -146,7 +146,7 @@ class InferenceManager(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._debug_options.logging.log_level, *inputs, **kwargs)
def _build_graph(self):
"""Build an optimized inference graph using the module_graph_builder"""

5
orttraining/orttraining/python/training/ortmodule/_torch_module_pytorch.py
View file

@ -2,14 +2,9 @@
# Licensed under the MIT License.
# _torch_module_pytorch.py
from . import _io
from .debug_options import DebugOptions
from ._graph_execution_manager_factory import GraphExecutionManagerFactory
from ._torch_module_interface import TorchModuleInterface
from ._fallback import _FallbackManager
from collections import OrderedDict
import functools
import torch
from typing import Iterator, Optional, Tuple, TypeVar, Callable

8
orttraining/orttraining/python/training/ortmodule/_training_manager.py
View file

@ -70,8 +70,7 @@ 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._debug_options.logging.log_level, *inputs, **kwargs)
try:
if self._first_skip_check_warning is True and self._skip_check.is_disabled() is False \
@ -282,10 +281,7 @@ 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._debug_options.logging.log_level, *inputs, **kwargs)
def _build_graph(self):
"""Build an optimized gradient graph using the module_graph_builder"""

18
orttraining/orttraining/python/training/ortmodule/_utils.py
View file

@ -6,6 +6,7 @@
from onnxruntime.capi.onnxruntime_inference_collection import OrtValue
from onnxruntime.capi import _pybind_state as C
from ._fallback_exceptions import ORTModuleDeviceException, wrap_exception
from ._torch_module_pytorch import TorchModulePytorch
import os
import copy
@ -208,3 +209,20 @@ def get_exception_as_string(exception):
raise exception
except:
return traceback.format_exc()
def switch_backend_to_pytorch(ortmodule, pytorch_module):
ortmodule._torch_module = TorchModulePytorch(pytorch_module)
# TODO: Rework by implementing the "__getattribute__" method.
# Assigning all default attributes from user's original torch.nn.Module into ORTModule
ortmodule._backward_hooks = pytorch_module._backward_hooks
ortmodule._forward_hooks = pytorch_module._forward_hooks
ortmodule._forward_pre_hooks = pytorch_module._forward_pre_hooks
ortmodule._parameters = pytorch_module._parameters
ortmodule._buffers = pytorch_module._buffers
ortmodule._non_persistent_buffers_set = pytorch_module._non_persistent_buffers_set
ortmodule._is_full_backward_hook = pytorch_module._is_full_backward_hook
ortmodule._state_dict_hooks = pytorch_module._state_dict_hooks
ortmodule._load_state_dict_pre_hooks = pytorch_module._load_state_dict_pre_hooks
ortmodule._modules = pytorch_module._modules
ortmodule.forward = pytorch_module.forward

25
orttraining/orttraining/python/training/ortmodule/ortmodule.py
View file

@ -59,7 +59,8 @@ class ORTModule(torch.nn.Module):
debug_options = DebugOptions()
# Fallback settings
self._fallback_manager = _FallbackManager(policy=ORTMODULE_FALLBACK_POLICY,
self._fallback_manager = _FallbackManager(pytorch_module=module,
policy=ORTMODULE_FALLBACK_POLICY,
retry=ORTMODULE_FALLBACK_RETRY)
try:
@ -101,26 +102,18 @@ class ORTModule(torch.nn.Module):
_utils.check_for_name_collisions_and_bind_methods_to_ortmodule(self, module)
except ORTModuleFallbackException as e:
self._torch_module = TorchModulePytorch(module)
# TODO: Rework by implementing the "__getattribute__" method.
# Assigning all default attributes from user's original torch.nn.Module into ORTModule
self._backward_hooks = module._backward_hooks
self._forward_hooks = module._forward_hooks
self._forward_pre_hooks = module._forward_pre_hooks
self._parameters = module._parameters
self._buffers = module._buffers
self._non_persistent_buffers_set = module._non_persistent_buffers_set
self._is_full_backward_hook = module._is_full_backward_hook
self._state_dict_hooks = module._state_dict_hooks
self._load_state_dict_pre_hooks = module._load_state_dict_pre_hooks
self._modules = module._modules
self.forward = module.forward
# Although backend is switched to PyTorch here,
# it is up to _FallbackManager to actually terminate execution or fallback
_utils.switch_backend_to_pytorch(self, module)
# Exceptions subject to fallback are handled here
self._fallback_manager.handle_exception(exception=e,
log_level=debug_options.logging.log_level)
except Exception as e:
self._torch_module = TorchModulePytorch(module)
# Although backend is switched to PyTorch here,
# it is up to _FallbackManager to actually terminate execution or fallback
_utils.switch_backend_to_pytorch(self, module)
# Catch-all FALLBACK_FORCE_TORCH_FORWARD fallback is handled here
self._fallback_manager.handle_exception(exception=e,
log_level=debug_options.logging.log_level,