From 70ca03d491e95b2799c1be6f3b4838290b086d4d Mon Sep 17 00:00:00 2001 From: baijumeswani Date: Tue, 31 Aug 2021 15:28:04 -0700 Subject: [PATCH] Correctly set the skip check flags for ORTModule (#8891) --- .../ortmodule/_graph_execution_manager.py | 1 + .../training/ortmodule/_inference_manager.py | 63 ++++++++++++------- .../training/ortmodule/_training_manager.py | 28 +++++---- 3 files changed, 58 insertions(+), 34 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index af955ef8d7..5ad652a10b 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -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 diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index 1a58f2e44b..a81f83422c 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -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, diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index 5d6bf4aaa0..6b1696981d 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -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()