mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
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:
This commit is contained in:
pengwa
GitHub
parent
347c963d5c
commit
1ebc5d3879
4 changed files with 142 additions and 31 deletions
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue