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

Correctly set the skip check flags for ORTModule (#8891)

Browse source
This commit is contained in:
baijumeswani 2021-08-31 15:28:04 -07:00 committed by GitHub
parent 69ab4670f7
commit 70ca03d491
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 58 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
View file

@ -95,6 +95,7 @@ class GraphExecutionManager(GraphExecutionInterface):
# indicators of some logic have been executed previously thus could be skipped for faster training
self._skip_check = _SkipCheck.SKIP_CHECK_DISABLED
self._first_skip_check_warning = True
# Graph transformer config
# Specify cast propagation strategy. Currently three strategies are available, NONE, INSERT-AND-REDUCE and FLOOD-FILL

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

@ -4,7 +4,7 @@
# --------------------------------------------------------------------------
from . import _utils, _io, _logger
from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo
from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck
from ._execution_agent import InferenceAgent
from .debug_options import DebugOptions
from ._fallback import ORTModuleFallbackException, _FallbackPolicy, _FallbackManager
@ -29,9 +29,6 @@ class InferenceManager(GraphExecutionManager):
def execution_session_run_forward(execution_session, onnx_model, device, *inputs):
"""Runs the forward graph on execution_session with given model inputs and device"""
# Assert that the input and model device match
_utils._check_same_device(device, "Input argument to forward", *inputs)
# TODO: Try to reuse the output buffers as some of the output tensors are same sizes,
# especially the backward graph outputs.
# REVIEW(codemzs): Consolidate Training Agent with InferenceAgent on C++ side to not
@ -49,10 +46,6 @@ class InferenceManager(GraphExecutionManager):
forward_output._ortvalue) for forward_output in forward_outputs)
state = None
# Assert that the outputs and model device match
_utils._check_same_device(
device, "Output argument from forward", *user_outputs)
output_info = [(output.shape, output.device, output.dtype)
for output in user_outputs]
run_info = _RunStateInfo(state, output_info)
@ -73,27 +66,51 @@ class InferenceManager(GraphExecutionManager):
return self._fallback_manager.fallback(self._original_module, self._debug_options.logging.log_level, *inputs, **kwargs)
try:
# Exporting module to ONNX for the first time
build_graph = self._export_model(*inputs, **kwargs)
if build_graph:
# If model was exported, then initialize the graph builder
self._initialize_graph_builder(training=False)
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.
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)
# Build the inference graph
if build_graph:
self._build_graph()
# 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 \
not self._onnx_models.exported_model:
# Exporting module to ONNX for the first time
build_graph = self._export_model(*inputs, **kwargs)
if build_graph:
# If model was exported, then initialize the graph builder
self._initialize_graph_builder(training=False)
# Build the inference graph
if build_graph:
self._build_graph()
# 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:
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
if self._device != module_device:
self._device = module_device
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
if self._device != module_device:
self._device = module_device
if create_execution_session:
# Create execution session creates the inference_session
self._create_execution_agent()
if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) == False:
# Assert that the input and model device match
_utils._check_same_device(self._device, "Input argument to forward", *inputs)
user_outputs, _ = InferenceManager.execution_session_run_forward(self._execution_agent,
self._onnx_models.optimized_model,
self._device,

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

@ -65,9 +65,20 @@ class TrainingManager(GraphExecutionManager):
return self._fallback_manager.fallback(self._original_module, self._debug_options.logging.log_level, *inputs, **kwargs)
try:
# Exporting module to ONNX for the first time
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.
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)
# 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:
if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT) == 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
@ -92,8 +103,11 @@ class TrainingManager(GraphExecutionManager):
if build_gradient_graph:
self._build_graph()
# 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:
if self._skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT) == 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
@ -106,14 +120,6 @@ class TrainingManager(GraphExecutionManager):
# Create execution session creates the training_session
self._create_execution_agent()
# disable some checks after execution session is created the first time
if self._skip_check.is_disabled() == False:
self._skip_check = _SkipCheck.SKIP_CHECK_BUILD_GRADIENT | _SkipCheck.SKIP_CHECK_EXECUTION_AGENT | _SkipCheck.SKIP_CHECK_DEVICE
if self._debug_options.logging.log_level <= _logger.LogLevel.WARNING:
warnings.warn("Fast path enabled - skipping checks for rebuilding gradient graph, execution agent creation, and device during training.",
UserWarning)
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()