From 1ebc5d38795c20412dbc049dfb312bcbbd1d6408 Mon Sep 17 00:00:00 2001 From: pengwa Date: Tue, 11 Jul 2023 14:11:29 +0800 Subject: [PATCH] Log ORTModule initialization overhead (#16529) ### Log ORTModule initialization overhead When profiling some model for example ``` torchrun --nproc_per_node=1 examples/onnxruntime/training/language-modeling/run_mlm.py --model_name_or_path microsoft/deberta-v3-large --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --num_train_epochs 10 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --do_train --overwrite_output_dir --output_dir ./outputs/ --seed 1137 --fp16 --report_to none --optim adamw_ort_fused --max_steps 200 --logging_steps 1 --use_module_with_loss {'train_runtime': 303.8711, 'train_samples_per_second': 0.658, 'train_steps_per_second': 0.658, 'train_loss': 6.569518616199494, 'epoch': 0.09} 100%|200/200 [05:03<00:00, 1.52s/it] ***** train metrics ***** epoch = 0.09 train_loss = 6.5695 train_runtime = 0:05:03.87 train_samples = 2223 train_samples_per_second = 0.658 train_steps_per_second = 0.658 ``` The end to end time is 303s (train_runtime=0:05:03.87), but the ORTModule first step initialization (including export, graph build, etc) takes about 255s, so when we compare the end to end time for a baseline ORT with an improved version of ORT, there is no perf gains, since the x% gains over (303-255) is diluted out among the overall 303s. This is misleading! So this PR outputs the ORTModule initialization overhead in the output, then we can manually compute the real compte time and get the perf gains. If the log level is >= WARNING, then only the total end to end time + export time is logged, otherwise, more details of break down is logged: ![image](https://github.com/microsoft/onnxruntime/assets/10530022/8e34283d-4868-4f22-b65b-9f00d10d8fb7) ![image](https://github.com/microsoft/onnxruntime/assets/10530022/c13bcfad-0d79-483d-a886-e238efcbe657) --- .../ortmodule/_graph_execution_manager.py | 54 ++++++----- .../training/ortmodule/_inference_manager.py | 14 ++- .../python/training/ortmodule/_logger.py | 94 ++++++++++++++++++- .../training/ortmodule/_training_manager.py | 11 ++- 4 files changed, 142 insertions(+), 31 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 3732e6cc4c..2e256eb241 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -65,7 +65,9 @@ class GraphExecutionManager(GraphExecutionInterface): self._logger = logger + # Management for ORTModule configuration. self._runtime_options = _RuntimeOptions(self._logger) + # Original and flattened (transformed) output module self._flattened_module = module @@ -84,8 +86,12 @@ class GraphExecutionManager(GraphExecutionInterface): self._first_skip_check_warning = True + # Inspector for runtime information, for example input data, memory usage, etc. self._runtime_inspector = RuntimeInspector(self._logger) + # Tracker for ORTModule model export, session creation overhead. + self.time_tracker = _logger.TimeTracker() + # Value can be either torch.onnx.TrainingMode.TRAINING or torch.onnx.TrainingMode.EVAL # To be instantiated in the concrete implementation of GraphExecutionManager self._export_mode = None @@ -98,7 +104,6 @@ class GraphExecutionManager(GraphExecutionInterface): 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) @@ -233,6 +238,7 @@ class GraphExecutionManager(GraphExecutionInterface): return session_options, providers, provider_options + @_logger.TrackTime(_logger.TimeTrackerPhase.EXPORT) def _export_model(self, *inputs, **kwargs) -> bool: # 1. Set the self._device from the user module # 2. Verify input schema matches the schema used on the previous model export @@ -284,10 +290,6 @@ 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 - - start = datetime.now() - # Setup dynamic axes for onnx model self._input_info = _io.parse_inputs_for_onnx_export( self._module_parameters, None, input_schema, inputs, kwargs @@ -363,9 +365,6 @@ class GraphExecutionManager(GraphExecutionInterface): 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 def _set_device_from_module(self, inputs, kwargs): @@ -388,6 +387,7 @@ class GraphExecutionManager(GraphExecutionInterface): graph_transformer_config.enable_compute_optimizer = self._runtime_options.enable_compute_optimizer return graph_transformer_config + @_logger.TrackTime(_logger.TimeTrackerPhase.GRAPH_BUILDER_INIT) def _initialize_graph_builder(self): """Creates a new OrtModuleGraphBuilder, initializes it and saves it to self._graph_builder""" @@ -459,6 +459,7 @@ class GraphExecutionManager(GraphExecutionInterface): _utils.reinitialize_graph_execution_manager(self) + @_logger.TrackTime(_logger.TimeTrackerPhase.DETECTION) def _enable_conditional_optimizations( self, graph_transformer_config: C.TrainingGraphTransformerConfiguration, inputs: Tuple, kwargs: Dict ): @@ -545,22 +546,23 @@ class GraphExecutionManager(GraphExecutionInterface): ), ] - if self._runtime_options.enable_compute_optimizer: - feature_map.extend( - [ - ( - "Compute Optimizer", - self._runtime_options.enable_compute_optimizer, - "Enable/Disable with env ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=1/0", - ), - ( - " -FLOPReduction", - self._runtime_options.enable_compute_optimizer, - "Reduce FLOPs by upstreaming shrinking-sized ops", - ), - ] - ) + # Add compute optimizer + feature_map.extend( + [ + ( + "Compute Optimizer", + self._runtime_options.enable_compute_optimizer, + "Enable/Disable with env ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=1/0", + ), + ( + " -FLOPReduction", + self._runtime_options.enable_compute_optimizer, + "Reduce FLOPs by upstreaming shrinking-sized ops", + ), + ] + ) + if self._runtime_options.enable_compute_optimizer: if len(self._runtime_options.label_sparsity_ratio) > 0: feature_map.append( (" -LabelSparsityOpt", True, f"Input density: {self._runtime_options.label_sparsity_ratio}") @@ -571,6 +573,7 @@ class GraphExecutionManager(GraphExecutionInterface): (" -EmbedSparsityOpt", True, f"Input density: {self._runtime_options.embed_sparsity_ratio}") ) + # Add fallback feature_map.append( ( "Auto Fallback", @@ -590,7 +593,10 @@ class GraphExecutionManager(GraphExecutionInterface): 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" + + # Collect ORTModule overheads for different phases. + stat += f"{self.time_tracker.to_string(self._debug_options.logging.log_level < LogLevel.WARNING)}\n" + stat += f"Versions: ONNX Runtime - {onnxruntime.__version__}, ONNX - {onnx.__version__}\n\n" stat += f"{_logger.LogColor.HEADER}************************************************************************{_logger.LogColor.ENDC}\n\n" diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index 7fe11c1510..1fbb735df4 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -15,6 +15,7 @@ from . import _are_deterministic_algorithms_enabled, _io, _use_deterministic_alg from ._execution_agent import InferenceAgent from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPolicy from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo +from ._logger import TimeTrackerPhase, TrackTime from .options import DebugOptions, _SkipCheck @@ -109,10 +110,12 @@ class InferenceManager(GraphExecutionManager): self._runtime_options.skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT) is False or not self._onnx_models.exported_model ): + self.time_tracker.start(TimeTrackerPhase.EndToEnd) + # 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 + # If model was exported, then initialize the graph builder. self._initialize_graph_builder() # Build the inference graph @@ -121,11 +124,9 @@ class InferenceManager(GraphExecutionManager): # Set the config according to input inspection. self._enable_conditional_optimizations(graph_transformer_config, inputs, kwargs) - # Build the gradient graph + # Build the 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 @@ -149,6 +150,9 @@ class InferenceManager(GraphExecutionManager): # Create execution session creates the inference_session self._create_execution_agent() + self.time_tracker.end(TimeTrackerPhase.EndToEnd) + self._log_feature_stats() + if self._runtime_options.skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) is False: # Assert that the input and model device match _utils._check_same_device(self._device, "Input argument to forward", *inputs) @@ -187,6 +191,7 @@ class InferenceManager(GraphExecutionManager): if self._fallback_manager.is_pending(): return self._fallback_manager.fallback(self._debug_options.logging.log_level, *inputs, **kwargs) + @TrackTime(TimeTrackerPhase.BUILD_GRAPH) def _build_graph(self, graph_transformer_config): """Build an inference graph using the module_graph_builder""" @@ -199,6 +204,7 @@ class InferenceManager(GraphExecutionManager): self._export_mode, ) + @TrackTime(TimeTrackerPhase.CREATE_SESSION) def _create_execution_agent(self): """Creates an InferenceAgent that can run forward graph on an inference model""" diff --git a/orttraining/orttraining/python/training/ortmodule/_logger.py b/orttraining/orttraining/python/training/ortmodule/_logger.py index e3c6f0074c..f075897434 100644 --- a/orttraining/orttraining/python/training/ortmodule/_logger.py +++ b/orttraining/orttraining/python/training/ortmodule/_logger.py @@ -6,9 +6,10 @@ import io import logging import sys +import time from contextlib import contextmanager from enum import IntEnum -from typing import Dict, List +from typing import Callable, Dict, List from onnxruntime.capi._pybind_state import Severity @@ -76,3 +77,94 @@ class LogColor: ENDC = "\033[0m" BOLD = "\033[1m" UNDERLINE = "\033[4m" + + +class TimeTrackerPhase(IntEnum): + EndToEnd = 0 # The total overhead of ORT first-time initialization + EXPORT = 1 # The latency of preparing and exporting the model to ONNX + GRAPH_BUILDER_INIT = 2 # The latency of initializing the graph builder + DETECTION = 3 # The latency of runtime detection + BUILD_GRAPH = 4 # The latency of optimizing forward graph (and building the gradient graph for training). + CREATE_SESSION = 5 # The latency of creating the session + + def to_string(self) -> str: + if self == TimeTrackerPhase.EndToEnd: + return "end to end" + if self == TimeTrackerPhase.EXPORT: + return "export" + elif self == TimeTrackerPhase.GRAPH_BUILDER_INIT: + return "graph builder init" + elif self == TimeTrackerPhase.DETECTION: + return "runtime detection" + elif self == TimeTrackerPhase.BUILD_GRAPH: + return "graph building" + elif self == TimeTrackerPhase.CREATE_SESSION: + return "session creation" + else: + return "invalid" + + +class TimeTracker: + """A simple class to track time spent in different phases of ORT backend first-time initialization.""" + + NOT_RECORD = -1.0 + + def __init__( + self, + ): + self.starts_: List[float] = [TimeTracker.NOT_RECORD] * len(TimeTrackerPhase) + self.ends_: List[float] = [TimeTracker.NOT_RECORD] * len(TimeTrackerPhase) + + def start(self, phase: TimeTrackerPhase): + self.starts_[phase] = time.time() + + def end(self, phase: TimeTrackerPhase): + self.ends_[phase] = time.time() + + def _get_duration(self, phase: TimeTrackerPhase): + if self.ends_[phase] == TimeTracker.NOT_RECORD or self.starts_[phase] == TimeTracker.NOT_RECORD: + return TimeTracker.NOT_RECORD + return self.ends_[phase] - self.starts_[phase] + + def to_string(self, log_details=False) -> str: + end_to_end_str = self._get_duration(TimeTrackerPhase.EndToEnd) + end_to_end_str = f"{end_to_end_str:.2f}" if end_to_end_str != TimeTracker.NOT_RECORD else "N/A" + export_str = self._get_duration(TimeTrackerPhase.EXPORT) + export_str = f"{export_str:.2f}" if export_str != TimeTracker.NOT_RECORD else "N/A" + overhead_title_str = ( + f"Total ORT initialization overhead is {end_to_end_str}s where export takes {export_str}s.\n" + ) + + if log_details is False: + return overhead_title_str + + duration_summaries = [] + for phase in TimeTrackerPhase: + _get_duration = self._get_duration(phase) + if phase in [TimeTrackerPhase.EndToEnd, TimeTrackerPhase.EXPORT]: + continue + + val = ( + f" {phase.to_string()} takes {_get_duration:.2f}s" if _get_duration != TimeTracker.NOT_RECORD else "N/A" + ) + duration_summaries.append(f"{val}") + + return f"{overhead_title_str}Other overhead details: {','.join(duration_summaries)}\n" + + +class TrackTime: + """A function decorator to track time spent in different phases of ORT backend first-time initialization.""" + + def __init__(self, phase: TimeTrackerPhase): + self.phase = phase + + def __call__(self, func: Callable): + def wrapper(graph_execution_manager, *args, **kwargs): + if not hasattr(graph_execution_manager, "time_tracker"): + raise RuntimeError("The class of the function to be tracked must have a 'time_tracker' attribute.") + graph_execution_manager.time_tracker.start(self.phase) + result = func(graph_execution_manager, *args, **kwargs) + graph_execution_manager.time_tracker.end(self.phase) + return result + + return wrapper diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index 42901d34a7..d8c28addb8 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -18,6 +18,7 @@ from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPo from ._gradient_accumulation_manager import GradientAccumulationManager from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo from ._io import _FlattenedModule, _InputInfo +from ._logger import TimeTrackerPhase, TrackTime from ._runtime_inspector import Phase from .options import DebugOptions, _SkipCheck @@ -243,7 +244,10 @@ class TrainingManager(GraphExecutionManager): self._runtime_options.skip_check.is_set(_SkipCheck.SKIP_CHECK_BUILD_GRADIENT) is False or not self._onnx_models.exported_model ): + self.time_tracker.start(TimeTrackerPhase.EndToEnd) + build_gradient_graph = self._export_model(*inputs, **kwargs) + if build_gradient_graph: # If model was exported, then initialize the graph builder self._initialize_graph_builder() @@ -269,8 +273,6 @@ 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 @@ -298,6 +300,9 @@ class TrainingManager(GraphExecutionManager): self._runtime_options.enable_grad_acc_optimization, self._flattened_module, self._graph_info ) + self.time_tracker.end(TimeTrackerPhase.EndToEnd) + self._log_feature_stats() + self._gradient_accumulation_manager.maybe_update_cache_before_run() prepared_input_list, _, _ = _io._combine_input_buffers_initializers( @@ -331,6 +336,7 @@ class TrainingManager(GraphExecutionManager): if self._fallback_manager.is_pending(): return self._fallback_manager.fallback(self._debug_options.logging.log_level, *inputs, **kwargs) + @TrackTime(TimeTrackerPhase.BUILD_GRAPH) def _build_graph(self, graph_transformer_config): """Build an optimized gradient graph using the module_graph_builder""" @@ -366,6 +372,7 @@ class TrainingManager(GraphExecutionManager): else: self._gradient_map.append(-1) + @TrackTime(TimeTrackerPhase.CREATE_SESSION) def _create_execution_agent(self): """Creates a TrainingAgent that can run the forward and backward graph on the training model"""