mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Consolidate ORTModule logging (#16078)
### Consolidate ORTModule logging There are few improvements for ORTModule loggings: - All ORTModule logging are used logger that is initialized in `ortmodule.py`. - Manage all export logs same way, e.g. use ` _logger.suppress_os_stream_output(log_level=self._debug_options.logging.log_level)` to control exporting related logs suppressing or not. If any warning or errors suppressed, `self._warning_log_detected_during_export` will be set to True, then when we log ORTModule feature matrix, we will also told users there are logs suppressed. - Downgrade some warnings. We had some warnings for years, and looks many models have them by default, no action we actually can take, so downgrade them to make user logging cleaner. - PyTorch export requires update of custom export function signature changes, otherwise, _symbolic_context_handler complains with warnings, so update custom export function adaption for version >=1.13 PyTorch. - Add ORTModule feature matrix summary, **this is supposed to be only places users see our logs by default** (unless they use INFO or VERBOSE). Features ON/OFF states are shown clearly to them in case they want to try some features in OFF states. This logs only shows up in rank 0 (if there are multiple rank), the intention is we want user to see a useful and clean output from ORTModule by default. The outputs shown as below:   - `reinitialize_ortmodule` in util.py is only used by ortmodule.py, moving it into ortmodule.py, then utils takes no dependency on `orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py`, then `_custom_op_symbolic_registry.py` can call functions defined in utils.py (without recursively include).
This commit is contained in:
pengwa
GitHub
parent
d19e5c0abb
commit
65b316a138
16 changed files with 400 additions and 295 deletions
|
|
@ -14,9 +14,9 @@ from torch.onnx import symbolic_helper
|
|||
from onnxruntime.capi._pybind_state import register_miscellaneous_const_input, register_torch_autograd_function
|
||||
from onnxruntime.training import ortmodule
|
||||
|
||||
from . import _logger
|
||||
from ._custom_op_symbolic_registry import pytorch_type_to_onnx
|
||||
from ._custom_op_symbolic_registry import pytorch_type_to_onnx, wrap_custom_export_function
|
||||
from ._fallback import ORTModuleONNXModelException, wrap_exception
|
||||
from ._utils import get_runtime_pytorch_version
|
||||
|
||||
# Some autograd.Function's shouldn't be exported as PythonOp.
|
||||
# If CheckpointFunction is exported as PythonOp, the checkpointed computation
|
||||
|
|
@ -26,28 +26,6 @@ from ._fallback import ORTModuleONNXModelException, wrap_exception
|
|||
# at all.
|
||||
BANNED_AUTOGRAD_FUNCTION_NAMES = frozenset([torch.utils.checkpoint.CheckpointFunction.__name__])
|
||||
|
||||
# Mapping from pytorch scalar type to onnx scalar type.
|
||||
_CAST_PYTORCH_TO_ONNX = {
|
||||
"Byte": torch.onnx.TensorProtoDataType.UINT8,
|
||||
"Char": torch.onnx.TensorProtoDataType.INT8,
|
||||
"Double": torch.onnx.TensorProtoDataType.DOUBLE,
|
||||
"Float": torch.onnx.TensorProtoDataType.FLOAT,
|
||||
"Half": torch.onnx.TensorProtoDataType.FLOAT16,
|
||||
"Int": torch.onnx.TensorProtoDataType.INT32,
|
||||
"Long": torch.onnx.TensorProtoDataType.INT64,
|
||||
"Short": torch.onnx.TensorProtoDataType.INT16,
|
||||
"Bool": torch.onnx.TensorProtoDataType.BOOL,
|
||||
"ComplexFloat": torch.onnx.TensorProtoDataType.COMPLEX64,
|
||||
"ComplexDouble": torch.onnx.TensorProtoDataType.COMPLEX128,
|
||||
"BFloat16": torch.onnx.TensorProtoDataType.BFLOAT16,
|
||||
# Not yet defined in torch.
|
||||
# "Float8E4M3FN": torch.onnx.TensorProtoDataType.FLOAT8E4M3FN,
|
||||
# "Float8E4M3FNUZ": torch.onnx.TensorProtoDataType.FLOAT8E4M3FNUZ,
|
||||
# "Float8E5M2": torch.onnx.TensorProtoDataType.FLOAT8E5M2,
|
||||
# "Float8E5M2FNUZ": torch.onnx.TensorProtoDataType.FLOAT8E5M2FNUZ,
|
||||
"Undefined": torch.onnx.TensorProtoDataType.UNDEFINED,
|
||||
}
|
||||
|
||||
|
||||
def _full_name(klass):
|
||||
module = klass.__module__
|
||||
|
|
@ -56,13 +34,6 @@ def _full_name(klass):
|
|||
return module + "." + klass.__qualname__
|
||||
|
||||
|
||||
def pytorch_type_to_onnx(scalar_type: str) -> torch.onnx.TensorProtoDataType: # noqa: F811
|
||||
try:
|
||||
return torch.onnx.JitScalarType.from_name(scalar_type).onnx_type()
|
||||
except AttributeError:
|
||||
return _CAST_PYTORCH_TO_ONNX[scalar_type]
|
||||
|
||||
|
||||
def _export_pt_1_10(g, n, *args, **kwargs):
|
||||
"""
|
||||
This function exports PythonOp (input: "n") into a graph
|
||||
|
|
@ -83,8 +54,7 @@ def _export_pt_1_10(g, n, *args, **kwargs):
|
|||
inplace = kwargs["inplace"]
|
||||
# TODO move to public API once exporter team exposes that
|
||||
training_mode = None
|
||||
runtime_pytorch_version = version.parse(torch.__version__.split("+")[0])
|
||||
if runtime_pytorch_version >= version.parse("1.12"):
|
||||
if get_runtime_pytorch_version() >= version.parse("1.12"):
|
||||
# FIXME: using privated modules
|
||||
from torch.onnx import _globals
|
||||
|
||||
|
|
@ -228,22 +198,10 @@ def _export_pt_1_10(g, n, *args, **kwargs):
|
|||
raise wrap_exception(ORTModuleONNXModelException, e) # noqa: B904
|
||||
|
||||
|
||||
# Starting from PyTorch 1.11, there has been a change to symbolic function signature
|
||||
# in terms of how additional context is accessed. More info at
|
||||
# https://github.com/pytorch/pytorch/blob/6b02648479d3615fa3260961e24f38dd0f22da94/torch/onnx/symbolic_helper.py#L48
|
||||
# This code can be cleaned up once support for PyTorch version < 1.11 is dropped.
|
||||
try:
|
||||
from torch.onnx import SymbolicContext
|
||||
|
||||
def _export(ctx: SymbolicContext, g, *args, **kwargs):
|
||||
n = ctx.cur_node
|
||||
return _export_pt_1_10(g, n, *args, **kwargs)
|
||||
|
||||
except ImportError:
|
||||
_export = _export_pt_1_10
|
||||
_export = wrap_custom_export_function(_export_pt_1_10)
|
||||
|
||||
|
||||
def _post_process_after_export(exported_model, enable_custom_autograd_function, log_level):
|
||||
def _post_process_after_export(exported_model, enable_custom_autograd_function):
|
||||
if enable_custom_autograd_function:
|
||||
return _post_process_enabling_autograd_fallback(exported_model)
|
||||
|
||||
|
|
@ -253,7 +211,7 @@ def _post_process_after_export(exported_model, enable_custom_autograd_function,
|
|||
is_pythonop_needed = True
|
||||
break
|
||||
|
||||
if is_pythonop_needed and log_level <= _logger.LogLevel.WARNING:
|
||||
if is_pythonop_needed:
|
||||
warnings.warn(
|
||||
"Detected autograd functions usage in current model, the run will fail \
|
||||
without enabling '_enable_custom_autograd_function'. Please enable it with: \
|
||||
|
|
|
|||
|
|
@ -3,16 +3,19 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import warnings # noqa: F401
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.onnx.symbolic_helper as sym_help
|
||||
from packaging import version
|
||||
from packaging.version import Version
|
||||
from torch.onnx import register_custom_op_symbolic
|
||||
from torch.onnx.symbolic_helper import _get_tensor_dim_size, _get_tensor_sizes, parse_args
|
||||
|
||||
from onnxruntime.training import ortmodule
|
||||
|
||||
from ._utils import get_runtime_pytorch_version
|
||||
|
||||
# Mapping from pytorch scalar type to onnx scalar type.
|
||||
_CAST_PYTORCH_TO_ONNX = {
|
||||
"Byte": torch.onnx.TensorProtoDataType.UINT8,
|
||||
|
|
@ -43,12 +46,34 @@ def pytorch_type_to_onnx(scalar_type: str) -> torch.onnx.TensorProtoDataType:
|
|||
return _CAST_PYTORCH_TO_ONNX[scalar_type]
|
||||
|
||||
|
||||
def wrap_custom_export_function(original_func):
|
||||
# Starting from PyTorch 1.11, there has been a change to symbolic function signature
|
||||
# in terms of how additional context is accessed. More info at
|
||||
# https://github.com/pytorch/pytorch/blob/6b02648479d3615fa3260961e24f38dd0f22da94/torch/onnx/symbolic_helper.py#L48
|
||||
# This code can be cleaned up once support for PyTorch version < 1.11 is dropped.
|
||||
try:
|
||||
def wrap_custom_export_function(original_func: Callable) -> Callable:
|
||||
"""This function is to wrap the custom export function to make sure it can be used by different versions of PyTorch.
|
||||
|
||||
Args:
|
||||
original_func: The original custom export function.
|
||||
|
||||
Note1:
|
||||
[PyTorch exporter breaking change] Starting from PyTorch 1.11, there has been a change to symbolic function
|
||||
signature in terms of how additional context is accessed. More info at
|
||||
https://github.com/pytorch/pytorch/blob/6b02648479d3615fa3260961e24f38dd0f22da94/torch/onnx/symbolic_helper.py#L48
|
||||
This code can be cleaned up once support for PyTorch version < 1.11 is dropped.
|
||||
Note2:
|
||||
[PyTorch exporter breaking change] Custom export function's first argument is SymbolicContext since 1.11, but
|
||||
is changed later, and will be deprecated in 1.13 as claimed. So we need to use GraphContext as the first
|
||||
argument instead.
|
||||
|
||||
"""
|
||||
runtime_pytorch_version = get_runtime_pytorch_version()
|
||||
|
||||
if runtime_pytorch_version >= version.parse("1.13"):
|
||||
from torch.onnx._internal import jit_utils
|
||||
|
||||
def _export_with_ctx(graph_context: jit_utils.GraphContext, *args, **kwargs):
|
||||
return original_func(graph_context, graph_context.original_node, *args, **kwargs)
|
||||
|
||||
return _export_with_ctx
|
||||
|
||||
elif runtime_pytorch_version >= version.parse("1.11"):
|
||||
from torch.onnx import SymbolicContext
|
||||
|
||||
def _export_with_ctx(ctx: SymbolicContext, graph, *args, **kwargs):
|
||||
|
|
@ -56,7 +81,7 @@ def wrap_custom_export_function(original_func):
|
|||
return original_func(graph, node, *args, **kwargs)
|
||||
|
||||
return _export_with_ctx
|
||||
except ImportError:
|
||||
else:
|
||||
|
||||
def _export_with_no_ctx(graph, *args, **kwargs):
|
||||
return original_func(graph, None, *args, **kwargs)
|
||||
|
|
|
|||
|
|
@ -4,8 +4,8 @@
|
|||
# --------------------------------------------------------------------------
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from enum import IntFlag
|
||||
from logging import Logger
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
|
@ -67,7 +67,7 @@ class _FallbackManager:
|
|||
be raised to the user, terminating execution
|
||||
"""
|
||||
|
||||
def __init__(self, pytorch_module: torch.nn.Module, policy: _FallbackPolicy, retry: bool):
|
||||
def __init__(self, pytorch_module: torch.nn.Module, policy: _FallbackPolicy, retry: bool, logger: Logger):
|
||||
self._original_module = pytorch_module
|
||||
|
||||
# Read policy from environment variable for testing purposes
|
||||
|
|
@ -103,6 +103,7 @@ class _FallbackManager:
|
|||
self.retry = retry
|
||||
self._exception = None
|
||||
self._raised_fallback_exception = False
|
||||
self._logger = logger
|
||||
|
||||
def handle_exception(
|
||||
self, exception: Exception, log_level: _logger.LogLevel, override_policy: Optional[_FallbackPolicy] = None
|
||||
|
|
@ -132,19 +133,15 @@ class _FallbackManager:
|
|||
and type(exception) in self._policy_exception_map[policy.value]
|
||||
)
|
||||
):
|
||||
if log_level <= _logger.LogLevel.INFO:
|
||||
warnings.warn(f"Fallback for policy {policy.name} is pending.", UserWarning)
|
||||
self._logger.info(f"Fallback for policy {policy.name} is pending.")
|
||||
|
||||
# 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,
|
||||
if policy == _FallbackPolicy.FALLBACK_BAD_INITIALIZATION:
|
||||
self._logger.warning(
|
||||
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)}"
|
||||
)
|
||||
|
||||
self._exception = exception
|
||||
|
|
@ -177,18 +174,15 @@ class _FallbackManager:
|
|||
exception_string = _utils.get_exception_as_string(self._exception)
|
||||
|
||||
# This warning will not be raised again if retry is not enabled
|
||||
warnings.warn(
|
||||
(
|
||||
"Fallback to PyTorch due to exception {} was triggered. "
|
||||
"Report this issue with a minimal repro at https://www.github.com/microsoft/onnxruntime. "
|
||||
"See details below:\n\n{}".format(exception_type, exception_string)
|
||||
),
|
||||
UserWarning,
|
||||
self._logger.warning(
|
||||
"Fallback to PyTorch due to exception {} was triggered. "
|
||||
"Report this issue with a minimal repro at https://www.github.com/microsoft/onnxruntime. "
|
||||
"See details below:\n\n{}".format(exception_type, exception_string)
|
||||
)
|
||||
|
||||
self._raised_fallback_exception = True
|
||||
|
||||
# Pending fallbacks are resetted to enforce retries
|
||||
# Pending fallbacks are reset to enforce retries
|
||||
if self.retry:
|
||||
self._raised_fallback_exception = False
|
||||
self._exception = None
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ import inspect
|
|||
import io
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod # noqa: F401
|
||||
from enum import IntFlag
|
||||
from functools import reduce
|
||||
|
|
@ -30,6 +29,7 @@ from ._fallback import (
|
|||
ORTModuleONNXModelException,
|
||||
ORTModuleTorchModelException,
|
||||
_FallbackManager,
|
||||
_FallbackPolicy,
|
||||
wrap_exception,
|
||||
)
|
||||
from ._gradient_accumulation_manager import GradientAccumulationManager
|
||||
|
|
@ -38,8 +38,6 @@ from ._io import _FlattenedModule, _InputInfo, _ModelInputOutputSchemaType
|
|||
from .debug_options import DebugOptions, LogLevel
|
||||
from .torch_cpp_extensions.cpu.aten_op_executor import load_aten_op_executor_cpp_extension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _RunStateInfo:
|
||||
def __init__(self, state, output_info: List[Tuple[torch.Size, torch.device, torch.dtype]]):
|
||||
|
|
@ -74,7 +72,13 @@ class _SkipCheck(IntFlag):
|
|||
|
||||
|
||||
class GraphExecutionManager(GraphExecutionInterface):
|
||||
def __init__(self, module: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
def __init__(
|
||||
self,
|
||||
module: _FlattenedModule,
|
||||
debug_options: DebugOptions,
|
||||
fallback_manager: _FallbackManager,
|
||||
logger: logging.Logger,
|
||||
):
|
||||
"""Manages construction and execution of ONNX graphs"""
|
||||
|
||||
super().__init__(module._original_module)
|
||||
|
|
@ -82,7 +86,8 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# IMPORTANT: Debug and Fallback must the configured first
|
||||
self._debug_options = debug_options
|
||||
self._fallback_manager = fallback_manager
|
||||
logger.setLevel(_logger.ortmodule_loglevel_to_python_loglevel(self._debug_options.logging.log_level))
|
||||
|
||||
self._logger = logger
|
||||
|
||||
# Original and flattened (transformed) output module
|
||||
self._flattened_module = module
|
||||
|
|
@ -119,7 +124,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
self._first_skip_check_warning = True
|
||||
|
||||
# Inspect embedding input index sparsity.
|
||||
self._rt_inspector = _runtime_inspector.RuntimeInspector(logger)
|
||||
self._rt_inspector = _runtime_inspector.RuntimeInspector(self._logger)
|
||||
|
||||
# Graph transformer config
|
||||
# Specify cast propagation strategy. Currently, three strategies are available, NONE, INSERT-AND-REDUCE and FLOOD-FILL
|
||||
|
|
@ -160,6 +165,8 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# Input and output infos (including schema) for exported model.
|
||||
self._input_info: Optional[_InputInfo] = None
|
||||
self._module_output_schema: Optional[_ModelInputOutputSchemaType] = None
|
||||
self._warning_log_detected_during_export = False
|
||||
self._export_duration_in_ms = 0
|
||||
|
||||
# Device where the model is placed.
|
||||
self._device: Optional[torch.device] = _utils.get_device_from_module(module)
|
||||
|
|
@ -172,8 +179,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# TODO: remove after PyTorch ONNX exporter supports VAR_KEYWORD parameters.
|
||||
for input_parameter in self._module_parameters:
|
||||
if input_parameter.kind == inspect.Parameter.VAR_KEYWORD:
|
||||
if self._debug_options.logging.log_level <= LogLevel.WARNING:
|
||||
logger.warning("The model's forward method has **kwargs parameter which has EXPERIMENTAL support!")
|
||||
self._logger.info("The model's forward method has **kwargs parameter which has EXPERIMENTAL support!")
|
||||
|
||||
self.is_rocm_pytorch = bool(torch.version.hip is not None and ROCM_HOME is not None)
|
||||
|
||||
|
|
@ -200,6 +206,8 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
|
||||
self._print_input_density = ortmodule._defined_from_envvar("ORTMODULE_PRINT_INPUT_DENSITY", 0, warn=True) == 1
|
||||
|
||||
self._enable_memory_optimizer = ortmodule._defined_from_envvar("ORTMODULE_MEMORY_OPT_CONFIG", "", warn=True)
|
||||
|
||||
# Flag to re-export the model due to attribute change on the original module.
|
||||
# Re-export will be avoided if _skip_check is enabled.
|
||||
self._original_model_has_changed = False
|
||||
|
|
@ -207,6 +215,21 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# Load ATen operator executor extension.
|
||||
load_aten_op_executor_cpp_extension()
|
||||
|
||||
self._feature_map: List[List[str]] = [
|
||||
["ATen Executor", "ON", "Dispatch ATen operators to ORT's ATen executor"],
|
||||
[
|
||||
"Cast Propagation",
|
||||
"ON" if self._propagate_cast_ops_level > 0 else "OFF",
|
||||
f"Level {self._propagate_cast_ops_level} enabled",
|
||||
],
|
||||
["Custom Function", "ON", "Support custom torch.autograd.Function export and execution"],
|
||||
[
|
||||
"Memory Optimizer",
|
||||
"ON" if self._enable_memory_optimizer else "OFF",
|
||||
"Enable with env ORTMODULE_MEMORY_OPT_CONFIG=<config>",
|
||||
],
|
||||
]
|
||||
|
||||
def _get_torch_gpu_allocator_function_addresses(self):
|
||||
if self._use_external_gpu_allocator and torch.cuda.is_available():
|
||||
# CPP extension to get torch GPU allocator's alloc and free function addresses
|
||||
|
|
@ -256,8 +279,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
"""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:
|
||||
logger.warning("ORTModule's determinism will be enabled because PyTorch's determinism is enabled.")
|
||||
self._logger.info("ORTModule's determinism will be enabled because PyTorch's determinism is enabled.")
|
||||
|
||||
providers = None
|
||||
provider_options = None
|
||||
|
|
@ -270,7 +292,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# Set Conv algo search mode to HEURISTIC by default, which is the same as PyTorch's default setting.
|
||||
conv_algo_search = ortmodule._defined_from_envvar("ORTMODULE_CONV_ALGO_SEARCH", "HEURISTIC", warn=True)
|
||||
if conv_algo_search not in ["HEURISTIC", "EXHAUSTIVE"]:
|
||||
warnings.warn("Invalid value of env CONV_ALGO_SEARCH. Must be HEURISTIC or EXHAUSTIVE.")
|
||||
self._logger.warning("Invalid value of env CONV_ALGO_SEARCH. Must be HEURISTIC or EXHAUSTIVE.")
|
||||
conv_algo_search = "HEURISTIC"
|
||||
provider_option_map["cudnn_conv_algo_search"] = conv_algo_search
|
||||
provider_option_map["cudnn_conv_use_max_workspace"] = "1"
|
||||
|
|
@ -331,7 +353,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# e.g., some sympy functions in symbolic_shape_infer will change Python's random state.
|
||||
random_states = _utils.get_random_states()
|
||||
|
||||
schema = _io._extract_schema({"args": copy.copy(inputs), "kwargs": copy.copy(kwargs)})
|
||||
schema = _io._extract_schema({"args": copy.copy(inputs), "kwargs": copy.copy(kwargs)}, self._logger)
|
||||
if (
|
||||
self._onnx_models.exported_model
|
||||
and schema == self._input_info.schema
|
||||
|
|
@ -365,72 +387,84 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
|
||||
TODO: How to support dynamic axes? Dimensions are determined by samples
|
||||
"""
|
||||
with _logger.suppress_os_stream_output(log_level=self._debug_options.logging.log_level) as suppress_output:
|
||||
from datetime import datetime
|
||||
|
||||
# Setup dynamic axes for onnx model
|
||||
self._input_info = _io.parse_inputs_for_onnx_export(self._module_parameters, None, input_schema, inputs, kwargs)
|
||||
(
|
||||
output_names,
|
||||
output_dynamic_axes,
|
||||
self._module_output_schema,
|
||||
) = _io.parse_outputs_for_onnx_export_and_extract_schema(
|
||||
self._original_module, inputs, kwargs, self._debug_options.logging.log_level
|
||||
)
|
||||
self._input_info.dynamic_axes.update(output_dynamic_axes)
|
||||
start = datetime.now()
|
||||
|
||||
# FlattenedModule needs _InputInfo to expand user input from *args to *args + **kwargs
|
||||
self._flattened_module._input_info = self._input_info
|
||||
|
||||
# Export torch.nn.Module to ONNX
|
||||
f = io.BytesIO()
|
||||
|
||||
# Deepcopy inputs, since input values may change after model run.
|
||||
# NOTE: Inputs may contain tensors that have attributes preventing their deepcopy (example grad_fn).
|
||||
# Therefore, deepcopy only the data component of the input tensors for export.
|
||||
sample_inputs_copy, sample_kwargs_copy = _io.deepcopy_model_input(*inputs, **kwargs)
|
||||
# NOTE: Flattening the input will change the 'input schema', resulting in a re-export
|
||||
sample_inputs_as_tuple = tuple(self._input_info.flatten(sample_inputs_copy, sample_kwargs_copy, self._device))
|
||||
# Ops behaving differently under train/eval mode need to be exported with the
|
||||
# correct training flag to reflect the expected behavior.
|
||||
# For example, the Dropout node in a model is dropped under eval mode.
|
||||
assert self._export_mode is not None, "Please use a concrete instance of ExecutionManager"
|
||||
|
||||
try:
|
||||
with torch.no_grad(), _logger.suppress_os_stream_output(log_level=self._debug_options.logging.log_level):
|
||||
required_export_kwargs = {
|
||||
"input_names": self._input_info.names,
|
||||
"output_names": output_names,
|
||||
"opset_version": ortmodule.ONNX_OPSET_VERSION,
|
||||
"do_constant_folding": False,
|
||||
"training": self._export_mode,
|
||||
"dynamic_axes": self._input_info.dynamic_axes,
|
||||
"verbose": self._debug_options.logging.log_level < LogLevel.WARNING,
|
||||
"export_params": False,
|
||||
"keep_initializers_as_inputs": True,
|
||||
}
|
||||
invalid_args = self._export_extra_kwargs.keys() & required_export_kwargs.keys()
|
||||
assert (
|
||||
len(invalid_args) == 0
|
||||
), f"The following PyTorch exporter arguments cannot be specified: '{invalid_args}'."
|
||||
torch.onnx.export(
|
||||
self._flattened_module,
|
||||
sample_inputs_as_tuple,
|
||||
f,
|
||||
**required_export_kwargs,
|
||||
**self._export_extra_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
raise wrap_exception( # noqa: B904
|
||||
ORTModuleONNXModelException,
|
||||
RuntimeError(
|
||||
f"There was an error while exporting the PyTorch model to ONNX: "
|
||||
f"\n\n{_utils.get_exception_as_string(e)}"
|
||||
),
|
||||
# Setup dynamic axes for onnx model
|
||||
self._input_info = _io.parse_inputs_for_onnx_export(
|
||||
self._module_parameters, None, input_schema, inputs, kwargs
|
||||
)
|
||||
exported_model = onnx.load_model_from_string(f.getvalue())
|
||||
(
|
||||
output_names,
|
||||
output_dynamic_axes,
|
||||
self._module_output_schema,
|
||||
) = _io.parse_outputs_for_onnx_export_and_extract_schema(
|
||||
self._original_module, inputs, kwargs, self._logger
|
||||
)
|
||||
self._input_info.dynamic_axes.update(output_dynamic_axes)
|
||||
|
||||
exported_model = _post_process_after_export(
|
||||
exported_model, self._enable_custom_autograd_function, self._debug_options.logging.log_level
|
||||
)
|
||||
# FlattenedModule needs _InputInfo to expand user input from *args to *args + **kwargs
|
||||
self._flattened_module._input_info = self._input_info
|
||||
|
||||
# Export torch.nn.Module to ONNX
|
||||
f = io.BytesIO()
|
||||
|
||||
# Deepcopy inputs, since input values may change after model run.
|
||||
# NOTE: Inputs may contain tensors that have attributes preventing their deepcopy (example grad_fn).
|
||||
# Therefore, deepcopy only the data component of the input tensors for export.
|
||||
sample_inputs_copy, sample_kwargs_copy = _io.deepcopy_model_input(*inputs, **kwargs)
|
||||
# NOTE: Flattening the input will change the 'input schema', resulting in a re-export
|
||||
sample_inputs_as_tuple = tuple(
|
||||
self._input_info.flatten(sample_inputs_copy, sample_kwargs_copy, self._device)
|
||||
)
|
||||
# Ops behaving differently under train/eval mode need to be exported with the
|
||||
# correct training flag to reflect the expected behavior.
|
||||
# For example, the Dropout node in a model is dropped under eval mode.
|
||||
assert self._export_mode is not None, "Please use a concrete instance of ExecutionManager"
|
||||
|
||||
try:
|
||||
with torch.no_grad():
|
||||
required_export_kwargs = {
|
||||
"input_names": self._input_info.names,
|
||||
"output_names": output_names,
|
||||
"opset_version": ortmodule.ONNX_OPSET_VERSION,
|
||||
"do_constant_folding": False,
|
||||
"training": self._export_mode,
|
||||
"dynamic_axes": self._input_info.dynamic_axes,
|
||||
"verbose": self._debug_options.logging.log_level < LogLevel.WARNING,
|
||||
"export_params": False,
|
||||
"keep_initializers_as_inputs": True,
|
||||
}
|
||||
invalid_args = self._export_extra_kwargs.keys() & required_export_kwargs.keys()
|
||||
assert (
|
||||
len(invalid_args) == 0
|
||||
), f"The following PyTorch exporter arguments cannot be specified: '{invalid_args}'."
|
||||
torch.onnx.export(
|
||||
self._flattened_module,
|
||||
sample_inputs_as_tuple,
|
||||
f,
|
||||
**required_export_kwargs,
|
||||
**self._export_extra_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
raise wrap_exception( # noqa: B904
|
||||
ORTModuleONNXModelException,
|
||||
RuntimeError(
|
||||
f"There was an error while exporting the PyTorch model to ONNX: "
|
||||
f"\n\n{_utils.get_exception_as_string(e)}"
|
||||
),
|
||||
)
|
||||
exported_model = onnx.load_model_from_string(f.getvalue())
|
||||
|
||||
exported_model = _post_process_after_export(exported_model, self._enable_custom_autograd_function)
|
||||
|
||||
if suppress_output.tell() > 0:
|
||||
self._warning_log_detected_during_export = True
|
||||
|
||||
end = datetime.now()
|
||||
self._export_duration_in_ms = (end - start).total_seconds() * 1000
|
||||
|
||||
return exported_model
|
||||
|
||||
|
|
@ -532,14 +566,27 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
Based on runtime inspection, enable conditional optimizations if applicable.
|
||||
|
||||
Input sparsity-based optimization workflows:
|
||||
1. Input density observer is enabled when label sparsity optimization is enabled.
|
||||
1. Input density observer is enabled if the sparse optimizer is ON or user wants to print input density.
|
||||
2. Input density observer inspects input tensors and returns sparsity results.
|
||||
3. If label or embedding input sparsity is found in sparsity results, graph transformer config is updated to
|
||||
enable sparsity-based optimization.
|
||||
|
||||
"""
|
||||
|
||||
# Enable data sparsity inspection if label sparsity optimization is enabled or user wants to print input density.
|
||||
self._feature_map.extend(
|
||||
[
|
||||
[
|
||||
"Compute Optimizer",
|
||||
"ON" if self._enable_compute_optimizer else "OFF",
|
||||
"Enable/Disable with env ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=1/0",
|
||||
],
|
||||
[
|
||||
" -FLOPReduction",
|
||||
"ON" if self._enable_compute_optimizer else "OFF",
|
||||
"Enable/Disable with env ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=1/0",
|
||||
],
|
||||
]
|
||||
)
|
||||
# Enable data sparsity inspection if sparse optimizer is ON or user wants to print input density.
|
||||
if self._enable_sparse_optimizer or self._print_input_density:
|
||||
self._rt_inspector.enable_input_inspector(
|
||||
self._onnx_models.exported_model, self._graph_builder.get_graph_info().user_input_names
|
||||
|
|
@ -563,14 +610,64 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
|
||||
# Enable sparsity-based optimization when applicable.
|
||||
if len(label_sparsity_results) > 0:
|
||||
graph_transformer_config.sparse_label_input_names = label_sparsity_results
|
||||
logger.info(f"Label sparsity based optimization is on for {label_sparsity_results}")
|
||||
graph_transformer_config.sparse_label_input_names = list(label_sparsity_results.keys())
|
||||
self._logger.info("Label sparsity-based optimization is ON for %s", label_sparsity_results)
|
||||
sparsity_stat_str = ",".join([f"{k}:{v:.0f}%" for k, v in label_sparsity_results.items()])
|
||||
self._feature_map.append(
|
||||
[
|
||||
" -LabelSparsityOpt",
|
||||
"ON",
|
||||
f"Input density: {sparsity_stat_str}, switch: ORTMODULE_ENABLE_SPARSE_OPTIMIZER=1/0",
|
||||
]
|
||||
)
|
||||
|
||||
if len(embed_sparsity_results) > 0:
|
||||
graph_transformer_config.sparse_embedding_input_names = embed_sparsity_results
|
||||
logger.info(f"Embedding sparsity based optimization is on for {embed_sparsity_results}")
|
||||
graph_transformer_config.sparse_embedding_input_names = list(embed_sparsity_results.keys())
|
||||
self._logger.info("Embedding sparsity-based optimization is ON for %s", embed_sparsity_results)
|
||||
sparsity_stat_str = ",".join([f"{k}:{v:.0f}%" for k, v in embed_sparsity_results.items()])
|
||||
self._feature_map.append(
|
||||
[
|
||||
" -EmbedSparsityOpt",
|
||||
"ON",
|
||||
f"Input density: {sparsity_stat_str}, switch: ORTMODULE_ENABLE_SPARSE_OPTIMIZER=1/0",
|
||||
]
|
||||
)
|
||||
|
||||
# If users don't want to print input density, disable the input density observer to avoid overhead
|
||||
# when looping through inputs during training.
|
||||
if not self._print_input_density:
|
||||
self._rt_inspector.disable_input_inspector()
|
||||
|
||||
def _log_feature_stats(self):
|
||||
rank = 0
|
||||
if torch.distributed.is_initialized():
|
||||
rank = torch.distributed.get_rank()
|
||||
|
||||
if rank != 0:
|
||||
return
|
||||
|
||||
self._feature_map.append(
|
||||
[
|
||||
"Auto Fallback",
|
||||
"ON" if self._fallback_manager.policy is not _FallbackPolicy.FALLBACK_DISABLE else "OFF",
|
||||
"Fallback to PyTorch when encountering unsupported ops",
|
||||
]
|
||||
)
|
||||
|
||||
mode = "training" if self._export_mode == torch.onnx.TrainingMode.TRAINING else "inference"
|
||||
mode = f"{_logger.LogColor.UNDERLINE}{mode}{_logger.LogColor.ENDC}"
|
||||
|
||||
stat = f"\n\n{_logger.LogColor.HEADER}***** ONNX Runtime Training (ORTModule) is accelerating your model *****{_logger.LogColor.ENDC}\n\n"
|
||||
stat += f"ORTModule is enabled with following features ON/OFF for [{mode}] mode:\n\n"
|
||||
for feature_tuple in self._feature_map:
|
||||
stat += f"{feature_tuple[0]:<20}:\t{feature_tuple[1]:<10}:\t{feature_tuple[2]:<80}\n"
|
||||
|
||||
# If anything was captured in fo, raise a single user warning letting users know that there was
|
||||
# any warning or error that was raised
|
||||
stat += f"\n{_logger.LogColor.WARNING}There were one or more warnings or errors raised while exporting the PyTorch model.\n"
|
||||
stat += f"Please enable INFO level logging with DebugOptions to view all warnings and errors.{_logger.LogColor.ENDC}\n\n"
|
||||
stat += f"Export duration: {self._export_duration_in_ms:.0f} milliseconds\n"
|
||||
stat += f"Versions: ONNX Runtime - {onnxruntime.__version__}, ONNX - {onnx.__version__}\n\n"
|
||||
stat += f"{_logger.LogColor.HEADER}************************************************************************{_logger.LogColor.ENDC}\n\n"
|
||||
|
||||
self._logger.warning(stat)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
from logging import Logger
|
||||
from typing import Union
|
||||
|
||||
from ._fallback import _FallbackManager
|
||||
|
|
@ -13,9 +14,11 @@ from .debug_options import DebugOptions
|
|||
|
||||
|
||||
class GraphExecutionManagerFactory:
|
||||
def __init__(self, module: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
self._training_manager = TrainingManager(module, debug_options, fallback_manager)
|
||||
self._inference_manager = InferenceManager(module, debug_options, fallback_manager)
|
||||
def __init__(
|
||||
self, module: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager, logger: Logger
|
||||
):
|
||||
self._training_manager = TrainingManager(module, debug_options, fallback_manager, logger)
|
||||
self._inference_manager = InferenceManager(module, debug_options, fallback_manager, logger)
|
||||
|
||||
def __call__(self, is_training) -> Union[InferenceManager, TrainingManager]:
|
||||
if is_training:
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import warnings
|
||||
from logging import Logger
|
||||
from typing import Tuple
|
||||
|
||||
import onnx
|
||||
|
|
@ -11,7 +11,7 @@ import torch
|
|||
|
||||
from onnxruntime.capi import _pybind_state as C
|
||||
|
||||
from . import _are_deterministic_algorithms_enabled, _io, _logger, _use_deterministic_algorithms, _utils
|
||||
from . import _are_deterministic_algorithms_enabled, _io, _use_deterministic_algorithms, _utils
|
||||
from ._execution_agent import InferenceAgent
|
||||
from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPolicy
|
||||
from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck
|
||||
|
|
@ -24,8 +24,8 @@ class InferenceManager(GraphExecutionManager):
|
|||
InferenceManager is responsible for building and running the forward graph of the inference model
|
||||
"""
|
||||
|
||||
def __init__(self, model, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
super().__init__(model, debug_options, fallback_manager)
|
||||
def __init__(self, model, debug_options: DebugOptions, fallback_manager: _FallbackManager, logger: Logger):
|
||||
super().__init__(model, debug_options, fallback_manager, logger)
|
||||
self._export_mode = torch.onnx.TrainingMode.EVAL
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -92,18 +92,14 @@ class InferenceManager(GraphExecutionManager):
|
|||
|
||||
try:
|
||||
# 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
|
||||
):
|
||||
if self._first_skip_check_warning is True and self._skip_check.is_disabled() is False:
|
||||
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,
|
||||
self._logger.warning(
|
||||
"Fast path enabled - skipping checks. rebuild gradient graph: %s, execution agent recreation: %s, "
|
||||
"device check: %s",
|
||||
self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT),
|
||||
self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT),
|
||||
self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE),
|
||||
)
|
||||
|
||||
# If exporting module to ONNX for the first time, this skip check will not take effect.
|
||||
|
|
@ -128,6 +124,8 @@ class InferenceManager(GraphExecutionManager):
|
|||
# Build the gradient graph
|
||||
self._build_graph(graph_transformer_config)
|
||||
|
||||
self._log_feature_stats()
|
||||
|
||||
# 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
|
||||
|
|
|
|||
|
|
@ -6,15 +6,14 @@
|
|||
import copy
|
||||
import gc
|
||||
import inspect
|
||||
import warnings
|
||||
from collections import abc
|
||||
from collections import OrderedDict, abc
|
||||
from logging import Logger
|
||||
from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Tuple, Union
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
|
||||
from ._fallback import ORTModuleIOError, ORTModuleONNXModelException, wrap_exception
|
||||
from ._logger import LogLevel, suppress_os_stream_output
|
||||
from ._runtime_inspector import RuntimeInspector
|
||||
from ._utils import warn_of_constant_inputs
|
||||
|
||||
|
|
@ -285,8 +284,8 @@ def _combine_input_buffers_initializers(
|
|||
_expand_inputs(kwargs, flattened_kwargs_inputs)
|
||||
buffer_names_dict = {buffer_name: inp for buffer_name, inp in named_buffer}
|
||||
result = []
|
||||
embed_sparsity_results = []
|
||||
label_sparsity_results = []
|
||||
embed_sparsity_results = OrderedDict()
|
||||
label_sparsity_results = OrderedDict()
|
||||
|
||||
for input_idx, name in enumerate(onnx_input_names):
|
||||
inp = None
|
||||
|
|
@ -319,12 +318,12 @@ def _combine_input_buffers_initializers(
|
|||
if _PrimitiveType.is_primitive_type(inp):
|
||||
inp = _PrimitiveType.get_tensor(inp, device)
|
||||
|
||||
found, embedding_is_sparse, label_is_sparse = rt_inspector.inspect_input(name, inp)
|
||||
found, embedding_density, label_density = rt_inspector.inspect_input(name, inp)
|
||||
if found:
|
||||
if embedding_is_sparse is True:
|
||||
embed_sparsity_results.append(name)
|
||||
if label_is_sparse is True:
|
||||
label_sparsity_results.append(name)
|
||||
if embedding_density < 100:
|
||||
embed_sparsity_results[name] = embedding_density
|
||||
if label_density < 100:
|
||||
label_sparsity_results[name] = label_density
|
||||
result.append(inp)
|
||||
else:
|
||||
raise wrap_exception(
|
||||
|
|
@ -405,7 +404,7 @@ def unflatten_user_output(output_schema: Optional[_ModelInputOutputSchemaType],
|
|||
return user_output
|
||||
|
||||
|
||||
def _extract_schema(data: _ModelInputOutputType) -> _ModelInputOutputSchemaType:
|
||||
def _extract_schema(data: _ModelInputOutputType, logger: Logger) -> _ModelInputOutputSchemaType:
|
||||
"""Extract the data schema by replacing every torch.Tensor value with _TensorStub.
|
||||
|
||||
Depth first traversal to iterate over the data:
|
||||
|
|
@ -413,7 +412,6 @@ def _extract_schema(data: _ModelInputOutputType) -> _ModelInputOutputSchemaType:
|
|||
> Replace None/str typed data with itself
|
||||
> Recreate tensor from data for other primitive types, and replace them with a stub.
|
||||
|
||||
TODO(pengwa): use a unified logger and log level control.
|
||||
Examples:
|
||||
Example 1, list:
|
||||
data = [torch.tensor(1), torch.tensor(2)]
|
||||
|
|
@ -452,11 +450,11 @@ def _extract_schema(data: _ModelInputOutputType) -> _ModelInputOutputSchemaType:
|
|||
if data is None:
|
||||
return data
|
||||
elif isinstance(data, str):
|
||||
warn_of_constant_inputs(data)
|
||||
warn_of_constant_inputs(data, logger)
|
||||
return data
|
||||
elif _PrimitiveType.is_primitive_type(data):
|
||||
if isinstance(data, bool):
|
||||
warn_of_constant_inputs(data)
|
||||
warn_of_constant_inputs(data, logger)
|
||||
return _TensorStub(dtype=_PrimitiveType.get_primitive_dtype(data), shape_dims=0)
|
||||
# Depth first traversal to iterate over the data to replace every tensor with a stub
|
||||
elif isinstance(data, torch.Tensor):
|
||||
|
|
@ -467,7 +465,7 @@ def _extract_schema(data: _ModelInputOutputType) -> _ModelInputOutputSchemaType:
|
|||
stubbed_schema: Optional[_ModelInputOutputSchemaType] = None
|
||||
if isinstance(data, abc.Sequence):
|
||||
sequence_type = type(data)
|
||||
stubbed_schema = [_extract_schema(val) for val in data]
|
||||
stubbed_schema = [_extract_schema(val, logger) for val in data]
|
||||
try:
|
||||
# namedtuple can be created by passing the list sequence to method _make
|
||||
stubbed_schema = sequence_type._make(stubbed_schema)
|
||||
|
|
@ -476,7 +474,7 @@ def _extract_schema(data: _ModelInputOutputType) -> _ModelInputOutputSchemaType:
|
|||
stubbed_schema = sequence_type(stubbed_schema)
|
||||
elif isinstance(data, abc.Mapping):
|
||||
dict_type = type(data)
|
||||
stubbed_schema = {key: _extract_schema(data[key]) for key in data}
|
||||
stubbed_schema = {key: _extract_schema(data[key], logger) for key in data}
|
||||
stubbed_schema = dict_type(**stubbed_schema)
|
||||
else:
|
||||
raise wrap_exception(
|
||||
|
|
@ -714,13 +712,13 @@ def parse_inputs_for_onnx_export(
|
|||
|
||||
|
||||
def parse_outputs_for_onnx_export_and_extract_schema(
|
||||
module, args: Sequence[_ModelInputOutputType], kwargs: Mapping[str, _ModelInputOutputType], log_level: LogLevel
|
||||
module, args: Sequence[_ModelInputOutputType], kwargs: Mapping[str, _ModelInputOutputType], logger: Logger
|
||||
):
|
||||
# Perform a forward call to grab outputs
|
||||
output_names = None
|
||||
output_dynamic_axes = None
|
||||
is_deepcopy = False
|
||||
with torch.no_grad(), suppress_os_stream_output(log_level=log_level):
|
||||
with torch.no_grad():
|
||||
# Deepcopy inputs, since input values may change after model run.
|
||||
sample_args_copy, sample_kwargs_copy = deepcopy_model_input(*args, **kwargs)
|
||||
try:
|
||||
|
|
@ -729,7 +727,7 @@ def parse_outputs_for_onnx_export_and_extract_schema(
|
|||
is_deepcopy = True
|
||||
except Exception:
|
||||
model_copy = module
|
||||
warnings.warn(
|
||||
logger.warning(
|
||||
"This model cannot be deep copied (or pickled), "
|
||||
"which is a required step for stateful models to be properly exported to ONNX."
|
||||
" Compute will continue, but unexpected results may occur!"
|
||||
|
|
@ -740,7 +738,7 @@ def parse_outputs_for_onnx_export_and_extract_schema(
|
|||
# Parse the output and extract the output_names and output_dynamic_axes to be used for onnx export
|
||||
output_names, output_dynamic_axes = _parse_outputs_and_extract_names_and_dynamic_axes(sample_outputs)
|
||||
|
||||
output_schema = _extract_schema(sample_outputs)
|
||||
output_schema = _extract_schema(sample_outputs, logger)
|
||||
if is_deepcopy:
|
||||
del model_copy
|
||||
gc.collect()
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@
|
|||
import io
|
||||
import logging
|
||||
import sys
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from enum import IntEnum
|
||||
from typing import Dict, List
|
||||
|
|
@ -42,22 +41,13 @@ def suppress_os_stream_output(suppress_stdout=True, suppress_stderr=True, log_le
|
|||
sys.stdout = fo
|
||||
if suppress_stderr and suppress_logs:
|
||||
sys.stderr = fo
|
||||
yield
|
||||
yield fo
|
||||
finally:
|
||||
if suppress_stdout:
|
||||
sys.stdout = stdout
|
||||
if suppress_stderr:
|
||||
sys.stderr = stderr
|
||||
|
||||
if fo.tell() > 0 and suppress_logs:
|
||||
# If anything was captured in fo, raise a single user warning letting users know that there was
|
||||
# some warning or error that was raised
|
||||
warnings.warn(
|
||||
"There were one or more warnings or errors raised while exporting the PyTorch "
|
||||
"model. Please enable INFO level logging to view all warnings and errors.",
|
||||
UserWarning,
|
||||
)
|
||||
|
||||
|
||||
ORTMODULE_LOG_LEVEL_MAP: Dict[LogLevel, List[int]] = {
|
||||
LogLevel.VERBOSE: [Severity.VERBOSE, logging.DEBUG],
|
||||
|
|
@ -74,3 +64,15 @@ def ortmodule_loglevel_to_onnxruntime_c_loglevel(loglevel: LogLevel) -> int:
|
|||
|
||||
def ortmodule_loglevel_to_python_loglevel(loglevel: LogLevel) -> int:
|
||||
return ORTMODULE_LOG_LEVEL_MAP.get(loglevel, [Severity.WARNING, logging.WARNING])[1]
|
||||
|
||||
|
||||
class LogColor:
|
||||
HEADER = "\033[95m"
|
||||
BLUE = "\033[94m"
|
||||
CYAN = "\033[96m"
|
||||
GREEN = "\033[92m"
|
||||
WARNING = "\033[93m"
|
||||
RED = "\033[91m"
|
||||
ENDC = "\033[0m"
|
||||
BOLD = "\033[1m"
|
||||
UNDERLINE = "\033[4m"
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@ class RuntimeInspector:
|
|||
|
||||
return self.input_density_ob.initialize(model, user_input_names)
|
||||
|
||||
def inspect_input(self, input_name, input_data) -> Tuple[bool, bool, bool]:
|
||||
def inspect_input(self, input_name, input_data) -> Tuple[bool, float, float]:
|
||||
"""
|
||||
Inspect input data and print statistics.
|
||||
|
||||
|
|
@ -51,13 +51,13 @@ class RuntimeInspector:
|
|||
Returns:
|
||||
found: Whether the input name is found in `_embedding_graph_input_to_padding_idx_map` and
|
||||
`_loss_label_graph_input_to_ignore_idx_map`.
|
||||
embedding_is_sparse: Whether the input is sparse.
|
||||
label_is_sparse: Whether the input is sparse.
|
||||
embed_input_density: Density for the inspected embedding input if found to be True; otherwise, return 100.
|
||||
label_input_density: Density for the inspected label input if found to be True; otherwise, return 100.
|
||||
"""
|
||||
if self.input_density_ob is not None:
|
||||
return self.input_density_ob.inspect_from_input_data(input_name, input_data)
|
||||
|
||||
return (False, False, False)
|
||||
return (False, 100, 100)
|
||||
|
||||
def disable_input_inspector(self) -> None:
|
||||
"""Disable input density inspector."""
|
||||
|
|
@ -282,7 +282,7 @@ class InputDensityObserver:
|
|||
[ignore_index, label_preprocess_func]
|
||||
)
|
||||
|
||||
def inspect_from_input_data(self, name: str, inp) -> Tuple[bool, bool, bool]:
|
||||
def inspect_from_input_data(self, name: str, inp) -> Tuple[bool, float, float]:
|
||||
"""
|
||||
Inspect input data and print statistics.
|
||||
|
||||
|
|
@ -292,15 +292,15 @@ class InputDensityObserver:
|
|||
Returns:
|
||||
found: Whether the input name is found in `_embedding_graph_input_to_padding_idx_map` and
|
||||
`_loss_label_graph_input_to_ignore_idx_map`.
|
||||
embedding_is_sparse: Whether the input is sparse.
|
||||
label_is_sparse: Whether the input is sparse.
|
||||
embed_input_density: Density for the inspected embedding input if found to be True; otherwise, return 100.
|
||||
label_input_density: Density for the inspected label input if found to be True; otherwise, return 100.
|
||||
"""
|
||||
if not self._is_initialized:
|
||||
return (False, False, False)
|
||||
return (False, 100, 100)
|
||||
|
||||
try:
|
||||
data = inp.clone()
|
||||
found, embedding_is_sparse, label_is_sparse = self._inspect_embed_label_input(name, data)
|
||||
found, embed_input_density, label_input_density = self._inspect_embed_label_input(name, data)
|
||||
if found:
|
||||
self._current_step += 1
|
||||
|
||||
|
|
@ -308,16 +308,15 @@ class InputDensityObserver:
|
|||
self._last_step = self._current_step
|
||||
self._print_embed_label_stats()
|
||||
|
||||
return (found, embedding_is_sparse, label_is_sparse)
|
||||
return (found, embed_input_density, label_input_density)
|
||||
except Exception as e:
|
||||
self._logger.warning(f"Failed to inspect input {name} due to {e}", UserWarning)
|
||||
return (False, False, False)
|
||||
return (False, 100, 100)
|
||||
|
||||
def _inspect_embed_label_input(self, name, data):
|
||||
found = False
|
||||
embedding_is_sparse = False
|
||||
label_is_sparse = False
|
||||
|
||||
min_embed_density = 100
|
||||
min_label_density = 100
|
||||
if (
|
||||
len(self._embedding_graph_input_to_padding_idx_map) > 0
|
||||
and name in self._embedding_graph_input_to_padding_idx_map
|
||||
|
|
@ -330,7 +329,8 @@ class InputDensityObserver:
|
|||
valid_token_per_batch = str(torch.count_nonzero(data - padding_idx, dim=1).tolist())
|
||||
total_token = data.numel()
|
||||
embed_density = float(valid_token) / float(total_token) * 100
|
||||
embedding_is_sparse = embedding_is_sparse or (embed_density < 90)
|
||||
if embed_density < 90:
|
||||
min_embed_density = min(min_embed_density, embed_density)
|
||||
self._stats.append(
|
||||
[
|
||||
self._current_step,
|
||||
|
|
@ -355,7 +355,8 @@ class InputDensityObserver:
|
|||
valid_token = torch.count_nonzero(data_preprocessed - ignore_index)
|
||||
total_token = data_preprocessed.numel()
|
||||
label_density = float(valid_token) / float(total_token) * 100
|
||||
label_is_sparse = label_is_sparse or (label_density < 90)
|
||||
if label_density < 90:
|
||||
min_label_density = min(min_label_density, label_density)
|
||||
self._stats.append(
|
||||
[
|
||||
self._current_step,
|
||||
|
|
@ -370,7 +371,7 @@ class InputDensityObserver:
|
|||
)
|
||||
found = True
|
||||
|
||||
return found, embedding_is_sparse, label_is_sparse
|
||||
return found, min_embed_density, min_label_density
|
||||
|
||||
def _print_embed_label_stats(self):
|
||||
if len(self._stats) > 0:
|
||||
|
|
|
|||
|
|
@ -2,13 +2,15 @@
|
|||
# Licensed under the MIT License.
|
||||
# _torch_module_factory.py
|
||||
|
||||
from logging import Logger
|
||||
|
||||
from ._fallback import _FallbackManager
|
||||
from ._torch_module_ort import TorchModuleORT
|
||||
from .debug_options import DebugOptions
|
||||
|
||||
|
||||
class TorchModuleFactory:
|
||||
def __call__(self, module, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
def __call__(self, module, debug_options: DebugOptions, fallback_manager: _FallbackManager, logger: Logger):
|
||||
"""Creates a TorchModule instance based on the input module."""
|
||||
|
||||
return TorchModuleORT(module, debug_options, fallback_manager)
|
||||
return TorchModuleORT(module, debug_options, fallback_manager, logger)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
# _torch_module_ort.py
|
||||
|
||||
from collections import OrderedDict
|
||||
from logging import Logger
|
||||
from typing import Callable, Iterator, Optional, Tuple, TypeVar
|
||||
|
||||
import torch
|
||||
|
|
@ -17,13 +18,17 @@ T = TypeVar("T", bound="torch.nn.Module")
|
|||
|
||||
|
||||
class TorchModuleORT(TorchModuleInterface):
|
||||
def __init__(self, module: torch.nn.Module, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
def __init__(
|
||||
self, module: torch.nn.Module, debug_options: DebugOptions, fallback_manager: _FallbackManager, logger: Logger
|
||||
):
|
||||
super().__init__(module)
|
||||
self._flattened_module = _io._FlattenedModule(module)
|
||||
|
||||
_utils.patch_torch_module_ort_forward_method(self)
|
||||
|
||||
self._execution_manager = GraphExecutionManagerFactory(self._flattened_module, debug_options, fallback_manager)
|
||||
self._execution_manager = GraphExecutionManagerFactory(
|
||||
self._flattened_module, debug_options, fallback_manager, logger
|
||||
)
|
||||
|
||||
def _apply(self, fn):
|
||||
"""Override original method to delegate execution to the flattened PyTorch user module"""
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import warnings
|
||||
from logging import Logger
|
||||
from typing import Tuple
|
||||
|
||||
import onnx
|
||||
|
|
@ -12,7 +12,7 @@ import torch
|
|||
from onnxruntime.capi import _pybind_state as C
|
||||
from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_type
|
||||
|
||||
from . import _are_deterministic_algorithms_enabled, _io, _logger, _use_deterministic_algorithms, _utils
|
||||
from . import _are_deterministic_algorithms_enabled, _io, _use_deterministic_algorithms, _utils
|
||||
from ._execution_agent import TrainingAgent
|
||||
from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPolicy
|
||||
from ._gradient_accumulation_manager import GradientAccumulationManager
|
||||
|
|
@ -27,8 +27,10 @@ class TrainingManager(GraphExecutionManager):
|
|||
TrainingManager is responsible for building and running the forward and backward graph of the training model.
|
||||
"""
|
||||
|
||||
def __init__(self, model: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
super().__init__(model, debug_options, fallback_manager)
|
||||
def __init__(
|
||||
self, model: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager, logger: Logger
|
||||
):
|
||||
super().__init__(model, debug_options, fallback_manager, logger)
|
||||
|
||||
self._export_mode = torch.onnx.TrainingMode.TRAINING
|
||||
self._forward_class = self._create_autofunction_class()
|
||||
|
|
@ -215,19 +217,14 @@ class TrainingManager(GraphExecutionManager):
|
|||
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
|
||||
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:
|
||||
# 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 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,
|
||||
self._logger.info(
|
||||
"Fast path enabled - skipping checks.Rebuild graph: %s, Execution agent: %s, Device check: %s",
|
||||
self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT),
|
||||
self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT),
|
||||
self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE),
|
||||
)
|
||||
|
||||
# If exporting module to ONNX for the first time, this skip check will not take effect.
|
||||
|
|
@ -263,6 +260,8 @@ class TrainingManager(GraphExecutionManager):
|
|||
# Build the gradient graph
|
||||
self._build_graph(graph_transformer_config)
|
||||
|
||||
self._log_feature_stats()
|
||||
|
||||
# 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
|
||||
|
|
|
|||
|
|
@ -7,11 +7,11 @@
|
|||
import copy
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import traceback
|
||||
import types
|
||||
import warnings
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
|
@ -22,11 +22,8 @@ from torch.utils.dlpack import to_dlpack
|
|||
|
||||
from onnxruntime.capi import _pybind_state as C
|
||||
from onnxruntime.capi.onnxruntime_inference_collection import OrtValue
|
||||
from onnxruntime.tools import pytorch_export_contrib_ops
|
||||
|
||||
from . import _onnx_models
|
||||
from ._custom_gradient_registry import CustomGradientRegistry
|
||||
from ._custom_op_symbolic_registry import CustomOpSymbolicRegistry
|
||||
from ._fallback_exceptions import ORTModuleDeviceException, ORTModuleIOError, wrap_exception
|
||||
from ._torch_module_pytorch import TorchModulePytorch
|
||||
from .torch_cpp_extensions.cpu.aten_op_executor import load_aten_op_executor_cpp_extension
|
||||
|
|
@ -215,7 +212,9 @@ def _create_iobinding(io_binding, inputs, model, device: torch.device):
|
|||
io_binding.bind_output(value_info.name, device.type, device_id=device_id)
|
||||
|
||||
|
||||
def check_for_name_collisions_and_bind_methods_to_ortmodule(ortmodule: torch.nn.Module, user_module: torch.nn.Module):
|
||||
def check_for_name_collisions_and_bind_methods_to_ortmodule(
|
||||
ortmodule: torch.nn.Module, user_module: torch.nn.Module, logger: logging.Logger
|
||||
):
|
||||
"""Warns if there are any common attributes between the user's model and ORTModule and binds user methods to ORTModule
|
||||
|
||||
If there are methods defined on the user's model that ORTModule does not recognize (custom methods),
|
||||
|
|
@ -255,7 +254,7 @@ def check_for_name_collisions_and_bind_methods_to_ortmodule(ortmodule: torch.nn.
|
|||
# This is a user defined/overridden method. Check for collisions.
|
||||
if attribute_name in ortmodule_attributes:
|
||||
# This is a user defined method, issue a warning.
|
||||
warnings.warn(
|
||||
logger.warning(
|
||||
f"User Module's attribute name {attribute_name} collides with ORTModule's attribute name. "
|
||||
"User Module's method may not be called upon invocation through ORTModule."
|
||||
)
|
||||
|
|
@ -269,7 +268,7 @@ def check_for_name_collisions_and_bind_methods_to_ortmodule(ortmodule: torch.nn.
|
|||
if attribute_name not in torch_module_attributes and attribute_name in ortmodule_attributes:
|
||||
# This is a user defined attribute that collides with ORTModule
|
||||
if attribute_name in ortmodule_attributes:
|
||||
warnings.warn(
|
||||
logger.warning(
|
||||
f"User Module's attribute name {attribute_name} collides with ORTModule's attribute name. "
|
||||
"User Module's attribute may not be returned when trying to retrieve the attribute through ORTModule."
|
||||
)
|
||||
|
|
@ -345,8 +344,8 @@ def switch_backend_to_pytorch(ortmodule, pytorch_module):
|
|||
ortmodule.forward = pytorch_module.forward
|
||||
|
||||
|
||||
def warn_of_constant_inputs(data):
|
||||
warnings.warn(
|
||||
def warn_of_constant_inputs(data, logger: logging.Logger):
|
||||
logger.info(
|
||||
f"Received input of type {type(data)} which may be treated as a constant by ORT by default."
|
||||
" Please consider moving constant arguments to the model constructor."
|
||||
)
|
||||
|
|
@ -389,19 +388,6 @@ def patch_ortmodule_forward_method(ortmodule):
|
|||
functools.update_wrapper(ortmodule.forward.__func__, ortmodule._torch_module.forward.__func__)
|
||||
|
||||
|
||||
def reinitialize_ortmodule(ortmodule):
|
||||
# Re-register contrib OPs
|
||||
pytorch_export_contrib_ops.register()
|
||||
CustomOpSymbolicRegistry.register_all()
|
||||
CustomGradientRegistry.register_all()
|
||||
|
||||
# Re-initialize the ORTModule forward method
|
||||
patch_ortmodule_forward_method(ortmodule)
|
||||
|
||||
# Re-bind users custom methods to ORTModule
|
||||
check_for_name_collisions_and_bind_methods_to_ortmodule(ortmodule, ortmodule.module)
|
||||
|
||||
|
||||
def reinitialize_torch_module_ort(torch_module):
|
||||
# Re-initialize the forward method
|
||||
patch_torch_module_ort_forward_method(torch_module)
|
||||
|
|
@ -427,3 +413,9 @@ def reinitialize_graph_execution_manager(graph_execution_manager):
|
|||
def reinitialize_training_manager(training_manager):
|
||||
# Redefine training managers forward_class
|
||||
training_manager._forward_class = training_manager._create_autofunction_class()
|
||||
|
||||
|
||||
def get_runtime_pytorch_version():
|
||||
from packaging import version
|
||||
|
||||
return version.parse(torch.__version__.split("+")[0])
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
# --------------------------------------------------------------------------
|
||||
# isort: skip_file
|
||||
# Import ordering is important in this module to aviod circular dependencies
|
||||
import logging
|
||||
from ._torch_module_factory import TorchModuleFactory
|
||||
from ._torch_module_ort import TorchModuleORT
|
||||
from ._custom_op_symbolic_registry import CustomOpSymbolicRegistry
|
||||
|
|
@ -11,6 +12,7 @@ from ._custom_gradient_registry import CustomGradientRegistry
|
|||
from . import _utils
|
||||
from .debug_options import DebugOptions
|
||||
from ._fallback import _FallbackManager, _FallbackPolicy, ORTModuleFallbackException
|
||||
from ._logger import ortmodule_loglevel_to_python_loglevel
|
||||
from onnxruntime.training import ortmodule
|
||||
|
||||
from onnxruntime.tools import pytorch_export_contrib_ops
|
||||
|
|
@ -51,9 +53,15 @@ class ORTModule(torch.nn.Module):
|
|||
if not debug_options:
|
||||
debug_options = DebugOptions()
|
||||
|
||||
self._logger = logging.getLogger(__name__)
|
||||
self._logger.setLevel(ortmodule_loglevel_to_python_loglevel(debug_options.logging.log_level))
|
||||
|
||||
# Fallback settings
|
||||
self._fallback_manager = _FallbackManager(
|
||||
pytorch_module=module, policy=ortmodule.ORTMODULE_FALLBACK_POLICY, retry=ortmodule.ORTMODULE_FALLBACK_RETRY
|
||||
pytorch_module=module,
|
||||
policy=ortmodule.ORTMODULE_FALLBACK_POLICY,
|
||||
retry=ortmodule.ORTMODULE_FALLBACK_RETRY,
|
||||
logger=self._logger,
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -63,7 +71,7 @@ class ORTModule(torch.nn.Module):
|
|||
|
||||
super().__init__()
|
||||
|
||||
self._torch_module = TorchModuleFactory()(module, debug_options, self._fallback_manager)
|
||||
self._torch_module = TorchModuleFactory()(module, debug_options, self._fallback_manager, self._logger)
|
||||
|
||||
_utils.patch_ortmodule_forward_method(self)
|
||||
|
||||
|
|
@ -75,7 +83,7 @@ class ORTModule(torch.nn.Module):
|
|||
# Warn user if there are name collisions between user model's and ORTModule attributes
|
||||
# And if there are custom methods defined on the user's model, copy and bind them to
|
||||
# ORTModule.
|
||||
_utils.check_for_name_collisions_and_bind_methods_to_ortmodule(self, module)
|
||||
_utils.check_for_name_collisions_and_bind_methods_to_ortmodule(self, module, self._logger)
|
||||
|
||||
except ORTModuleFallbackException as e:
|
||||
# Although backend is switched to PyTorch here,
|
||||
|
|
@ -310,4 +318,13 @@ class ORTModule(torch.nn.Module):
|
|||
def __setstate__(self, state):
|
||||
self.__dict__.update(state)
|
||||
|
||||
_utils.reinitialize_ortmodule(self)
|
||||
# Re-register contrib OPs
|
||||
pytorch_export_contrib_ops.register()
|
||||
CustomOpSymbolicRegistry.register_all()
|
||||
CustomGradientRegistry.register_all()
|
||||
|
||||
# Re-initialize the ORTModule forward method
|
||||
_utils.patch_ortmodule_forward_method(self)
|
||||
|
||||
# Re-bind users custom methods to ORTModule
|
||||
_utils.check_for_name_collisions_and_bind_methods_to_ortmodule(self, self.module, self._logger)
|
||||
|
|
|
|||
|
|
@ -4214,18 +4214,20 @@ def test_hf_save_pretrained():
|
|||
assert p1.data.ne(p2.data).sum() == 0
|
||||
|
||||
|
||||
def test_ortmodule_string_inputs_are_ignored():
|
||||
def test_ortmodule_string_inputs_are_ignored(caplog):
|
||||
pt_model = MyStrNet()
|
||||
ort_model = ORTModule(copy.deepcopy(pt_model))
|
||||
ort_model = ORTModule(copy.deepcopy(pt_model), DebugOptions(log_level=LogLevel.INFO))
|
||||
x = torch.randn(1, 2)
|
||||
|
||||
with pytest.warns(UserWarning) as warning_record:
|
||||
out = ort_model(x, "hello")
|
||||
out = ort_model(x, "hello")
|
||||
|
||||
assert (
|
||||
"Received input of type <class 'str'> which may be treated as a constant by ORT by default."
|
||||
in warning_record[1].message.args[0]
|
||||
)
|
||||
target_str = "Received input of type <class 'str'> which may be treated as a constant by ORT by default."
|
||||
found_target_str = False
|
||||
for record in caplog.records:
|
||||
if target_str in record.message:
|
||||
found_target_str = True
|
||||
|
||||
assert found_target_str
|
||||
_test_helpers.assert_values_are_close(out, x + 1)
|
||||
|
||||
|
||||
|
|
@ -4803,8 +4805,7 @@ def test_ortmodule_setattr_signals_model_changed():
|
|||
del os.environ["ORTMODULE_SKIPCHECK_POLICY"]
|
||||
|
||||
|
||||
@pytest.mark.skipif(Version(torch.__version__) < Version("1.13.0"), reason="PyTorch 1.12 incompatible")
|
||||
def test_ortmodule_attribute_name_collision_warning():
|
||||
def test_ortmodule_attribute_name_collision_warning(caplog):
|
||||
class UserNet(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
|
@ -4819,18 +4820,14 @@ def test_ortmodule_attribute_name_collision_warning():
|
|||
|
||||
device = "cuda"
|
||||
pt_model = UserNet().to(device)
|
||||
with pytest.warns(UserWarning) as warning_record:
|
||||
ORTModule(pt_model)
|
||||
|
||||
# FutureWarning('The first argument to symbolic functions is deprecated in 1.13 and will be removed in the future.
|
||||
# Please annotate treat the first argument (g) as GraphContext and use context information from the object
|
||||
# instead.')
|
||||
# TODO(bmeswani): Check with the exporter team as to what this might mean for ortmodule.
|
||||
# For ROCm EP, the log above appears twice.
|
||||
assert len(warning_record) == 3 or len(warning_record) == 4
|
||||
ORTModule(pt_model)
|
||||
warning_record = [record.message for record in caplog.records if record.levelname == "WARNING"]
|
||||
|
||||
assert "_torch_module collides with ORTModule's attribute name." in warning_record[-2].message.args[0]
|
||||
assert "load_state_dict collides with ORTModule's attribute name." in warning_record[-1].message.args[0]
|
||||
assert len(warning_record) == 2
|
||||
|
||||
assert "_torch_module collides with ORTModule's attribute name." in warning_record[-2]
|
||||
assert "load_state_dict collides with ORTModule's attribute name." in warning_record[-1]
|
||||
|
||||
|
||||
def test_ortmodule_ortmodule_method_attribute_copy():
|
||||
|
|
@ -5733,10 +5730,10 @@ def test_runtime_inspector_label_and_embed_sparsity_detection(embed_is_sparse, l
|
|||
found_embed_is_sparse = False
|
||||
found_label_is_sparse = False
|
||||
for record in caplog.records:
|
||||
if "Label sparsity based optimization is on for" in record.getMessage():
|
||||
if "Label sparsity-based optimization is ON for" in record.getMessage():
|
||||
found_label_is_sparse = True
|
||||
|
||||
if "Embedding sparsity based optimization is on for" in record.getMessage():
|
||||
if "Embedding sparsity-based optimization is ON for" in record.getMessage():
|
||||
found_embed_is_sparse = True
|
||||
|
||||
if label_is_sparse:
|
||||
|
|
|
|||
|
|
@ -533,7 +533,7 @@ def test_ortmodule_fallback_onnx_model__missing_op(is_training, fallback_enabled
|
|||
|
||||
|
||||
@pytest.mark.parametrize("is_training,persist_fallback", list(itertools.product([True, False], repeat=2)))
|
||||
def test_ortmodule_fallback_warn_message(is_training, persist_fallback):
|
||||
def test_ortmodule_fallback_warn_message(is_training, persist_fallback, caplog):
|
||||
# is_training: True for torch.nn.Module training model, eval mode otherwise
|
||||
|
||||
policy = "FALLBACK_UNSUPPORTED_DEVICE"
|
||||
|
|
@ -555,18 +555,35 @@ def test_ortmodule_fallback_warn_message(is_training, persist_fallback):
|
|||
inputs = torch.randn(N, D_in, device=data_device)
|
||||
|
||||
for i in range(3):
|
||||
with pytest.raises(RuntimeError), pytest.warns(UserWarning) as warning_record:
|
||||
# The run MUST fail no matter with ORT or PyTorch, so we catch the RuntimeError here in case it breaks the test.
|
||||
# Be noted, the logs will be caught by caplog.
|
||||
with pytest.raises(RuntimeError):
|
||||
ort_model(inputs)
|
||||
# For retries, the warn message will always be logged
|
||||
if not persist_fallback:
|
||||
assert "Fallback to PyTorch due to exception" in str(warning_record[0].message.args[0])
|
||||
if i == 0:
|
||||
# For the first time, run ORTModule, feature map is logged as warning
|
||||
# And the fallback warning is logged.
|
||||
assert len(caplog.records) == 2
|
||||
else:
|
||||
# For the other time, only the fallback warning is logged.
|
||||
assert len(caplog.records) == 1
|
||||
assert "Fallback to PyTorch due to exception" in caplog.records[-1].message
|
||||
caplog.clear()
|
||||
continue
|
||||
|
||||
# If `retries` is not enabled, only log the warn message once
|
||||
if i == 0:
|
||||
assert "Fallback to PyTorch due to exception" in str(warning_record[0].message.args[0])
|
||||
# For the first time, run ORTModule, feature map is logged as warning
|
||||
# And the fallback warning is logged.
|
||||
assert len(caplog.records) == 2
|
||||
assert "Fallback to PyTorch due to exception" in caplog.records[1].message
|
||||
caplog.clear()
|
||||
else:
|
||||
assert warning_record.list == []
|
||||
# For the other time, no fallback warning will be logged because
|
||||
# we are running with PyTorch.
|
||||
assert len(caplog.records) == 0
|
||||
caplog.clear()
|
||||
|
||||
del os.environ["ORTMODULE_SKIPCHECK_POLICY"]
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue