diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 335830ea36..3c7dacb1e4 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -12,7 +12,7 @@ import warnings from abc import ABC, abstractmethod # noqa: F401 from enum import IntFlag from functools import reduce -from typing import Dict, Tuple +from typing import Dict, List, Optional, Tuple import onnx import torch @@ -34,6 +34,7 @@ from ._fallback import ( ) from ._gradient_accumulation_manager import GradientAccumulationManager from ._graph_execution_interface import GraphExecutionInterface +from ._io import _FlattenedModule, _InputInfo, _ModelInputOutputSchemaType from .debug_options import DebugOptions, LogLevel from .torch_cpp_extensions.cpu.aten_op_executor import load_aten_op_executor_cpp_extension @@ -41,7 +42,7 @@ logger = logging.getLogger(__name__) class _RunStateInfo: - def __init__(self, state, output_info): + def __init__(self, state, output_info: List[Tuple[torch.Size, torch.device, torch.dtype]]): """ :param state: State of partial run that contains intermediate tensors needed to resume the run later. :param output_info: Output info. @@ -73,7 +74,7 @@ class _SkipCheck(IntFlag): class GraphExecutionManager(GraphExecutionInterface): - def __init__(self, module, debug_options: DebugOptions, fallback_manager: _FallbackManager): + def __init__(self, module: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager): """Manages construction and execution of ONNX graphs""" super().__init__(module._original_module) @@ -92,9 +93,9 @@ class GraphExecutionManager(GraphExecutionInterface): # Model after inference optimization or gradient building. self._graph_builder = None self._graph_info = None - self._graph_initializer_names = None - self._graph_initializer_names_to_train = None - self._graph_initializers = None + self._graph_initializer_names = set() + self._graph_initializer_names_to_train = set() + self._graph_initializers: List[torch.nn.parameter.Parameter] = [] # Update constant ONNX_OPSET_VERSION with env var ORTMODULE_ONNX_OPSET_VERSION # if defined. @@ -105,7 +106,7 @@ class GraphExecutionManager(GraphExecutionInterface): # TrainingAgent or InferenceAgent self._execution_agent = None - # indicators of some logic have been executed previously and thus could be skipped for faster training + # Indicators of some logic have been executed previously and thus could be skipped for faster training # default is enabled, if not defined in os env self._skip_check = _SkipCheck( _SkipCheck.SKIP_CHECK_DEVICE | _SkipCheck.SKIP_CHECK_BUILD_GRADIENT | _SkipCheck.SKIP_CHECK_EXECUTION_AGENT @@ -144,8 +145,6 @@ class GraphExecutionManager(GraphExecutionInterface): # It cannot overlap with required/immutable arguments (validated in runtime) self._export_extra_kwargs = {} - # Related to training graph shape inference - self._current_input_shape = None # default execution order is priority-based for both dynamic/static shape input for now # if we observe the benefit of static shape, we can expose this flag to the user self._use_static_shape = False @@ -158,11 +157,17 @@ class GraphExecutionManager(GraphExecutionInterface): self._enable_custom_autograd_function = custom_autograd_function_enabler.state - self._input_info = None - self._module_output_schema = None - self._device = _utils.get_device_from_module(module) + # Input and output infos (including schema) for exported model. + self._input_info: Optional[_InputInfo] = None + self._module_output_schema: Optional[_ModelInputOutputSchemaType] = None - self._module_parameters = list(inspect.signature(self._original_module.forward).parameters.values()) + # Device where the model is placed. + self._device: Optional[torch.device] = _utils.get_device_from_module(module) + + # Forward function input parameters of the original module. + self._module_parameters: List[inspect.Parameter] = list( + inspect.signature(self._original_module.forward).parameters.values() + ) # TODO: remove after PyTorch ONNX exporter supports VAR_KEYWORD parameters. for input_parameter in self._module_parameters: @@ -230,26 +235,6 @@ class GraphExecutionManager(GraphExecutionInterface): ), ) - @staticmethod - def execution_session_run_forward(execution_session, onnx_model, device, *inputs): - """Runs the forward pass on `execution_session` with given `onnx_model`, `device` and `inputs` - - This is a helper that can be called by the actual `GraphExecutionManager.forward` method - - Args: - execution_session (InferenceAgent or InferenceAgent): Agent which runs either inference or train - onnx_model (onnx.ModelProto): ONNX model - device (torch.device): PyTorch device - inputs: (torch.Tensor or a container of): User input - - Returns: - Returns a tuple (user_outputs, run_info): - user_outputs: The model output (either torch.Tensor or a container of torch.Tensor) - run_info: A _RunStateInfo which contains extra information about the execution of the graph - """ - - raise NotImplementedError - @abstractmethod def forward(self): """Executes the forward method for ORTModule @@ -330,7 +315,7 @@ class GraphExecutionManager(GraphExecutionInterface): return session_options, providers, provider_options - def _export_model(self, *inputs, **kwargs): + 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 # 3. Export the user model under self._export_training_flag mode @@ -374,8 +359,9 @@ class GraphExecutionManager(GraphExecutionInterface): return True - def _get_exported_model(self, input_schema, *inputs, **kwargs): - """Exports PyTorch `self._flattened_module` to ONNX for inferencing or training, using `*inputs` and `**kwargs` as input + def _get_exported_model(self, input_schema: _ModelInputOutputSchemaType, *inputs, **kwargs) -> onnx.ModelProto: + """Exports PyTorch `self._flattened_module` to ONNX for inferencing or training, + using `*inputs` and `**kwargs` as input TODO: How to support dynamic axes? Dimensions are determined by samples """ @@ -459,7 +445,7 @@ class GraphExecutionManager(GraphExecutionInterface): ORTModuleDeviceException, RuntimeError("A device must be specified in the model or inputs!") ) - def _get_graph_transformer_config(self): + def _get_graph_transformer_config(self) -> C.TrainingGraphTransformerConfiguration: graph_transformer_config = C.TrainingGraphTransformerConfiguration() graph_transformer_config.propagate_cast_ops_config = C.PropagateCastOpsConfiguration() graph_transformer_config.propagate_cast_ops_config.level = self._propagate_cast_ops_level diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py index 95849a3335..2bb9001c60 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py @@ -3,18 +3,21 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +from typing import Union + from ._fallback import _FallbackManager from ._inference_manager import InferenceManager +from ._io import _FlattenedModule from ._training_manager import TrainingManager from .debug_options import DebugOptions class GraphExecutionManagerFactory: - def __init__(self, module, debug_options: DebugOptions, fallback_manager: _FallbackManager): + def __init__(self, module: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager): self._training_manager = TrainingManager(module, debug_options, fallback_manager) self._inference_manager = InferenceManager(module, debug_options, fallback_manager) - def __call__(self, is_training): + def __call__(self, is_training) -> Union[InferenceManager, TrainingManager]: if is_training: return self._training_manager else: diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index ff333691ce..e3e211d749 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -4,6 +4,7 @@ # -------------------------------------------------------------------------- import warnings +from typing import Tuple import onnx import torch @@ -20,7 +21,7 @@ from .debug_options import DebugOptions class InferenceManager(GraphExecutionManager): """Concrete instance of GraphExecutionManager that is able to manage the inference model - InferenceManager is resposible for building and running the forward graph of the inference model + InferenceManager is responsible for building and running the forward graph of the inference model """ def __init__(self, model, debug_options: DebugOptions, fallback_manager: _FallbackManager): @@ -28,9 +29,25 @@ class InferenceManager(GraphExecutionManager): self._export_mode = torch.onnx.TrainingMode.EVAL @staticmethod - def execution_session_run_forward(execution_session, onnx_model, device, *inputs): - """Runs the forward graph on execution_session with given model inputs and device""" + def execution_session_run_forward( + execution_session, + onnx_model: onnx.ModelProto, + device: torch.device, + *inputs, + ) -> Tuple[Tuple[torch.Tensor, ...], _RunStateInfo]: + """Runs the forward pass on `execution_session` with given `onnx_model`, `device` and `inputs` + Args: + execution_session InferenceAgent: Agent which runs inference + onnx_model (onnx.ModelProto): ONNX model + device (torch.device): PyTorch device + inputs: (torch.Tensor or a container of): User inputs passed from ORTModule.forward(). + + Returns: + Returns a tuple (user_outputs, run_info): + user_outputs: The model output (either torch.Tensor or a container of torch.Tensor) + run_info: A _RunStateInfo which contains extra information about the execution of the graph + """ # 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 @@ -58,6 +75,14 @@ class InferenceManager(GraphExecutionManager): ONNX model is exported the first time this method is executed. Next, we build an optimized inference graph with module_graph_builder. Finally, we instantiate the ONNX Runtime InferenceSession through the InferenceAgent. + + The call stack is as follows: + ORTModule.forward(*inputs, **kwargs) -> + ORTModule._torch_module.forward(*inputs, **kwargs) where _torch_module is a TorchModuleORT instance -> + ORTModule._torch_module._execution_manager(is_training()).forward(*inputs, **kwargs) where: + TorchModuleORT._execution_manager(true) is a TrainingManager instance; + and TorchModuleORT._execution_manager(false) is an InferenceManager instance. + """ # Fallback to PyTorch due to failures *external* to forward(), diff --git a/orttraining/orttraining/python/training/ortmodule/_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py index e264bf28b8..419d63f090 100644 --- a/orttraining/orttraining/python/training/ortmodule/_io.py +++ b/orttraining/orttraining/python/training/ortmodule/_io.py @@ -8,14 +8,84 @@ import gc import inspect import warnings from collections import abc +from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Tuple, Union +import onnx import torch from ._fallback import ORTModuleIOError, ORTModuleONNXModelException, wrap_exception -from ._logger import suppress_os_stream_output +from ._logger import LogLevel, suppress_os_stream_output from ._runtime_inspector import RuntimeInspector from ._utils import warn_of_constant_inputs +_ModelInputOutputType = Union[ + None, + str, + int, + bool, + float, + torch.Tensor, + Sequence["_ModelInputOutputType"], + Mapping[str, "_ModelInputOutputType"], +] + + +class _TensorStub: + """Tensor stub class used to represent model's input or output""" + + __slots__ = ["name", "dtype", "shape", "shape_dims"] + + def __init__( + self, name: Optional[str] = None, dtype: Optional[str] = None, shape=None, shape_dims: Optional[int] = None + ): + self.name: Optional[str] = name + self.dtype: Optional[str] = dtype + self.shape = shape + self.shape_dims: Optional[int] = shape_dims # r.g. rank. + + def __repr__(self) -> str: + result = "_TensorStub(" + if self.name is not None: + result += f"name={self.name}" + if self.dtype is not None: + if result[-1] != "(": + result += ", " + result += f"dtype={self.dtype}" + if self.shape is not None: + if result[-1] != "(": + result += ", " + result += f"shape={self.shape}" + if self.shape_dims is not None: + if result[-1] != "(": + result += ", " + result += f"shape_dims={self.shape_dims}" + result += ")" + return result + + def __eq__(self, other): + if not other: + return False + elif not isinstance(other, _TensorStub): + raise NotImplementedError("_TensorStub must only be compared to another _TensorStub instance!") + elif self.name != other.name: + return False + elif self.dtype != other.dtype: + return False + elif self.shape != other.shape: + return False + elif self.shape_dims != other.shape_dims: + return False + return True + + +_ModelInputOutputSchemaType = Union[ + None, + str, + _TensorStub, + Sequence["_ModelInputOutputSchemaType"], + Mapping[str, "_ModelInputOutputSchemaType"], +] + class _OutputIdentityOp(torch.autograd.Function): """Internal class used to prepend Identity ops in model's outputs @@ -76,7 +146,7 @@ class _PrimitiveType: return type(value) in _PrimitiveType._primitive_types @staticmethod - def get_tensor(value, device): + def get_tensor(value, device) -> torch.Tensor: return torch.tensor(value, device=device) @staticmethod @@ -115,18 +185,18 @@ def flatten_kwargs(kwargs, device): class _InputInfo: def __init__( self, - names, - shape, - require_grad_names=None, - dynamic_axes=None, - schema=None, + names: List[str], + shape: List[List[int]], + require_grad_names: Optional[List[str]] = None, + dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None, + schema: Optional[_ModelInputOutputSchemaType] = None, num_positionals=0, ): - self.names = names - self.shape = shape - self.require_grad_names = require_grad_names if require_grad_names else [] - self.dynamic_axes = dynamic_axes if dynamic_axes else {} - self.schema = schema if schema else [] + self.names: List[str] = names + self.shape: List[List[int]] = shape + self.require_grad_names: List[str] = require_grad_names if require_grad_names else [] + self.dynamic_axes: Dict[str, Dict[int, str]] = dynamic_axes if dynamic_axes else {} + self.schema: _ModelInputOutputSchemaType = schema if schema else [] self.num_positionals = num_positionals self.kwargs = None @@ -139,7 +209,9 @@ class _InputInfo: \tSchema: {self.schema} \t#Positionals (total): {self.num_positionals}""" - def flatten(self, args, kwargs, device): + def flatten( + self, args: Sequence[_ModelInputOutputType], kwargs: Mapping[str, _ModelInputOutputType], device: torch.device + ) -> Sequence[_ModelInputOutputType]: """Flatten args and kwargs in a single tuple of tensors with strict ordering""" ret = [_PrimitiveType.get_tensor(arg, device) if _PrimitiveType.is_primitive_type(arg) else arg for arg in args] @@ -154,7 +226,9 @@ class _InputInfo: return ret - def unflatten(self, flat_args): + def unflatten( + self, flat_args: Sequence[_ModelInputOutputType] + ) -> Tuple[Sequence[_ModelInputOutputType], Mapping[str, _ModelInputOutputType]]: """Unflatten tuple of tensors into args and kwargs""" args = tuple(flat_args[: self.num_positionals]) @@ -164,13 +238,20 @@ class _InputInfo: def _combine_input_buffers_initializers( - params, onnx_input_names, input_info, buffer_names, inputs, kwargs, device, rt_inspector: RuntimeInspector + params: List[torch.nn.parameter.Parameter], + onnx_input_names: List[str], + input_info: Optional[_InputInfo], + named_buffer: Iterator[Tuple[str, torch.Tensor]], + inputs: Sequence[_ModelInputOutputType], + kwargs: Mapping[str, _ModelInputOutputType], + device: torch.device, + rt_inspector: RuntimeInspector, ): """Creates forward `*inputs` list from user input and PyTorch initializers ONNX Runtime forward requires an ordered list of: * User input: computed from forward InferenceSession - * Initializers: computed from original PyTorch model parameters + * Initializers: computed from original PyTorch model parameters. """ def _expand_inputs(current_input, non_none_inputs, name=""): @@ -202,7 +283,7 @@ def _combine_input_buffers_initializers( _expand_inputs(inputs, non_none_inputs) flattened_kwargs_inputs = {} _expand_inputs(kwargs, flattened_kwargs_inputs) - buffer_names_dict = {buffer_name: inp for buffer_name, inp in buffer_names} + buffer_names_dict = {buffer_name: inp for buffer_name, inp in named_buffer} result = [] embed_sparsity_results = [] label_sparsity_results = [] @@ -256,7 +337,9 @@ def _combine_input_buffers_initializers( return result, embed_sparsity_results, label_sparsity_results -def deepcopy_model_input(*inputs, **kwargs): +def deepcopy_model_input( + *args, **kwargs +) -> Tuple[Sequence[_ModelInputOutputType], Mapping[str, _ModelInputOutputType]]: def extract_tensor(value): if isinstance(value, torch.Tensor): if value.requires_grad: @@ -266,64 +349,18 @@ def deepcopy_model_input(*inputs, **kwargs): else: return value - sample_inputs_copy = [extract_tensor(value) for value in inputs] - sample_inputs_copy = copy.deepcopy(tuple(sample_inputs_copy)) + sample_args_copy: Sequence[_ModelInputOutputType] = [extract_tensor(value) for value in args] + sample_args_copy = copy.deepcopy(tuple(sample_args_copy)) - sample_kwargs_copy = {} + sample_kwargs_copy: Mapping[str, _ModelInputOutputType] = {} for name, value in kwargs.items(): sample_kwargs_copy[name] = extract_tensor(value) sample_kwargs_copy = copy.deepcopy(sample_kwargs_copy) - return sample_inputs_copy, sample_kwargs_copy + return sample_args_copy, sample_kwargs_copy -class _TensorStub: - """Tensor stub class used to represent model's input or output""" - - __slots__ = ["name", "dtype", "shape", "shape_dims"] - - def __init__(self, name=None, dtype=None, shape=None, shape_dims=None): - self.name = name - self.dtype = dtype - self.shape = shape - self.shape_dims = shape_dims - - def __repr__(self) -> str: - result = "_TensorStub(" - if self.name is not None: - result += f"name={self.name}" - if self.dtype is not None: - if result[-1] != "(": - result += ", " - result += f"dtype={self.dtype}" - if self.shape is not None: - if result[-1] != "(": - result += ", " - result += f"shape={self.shape}" - if self.shape_dims is not None: - if result[-1] != "(": - result += ", " - result += f"shape_dims={self.shape_dims}" - result += ")" - return result - - def __eq__(self, other): - if not other: - return False - elif not isinstance(other, _TensorStub): - raise NotImplementedError("_TensorStub must only be compared to another _TensorStub instance!") - elif self.name != other.name: - return False - elif self.dtype != other.dtype: - return False - elif self.shape != other.shape: - return False - elif self.shape_dims != other.shape_dims: - return False - return True - - -def unflatten_user_output(output_schema, outputs): +def unflatten_user_output(output_schema: Optional[_ModelInputOutputSchemaType], outputs): """Follows the schema to generate an output that is expected by the user""" def _replace_stub_with_tensor_value(user_output, outputs, output_idx): @@ -368,8 +405,49 @@ def unflatten_user_output(output_schema, outputs): return user_output -def _extract_schema(data): - """Extract the data schema by replacing every torch.Tensor value with _TensorStub""" +def _extract_schema(data: _ModelInputOutputType) -> _ModelInputOutputSchemaType: + """Extract the data schema by replacing every torch.Tensor value with _TensorStub. + + Depth first traversal to iterate over the data: + > Replace every tensor with a stub + > Replace None/str typed data with itself + > Recreate tensor from data for other primitive types, and replace them with a stub. + + TODO(pengwa): use a unified logger and log level control. + Examples: + Example 1, list: + data = [torch.tensor(1), torch.tensor(2)] + schema = [_TensorStub(shape=()), _TensorStub(shape=())] + + Example 2, dict: + data = {"a": torch.tensor(1), "b": torch.tensor(2)} + schema = {"a": _TensorStub(shape=()), "b": _TensorStub(shape=())} + + Example 3, dict of list: + data = {"a": [torch.tensor(1), torch.tensor(2)], "b": [torch.tensor(3), torch.tensor(4)]} + schema = {"a": [_TensorStub(shape=()), _TensorStub(shape=())], + "b": [_TensorStub(shape=()), _TensorStub(shape=())]} + + Example 4, nested dict: + data = {"a": {"b": torch.tensor(1), "c": torch.tensor(2)}, + "d": {"e": torch.tensor(3), "f": torch.tensor(4)}} + schema = {"a": {"b": _TensorStub(shape=()), "c": _TensorStub(shape=())}, + "d": {"e": _TensorStub(shape=()), "f": _TensorStub(shape=())}} + + Example 5, dict of mixed list and dict: + data = {"a": [torch.tensor(1), torch.tensor(2)], "b": {"c": torch.tensor(3), "d": torch.tensor(4)}} + schema = {"a": [_TensorStub(shape=()), _TensorStub(shape=())], + "b": {"c": _TensorStub(shape=()), "d": _TensorStub(shape=())}} + + Args: + data: The data to extract the schema from, which can be in any kind of nested structure, + including Sequence and Mapping. + + + Returns: + The schema of the data, which has the same structure as the data. + + """ if data is None: return data @@ -386,7 +464,7 @@ def _extract_schema(data): # Instead of replacing the tensor with a stub in the original user input, build the stubbed_schema # from scratch from the user input. - stubbed_schema = None + stubbed_schema: Optional[_ModelInputOutputSchemaType] = None if isinstance(data, abc.Sequence): sequence_type = type(data) stubbed_schema = [_extract_schema(val) for val in data] @@ -407,10 +485,12 @@ def _extract_schema(data): return stubbed_schema -def _parse_outputs_and_extract_names_and_dynamic_axes(module_output): +def _parse_outputs_and_extract_names_and_dynamic_axes(module_output) -> Tuple[List[str], Dict[str, Dict[int, str]]]: """Parses through the module output and returns output names and dynamic axes""" - def _populate_output_names_and_dynamic_axes(output, output_names, output_dynamic_axes, output_idx): + def _populate_output_names_and_dynamic_axes( + output, output_names: List[str], output_dynamic_axes: Dict[str, Dict[int, str]], output_idx: List[int] + ): # Depth first traversal to traverse through the entire output collecting output names and dynamic axes if output is None: @@ -438,9 +518,9 @@ def _parse_outputs_and_extract_names_and_dynamic_axes(module_output): TypeError(f"ORTModule does not support the following model output type {type(output)}"), ) - output_names = [] - output_dynamic_axes = {} - output_idx = [0] + output_names: List[str] = [] + output_dynamic_axes: Dict[str, Dict[int, str]] = {} + output_idx: List[int] = [0] _populate_output_names_and_dynamic_axes(module_output, output_names, output_dynamic_axes, output_idx) return output_names, output_dynamic_axes @@ -474,21 +554,53 @@ def _transform_output_to_flat_tuple(data): class _FlattenedModule(torch.nn.Module): - def __init__(self, original_module): + def __init__(self, original_module: torch.nn.Module): super().__init__() - self._original_module = original_module + self._original_module: torch.nn.Module = original_module # Before `forward` is called, _ort_module must be assigned # Updated input info is needed to expand args into *args, **kwargs - self._input_info = None + self._input_info: Optional[_InputInfo] = None def forward(self, *args): new_args, new_kwargs = self._input_info.unflatten(args) return _transform_output_to_flat_tuple(self._original_module(*new_args, **new_kwargs)) -def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, inputs, kwargs): - def _add_dynamic_shape(name, input): +def parse_inputs_for_onnx_export( + all_input_parameters: List[inspect.Parameter], + onnx_graph: Optional[onnx.ModelProto], + schema: _ModelInputOutputSchemaType, + args: Sequence[_ModelInputOutputType], + kwargs: Mapping[str, _ModelInputOutputType], +) -> _InputInfo: + """Parses through the model inputs and returns _InputInfo. + + Loop through all input parameters, try to flatten them into a 1-D list of inputs. For nested data in the inputs, + construct the name in hierarchical order. + + Example 1, arg is a list, kwarg is a dict: + args = [arg1, arg2], kwargs = {"a": 4, "b": 5}, + input_names = ["arg1", "arg2", "a", "b"]. + + Example 2, arg is a list, kwarg is a dict of mixed list and scalar: + args = [arg1, arg2], kwargs = {"a": [4, 5], "b": 6}, + input_names = ["arg1", "arg2", "a_0", "a_1", "b"]. + + Example 3, arg is a list, kwarg is a dict of mixed dict and scalar: + args = [arg1, arg2], kwargs = {"a": {"c": 4, "d": 5}, "b": 6}, + input_names = ["arg1", "arg2", "a_c", "a_d", "b"]. + + Args: + all_input_parameters: All inspected input parameters from the original model forward function signature. + onnx_graph: (optional) The onnx graph. + schema: The schema extracted from the positional arguments and keyword arguments of the model. + args: The positional arguments of the model. + kwargs: The keyword arguments of the model. + + """ + + def _add_dynamic_shape(name, input) -> Dict[str, Dict[int, str]]: dynamic_axes[name] = {} for dim_idx in range(len(input.shape)): dynamic_axes[name].update({dim_idx: f"{name}_dim{dim_idx}"}) @@ -534,24 +646,43 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, input # Ignore optional inputs explicitly specified as None # ONNX exporter may remove unused inputs - onnx_graph_input_names = [] + onnx_graph_input_names: List[str] = [] if onnx_graph is not None: onnx_graph_input_names = {inp.name for inp in onnx_graph.graph.input} - input_names = [] - dynamic_axes = {} - input_names_require_grad = [] - input_shape = [] + input_names: List[str] = [] + dynamic_axes: Dict[str, Dict[int, str]] = {} + input_names_require_grad: List[str] = [] + input_shape: List[List[int]] = [] var_positional_idx = 0 + # Be noted, all_input_parameters is a list of inspect.Parameters parsed from the original module's forward method. + # While the execution manger's forward function will map all given model inputs to *args and **kwargs, so it is + # possible the input parameter list cannot represent the real model inputs given here (e.g., *args, **kwargs). + # But it is still fine to use all_input_parameters to make sure all model inputs are covered. + # + # Here is an example caused by the mismatch between all_input_parameters and real model inputs. + # def foo(*args, named_kwargs, **kwargs): + # ... print("foo") + # From inspection, + # > ('args', <_ParameterKind.VAR_POSITIONAL: 2>) + # > ('named_kwargs', <_ParameterKind.KEYWORD_ONLY: 3>) + # > ('kwargs', <_ParameterKind.VAR_KEYWORD: 4>) + # + # At this point, 'named_kwargs' exists in **kwargs as a result of ORTModule's forward parse all original + # model inputs in to *args and **kwargs. + # When we loop `all_input_parameters``, for the `named_kwargs`, we will try to handle it in KEYWORD_ONLY branch. + # Additionally in VAR_KEYWORD branch, we will get the `named_kwargs` value again, because its name exists in the + # `kwargs`. So _add_input avoids handling the `named_kwargs` twice, check test case `test_named_kwargs_dict_input` + # for the details. for input_idx, input_parameter in enumerate(all_input_parameters): if input_parameter.kind == inspect.Parameter.VAR_POSITIONAL: # VAR_POSITIONAL parameter carries all *args parameters from original forward method - for args_i in range(input_idx, len(inputs)): + for args_i in range(input_idx, len(args)): name = f"{input_parameter.name}_{var_positional_idx}" var_positional_idx += 1 - inp = inputs[args_i] + inp = args[args_i] _add_input(name, inp, onnx_graph, onnx_graph_input_names) elif ( input_parameter.kind == inspect.Parameter.POSITIONAL_ONLY @@ -562,8 +693,8 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, input name = input_parameter.name inp = None input_idx += var_positional_idx # noqa: PLW2901 - if input_idx < len(inputs) and inputs[input_idx] is not None: - inp = inputs[input_idx] + if input_idx < len(args) and args[input_idx] is not None: + inp = args[input_idx] elif name in kwargs and kwargs[name] is not None: inp = kwargs[name] _add_input(name, inp, onnx_graph, onnx_graph_input_names) @@ -578,18 +709,20 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, input require_grad_names=input_names_require_grad, dynamic_axes=dynamic_axes, schema=schema, - num_positionals=len(inputs), + num_positionals=len(args), ) -def parse_outputs_for_onnx_export_and_extract_schema(module, inputs, kwargs, log_level): +def parse_outputs_for_onnx_export_and_extract_schema( + module, args: Sequence[_ModelInputOutputType], kwargs: Mapping[str, _ModelInputOutputType], log_level: LogLevel +): # Perform a forward call to grab outputs output_names = None output_dynamic_axes = None is_deepcopy = False with torch.no_grad(), suppress_os_stream_output(log_level=log_level): # Deepcopy inputs, since input values may change after model run. - sample_inputs_copy, sample_kwargs_copy = deepcopy_model_input(*inputs, **kwargs) + sample_args_copy, sample_kwargs_copy = deepcopy_model_input(*args, **kwargs) try: # Deepcopy model, in case model is stateful and changes after model run. model_copy = copy.deepcopy(module) @@ -602,7 +735,7 @@ def parse_outputs_for_onnx_export_and_extract_schema(module, inputs, kwargs, log " Compute will continue, but unexpected results may occur!" ) - sample_outputs = model_copy(*sample_inputs_copy, **sample_kwargs_copy) + sample_outputs = model_copy(*sample_args_copy, **sample_kwargs_copy) # Parse the output and extract the output_names and output_dynamic_axes to be used for onnx export output_names, output_dynamic_axes = _parse_outputs_and_extract_names_and_dynamic_axes(sample_outputs) diff --git a/orttraining/orttraining/python/training/ortmodule/_onnx_models.py b/orttraining/orttraining/python/training/ortmodule/_onnx_models.py index 1dd0a12dea..45e9992692 100644 --- a/orttraining/orttraining/python/training/ortmodule/_onnx_models.py +++ b/orttraining/orttraining/python/training/ortmodule/_onnx_models.py @@ -4,6 +4,7 @@ import os from dataclasses import dataclass +from typing import Optional import onnx import torch @@ -31,9 +32,9 @@ class ONNXModels: InferenceSession and is saved by the InferenceSession. """ - exported_model: onnx.ModelProto = None - optimized_model: onnx.ModelProto = None - optimized_pre_grad_model: onnx.ModelProto = None + exported_model: Optional[onnx.ModelProto] = None + optimized_model: Optional[onnx.ModelProto] = None + optimized_pre_grad_model: Optional[onnx.ModelProto] = None def save_exported_model(self, path, name_prefix, export_mode): # save the ortmodule exported model diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index f07a225bf8..65ebfce7c0 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -4,6 +4,7 @@ # -------------------------------------------------------------------------- import warnings +from typing import Tuple import onnx import torch @@ -14,24 +15,46 @@ from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_typ from . import _are_deterministic_algorithms_enabled, _io, _logger, _use_deterministic_algorithms, _utils from ._execution_agent import TrainingAgent from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPolicy +from ._gradient_accumulation_manager import GradientAccumulationManager from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck +from ._io import _FlattenedModule, _InputInfo from .debug_options import DebugOptions class TrainingManager(GraphExecutionManager): """Concrete instance of GraphExecutionManager that is able to manage the training model - TrainingManager is responsible for building and running the forward and backward graph of the training model + TrainingManager is responsible for building and running the forward and backward graph of the training model. """ - def __init__(self, model, debug_options: DebugOptions, fallback_manager: _FallbackManager): + def __init__(self, model: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager): super().__init__(model, debug_options, fallback_manager) + self._export_mode = torch.onnx.TrainingMode.TRAINING self._forward_class = self._create_autofunction_class() @staticmethod - def execution_session_run_forward(execution_session, onnx_model, device, gradient_accumulation_manager, *inputs): - """Runs the forward graph on execution_session with given model inputs and device""" + def execution_session_run_forward( + execution_session, + onnx_model: onnx.ModelProto, + device: torch.device, + gradient_accumulation_manager: GradientAccumulationManager, + *inputs, + ) -> Tuple[Tuple[torch.Tensor, ...], _RunStateInfo]: + """Runs the forward pass on `execution_session` with given `onnx_model`, `device` and `inputs` + + Args: + execution_session (InferenceAgent or TrainingAgent): Agent which runs training. + onnx_model (onnx.ModelProto): ONNX model + device (torch.device): PyTorch device + gradient_accumulation_manager (GradientAccumulationManager): Gradient accumulation manager + inputs: (torch.Tensor or a container of): User inputs passed from ORTModule.forward(). + + Returns: + Returns a tuple (user_outputs, run_info): + user_outputs: The model output (either torch.Tensor or a container of torch.Tensor) + run_info: A _RunStateInfo which contains extra information about the execution of the graph + """ # TODO: Try to reuse the output buffers as some of the output tensors are same sizes, # especially the backward graph outputs. @@ -54,7 +77,9 @@ class TrainingManager(GraphExecutionManager): # Run and return module outputs. execution_session.run_forward(forward_inputs, forward_outputs, state, gradient_accumulation_manager.cache) - user_outputs = gradient_accumulation_manager.extract_outputs_and_maybe_update_cache(forward_outputs, device) + user_outputs: Tuple[torch.Tensor, ...] = gradient_accumulation_manager.extract_outputs_and_maybe_update_cache( + forward_outputs, device + ) output_info = [(output.shape, output.device, output.dtype) for output in user_outputs] run_info = _RunStateInfo(state, output_info) @@ -174,6 +199,14 @@ class TrainingManager(GraphExecutionManager): ONNX model is exported the first time this method is executed. Next, we build a full training graph with module_graph_builder. Finally, we instantiate the ONNX Runtime InferenceSession. + + The call stack is as follows: + ORTModule.forward(*inputs, **kwargs) -> + ORTModule._torch_module.forward(*inputs, **kwargs) where _torch_module is a TorchModuleORT instance -> + ORTModule._torch_module._execution_manager(is_training()).forward(*inputs, **kwargs) where: + TorchModuleORT._execution_manager(true) is a TrainingManager instance; + and TorchModuleORT._execution_manager(false) is an InferenceManager instance. + """ # Fallback to PyTorch due to failures *external* to forward(), @@ -209,7 +242,7 @@ class TrainingManager(GraphExecutionManager): # If model was exported, then initialize the graph builder self._initialize_graph_builder() - # since the schema was just extracted while trying to export the model and it was either + # Since the schema was just extracted while trying to export the model and it was either # saved to self._input_info.schema or checked for equality with the self._input_info.schema # it should not need to be updated again. Pass it inside parse_inputs_for_onnx_export. input_info = _io.parse_inputs_for_onnx_export( @@ -366,7 +399,7 @@ class TrainingManager(GraphExecutionManager): local_device_rank, ) - def _reinitialize_graph_builder(self, input_info): + def _reinitialize_graph_builder(self, input_info: _InputInfo): """Return true if the module graph builder was reinitialized""" # Model may have unused params dropped after export and not part of self._graph_initializer_names_to_train diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py index 0151cd532a..0839a25f53 100644 --- a/orttraining/orttraining/python/training/ortmodule/_utils.py +++ b/orttraining/orttraining/python/training/ortmodule/_utils.py @@ -12,7 +12,7 @@ import random import traceback import types import warnings -from typing import List # noqa: F401 +from typing import List, Optional, Tuple, Union import numpy as np import torch @@ -49,7 +49,8 @@ def set_random_states(states): torch.cuda.set_rng_state(torch_cuda_state) -def _ortvalue_from_torch_tensor(torch_tensor): +def _ortvalue_from_torch_tensor(torch_tensor: torch.Tensor) -> C.OrtValue: + """Converts a torch tensor to an OrtValue.""" # TODO: Current DLPack doesn't support bool and PyTorch disables converting bool tensor to DLPack in recent commit. # https://github.com/pytorch/pytorch/blob/7e7be526c9d9179f35084e9cca5b5c5ad5172100/aten/src/ATen/DLConvertor.cpp#L41 # We need to convert bool tensor to unit8 tensor to workaround this. @@ -63,7 +64,9 @@ def _ortvalue_from_torch_tensor(torch_tensor): return C.OrtValue.from_dlpack(to_dlpack(torch_tensor), is_bool_tensor) -def _ortvalues_to_torch_tensor(ortvalues, device=None): +def _ortvalues_to_torch_tensor( + ortvalues: C.OrtValueVector, device: Optional[torch.device] = None +) -> Tuple[torch.Tensor, ...]: if len(ortvalues) == 0: return tuple() @@ -75,7 +78,7 @@ def _ortvalues_to_torch_tensor(ortvalues, device=None): if not isinstance(ortvalues, C.OrtValueVector): raise TypeError("ortvalues must be an instance of OrtValueVector not %r." % type(ortvalues)) - res = ortvalues.to_dlpacks(_from_dlpack) + res: List[torch.Tensor] = ortvalues.to_dlpacks(_from_dlpack) bool_indices = ortvalues.bool_tensor_indices() if len(bool_indices): # DLPack structure does not know for sure if it stores boolean @@ -97,7 +100,7 @@ def _ortvalues_to_torch_tensor(ortvalues, device=None): return tuple(res) -def _torch_tensor_to_dlpack(tensor): +def _torch_tensor_to_dlpack(tensor: torch.Tensor): # TODO: Current DLPack doesn't support bool and PyTorch disables converting bool tensor to DLPack in recent commit. # https://github.com/pytorch/pytorch/blob/7e7be526c9d9179f35084e9cca5b5c5ad5172100/aten/src/ATen/DLConvertor.cpp#L41 # We need to convert bool tensor to unit8 tensor to workaround this. @@ -110,7 +113,7 @@ def _torch_tensor_to_dlpack(tensor): return to_dlpack(tensor) -def _check_same_device(device, argument_str, *args): +def _check_same_device(device: torch.device, argument_str: str, *args): """Check that all tensor arguments in *args reside on the same device as the input device""" assert isinstance(device, torch.device), "`device` must be a valid `torch.device` object" @@ -126,7 +129,7 @@ def _check_same_device(device, argument_str, *args): ) -def get_device_index(device): +def get_device_index(device: Union[str, int, torch.device]) -> int: if isinstance(device, str): # could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0 device = torch.device(device) @@ -135,7 +138,7 @@ def get_device_index(device): return 0 if device.index is None else device.index -def get_device_str(device): +def get_device_str(device: Union[str, int, torch.device]) -> str: if isinstance(device, str): # could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0 if device.find(":") == -1: @@ -152,12 +155,15 @@ def get_device_str(device): return device -def get_device_from_module(module): +def get_device_from_module(module) -> Optional[torch.device]: """Returns the first device found in the `module`'s parameters or None Args: module (torch.nn.Module): PyTorch model to extract device from + Returns: + torch.device: Device extracted from `module`'s parameters or None + Raises: ORTModuleFallbackException: When more than one device is found at `module` """ @@ -175,12 +181,15 @@ def get_device_from_module(module): return device -def get_device_from_inputs(args, kwargs): +def get_device_from_inputs(args, kwargs) -> Optional[torch.device]: """Returns device from first PyTorch Tensor within args or kwargs Args: args: List with inputs kwargs: Dictionary with inputs + + Returns: + torch.device: Device extracted from `args` or `kwargs` or None """ device = None @@ -191,7 +200,7 @@ def get_device_from_inputs(args, kwargs): return device -def _create_iobinding(io_binding, inputs, model, device): +def _create_iobinding(io_binding, inputs, model, device: torch.device): """Creates IO binding for a `model` inputs and output""" for idx, value_info in enumerate(model.graph.input): io_binding.bind_ortvalue_input( diff --git a/orttraining/orttraining/python/training/ortmodule/ortmodule.py b/orttraining/orttraining/python/training/ortmodule/ortmodule.py index 16a8fc8633..bb80283c14 100644 --- a/orttraining/orttraining/python/training/ortmodule/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule/ortmodule.py @@ -33,7 +33,7 @@ class ORTModule(torch.nn.Module): debug_options (:obj:`DebugOptions`, optional): debugging options for ORTModule. """ - def __init__(self, module, debug_options=None): + def __init__(self, module: torch.nn.Module, debug_options: Optional[DebugOptions] = None): # NOTE: torch.nn.Modules that call setattr on their internal attributes regularly # (for example PyTorch Lightning), will trigger regular re-exports. This is # because ORTModule auto detects such setattrs on the original module and