diff --git a/orttraining/orttraining/python/training/ortmodule/__init__.py b/orttraining/orttraining/python/training/ortmodule/__init__.py index 0fd1460e2f..97b39ba252 100644 --- a/orttraining/orttraining/python/training/ortmodule/__init__.py +++ b/orttraining/orttraining/python/training/ortmodule/__init__.py @@ -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.")) diff --git a/orttraining/orttraining/python/training/ortmodule/_fallback.py b/orttraining/orttraining/python/training/ortmodule/_fallback.py index 4c316e68b4..832c40526d 100644 --- a/orttraining/orttraining/python/training/ortmodule/_fallback.py +++ b/orttraining/orttraining/python/training/ortmodule/_fallback.py @@ -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) diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index d4e2e59caa..b1f0184ea3 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -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""" diff --git a/orttraining/orttraining/python/training/ortmodule/_torch_module_pytorch.py b/orttraining/orttraining/python/training/ortmodule/_torch_module_pytorch.py index c1f6cc7192..29066738f3 100644 --- a/orttraining/orttraining/python/training/ortmodule/_torch_module_pytorch.py +++ b/orttraining/orttraining/python/training/ortmodule/_torch_module_pytorch.py @@ -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 diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index b5ef430dbf..f3f80983a7 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -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""" diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py index 1cf035c37b..8edfae3419 100644 --- a/orttraining/orttraining/python/training/ortmodule/_utils.py +++ b/orttraining/orttraining/python/training/ortmodule/_utils.py @@ -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 diff --git a/orttraining/orttraining/python/training/ortmodule/ortmodule.py b/orttraining/orttraining/python/training/ortmodule/ortmodule.py index a642261600..1c92ffcaaf 100644 --- a/orttraining/orttraining/python/training/ortmodule/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule/ortmodule.py @@ -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,