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"""