From c60ef62190c911f63ca1e075e9d9af5c2e9cc78e Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Wed, 17 Mar 2021 16:32:32 -0700 Subject: [PATCH] Update ORTModule feature with remaining PRs from feature branch (#7040) * Liqun/ort module perf1 (#6806) add mysql script to log perf data Co-authored-by: liqun * Resolve HTTP Error 503: Service Unavailable for MNIST dataset (#6989) * Reduce logging for ORTModule for the end user (#6982) * Support none types in forward output (#7001) * Missed test case for none type output (#7014) * Fix code style according to autopep8 Co-authored-by: liqunfu Co-authored-by: baijumeswani --- .../core/framework/fallback_cpu_capability.cc | 2 +- .../providers/cuda/cuda_execution_provider.cc | 2 +- .../orttraining/core/agent/training_agent.cc | 2 +- .../core/framework/gradient_graph_builder.cc | 4 +- .../core/graph/gradient_builder_base.cc | 6 +- .../python/orttraining_pybind_state.cc | 4 +- .../_ortmodule_output_transformation.py | 16 +- .../orttraining/python/training/ortmodule.py | 223 +++++++++++------- .../python/orttraining_test_ortmodule_api.py | 23 ++ .../perf_log/ort_module_perf_test_tools.py | 206 ++++++++++++++++ 10 files changed, 394 insertions(+), 94 deletions(-) create mode 100644 orttraining/orttraining/test/python/perf_log/ort_module_perf_test_tools.py diff --git a/onnxruntime/core/framework/fallback_cpu_capability.cc b/onnxruntime/core/framework/fallback_cpu_capability.cc index 1653d88f43..b25c5407f1 100644 --- a/onnxruntime/core/framework/fallback_cpu_capability.cc +++ b/onnxruntime/core/framework/fallback_cpu_capability.cc @@ -132,7 +132,7 @@ std::unordered_set GetCpuPreferredNodes(const onnxruntime::GraphViewe if (place_in_cpu) { cpu_nodes.insert(cur); - LOGS_DEFAULT(WARNING) << "Force fallback to CPU execution for node: " << node->Name(); + LOGS_DEFAULT(INFO) << "Force fallback to CPU execution for node: " << node->Name(); for (auto* output : node->OutputDefs()) { cpu_output_args.insert(output); } diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 60e9d9d315..1ab5ebf21c 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -1988,7 +1988,7 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, // none of the provided registries has a CUDA kernel for this node if (cuda_kernel_def == nullptr) { - LOGS_DEFAULT(WARNING) << "CUDA kernel not found in registries for Op type: " << node.OpType() << " node name: " << node.Name(); + LOGS_DEFAULT(INFO) << "CUDA kernel not found in registries for Op type: " << node.OpType() << " node name: " << node.Name(); continue; } diff --git a/orttraining/orttraining/core/agent/training_agent.cc b/orttraining/orttraining/core/agent/training_agent.cc index 1d1a978b86..88f877edd0 100644 --- a/orttraining/orttraining/core/agent/training_agent.cc +++ b/orttraining/orttraining/core/agent/training_agent.cc @@ -115,7 +115,7 @@ common::Status TrainingAgent::RunBackward(int64_t run_id, const std::vectorGetLogger(), WARNING) << "Canceling background task with run_id " << run_id; + LOGS(*inference_session_->GetLogger(), INFO) << "Canceling background task with run_id " << run_id; // resume background thread with terminate = true onnxruntime::contrib::OrtTasks::GetInstance().SetBackwardInputs(run_id, {}, true); diff --git a/orttraining/orttraining/core/framework/gradient_graph_builder.cc b/orttraining/orttraining/core/framework/gradient_graph_builder.cc index 08811af545..c7721317ed 100644 --- a/orttraining/orttraining/core/framework/gradient_graph_builder.cc +++ b/orttraining/orttraining/core/framework/gradient_graph_builder.cc @@ -105,8 +105,8 @@ NodeSet GradientGraphBuilder::ReverseBFS(const NodeSet& nodes) const { for (auto edge_it = n->InputEdgesBegin(); edge_it != n->InputEdgesEnd(); ++edge_it) { auto it = STOP_GRADIENT_EDGES.find(n->OpType()); if (it != STOP_GRADIENT_EDGES.end() && it->second.count(edge_it->GetDstArgIndex())) { - LOGS(logger_, WARNING) << "Skip building gradient for input_" << edge_it->GetDstArgIndex() - << " of node: " << n->Name(); + LOGS(logger_, INFO) << "Skip building gradient for input_" << edge_it->GetDstArgIndex() + << " of node: " << n->Name(); continue; } diff --git a/orttraining/orttraining/core/graph/gradient_builder_base.cc b/orttraining/orttraining/core/graph/gradient_builder_base.cc index 985dfbc82b..342e6c67e7 100644 --- a/orttraining/orttraining/core/graph/gradient_builder_base.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_base.cc @@ -59,7 +59,7 @@ void ComputeBroadcastBackwardAxes( auto A_dim = A_dims[i].dim_param(), B_dim = B_dims[j].dim_param(); if (A_dim != B_dim) { - LOGS_DEFAULT(WARNING) << "Gradient building for node " << node_name << ": symbolic dimension expects to match. " << + LOGS_DEFAULT(INFO) << "Gradient building for node " << node_name << ": symbolic dimension expects to match. " << "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << " This is a relaxing case, and the kernel might run into problem later if A_dims and B_dims turns out not broadcastable."; } @@ -68,7 +68,7 @@ void ComputeBroadcastBackwardAxes( auto B_dim = B_dims[j].dim_value(); if (B_dim != 1) { - LOGS_DEFAULT(WARNING) << "Gradient building for node " << node_name << ": symbolic broadcasting expects the B_dimension to be 1. " << + LOGS_DEFAULT(INFO) << "Gradient building for node " << node_name << ": symbolic broadcasting expects the B_dimension to be 1. " << "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << " This is a relaxing case, and the kernel might run into problem later if A_dims and B_dims turns out not broadcastable."; } else { @@ -81,7 +81,7 @@ void ComputeBroadcastBackwardAxes( auto B_dim = B_dims[j].dim_param(); if (A_dim != 1) { - LOGS_DEFAULT(WARNING) << "Gradient building for node " << node_name << ": symbolic broadcasting expects the A_dimension to be 1. " << + LOGS_DEFAULT(INFO) << "Gradient building for node " << node_name << ": symbolic broadcasting expects the A_dimension to be 1. " << "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << " This is a relaxing case, and the kernel might run into problem later if A_dims and B_dims turns out not broadcastable."; } else { diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index 9c51ca8794..5290ba8166 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -506,9 +506,9 @@ py::class_(m, "TrainingAgent", R"pbdoc(This is the main class use .def_readwrite("use_invertible_layernorm_grad", &ModuleGradientGraphBuilderConfiguration::use_invertible_layernorm_grad); - py::class_ split_graphs_info(m, "TrainingGraphInfo", + py::class_ training_graph_info(m, "TrainingGraphInfo", R"pbdoc(The information of split graphs for frontend.)pbdoc"); - split_graphs_info.def(py::init()) + training_graph_info.def(py::init()) .def_readwrite("user_input_names", &TrainingGraphInfo::user_input_names) .def_readwrite("user_input_grad_names", &TrainingGraphInfo::user_input_grad_names) .def_readwrite("initializer_names_to_train", &TrainingGraphInfo::initializer_names_to_train) diff --git a/orttraining/orttraining/python/training/_ortmodule_output_transformation.py b/orttraining/orttraining/python/training/_ortmodule_output_transformation.py index 0e28b7e76a..ae7a6e1d5d 100644 --- a/orttraining/orttraining/python/training/_ortmodule_output_transformation.py +++ b/orttraining/orttraining/python/training/_ortmodule_output_transformation.py @@ -29,7 +29,9 @@ def populate_user_output_from_schema_and_outputs(output_schema, output_names, ou # Recursively traverse across user_output and replace all _TensorStub # with torch.Tensor values from outputs following output_idx - if isinstance(user_output, _TensorStub): + if user_output is None: + return None + elif isinstance(user_output, _TensorStub): output_idx[0] += 1 return outputs[output_idx[0]-1] @@ -65,8 +67,10 @@ def populate_user_output_from_schema_and_outputs(output_schema, output_names, ou def _extract_output_schema(output): """Extract the output schema by replacing every torch.Tensor value with _TensorStub""" + if output is None: + return None # Depth first traversal to iterate over the output to replace every tensor with a stub - if isinstance(output, torch.Tensor): + elif isinstance(output, torch.Tensor): return _TensorStub() if isinstance(output, abc.Sequence): @@ -94,7 +98,9 @@ def _parse_outputs_and_extract_names_and_dynamic_axes(module_output): def _populate_output_names_and_dynamic_axes(output, output_names, output_dynamic_axes, output_idx): # Depth first traversal to traverse through the entire output collecting output names and dynamic axes - if isinstance(output, torch.Tensor): + if output is None: + return + elif isinstance(output, torch.Tensor): output_name = f'output{output_idx[0]}' output_idx[0] += 1 output_names.append(output_name) @@ -128,7 +134,9 @@ def get_flattened_output_module(original_module): def _flatten_output(output, flat_output): # Recursively traverse over the output and populate the flat_output with torch.Tensors - if isinstance(output, torch.Tensor): + if output is None: + return + elif isinstance(output, torch.Tensor): flat_output.append(output) elif isinstance(output, abc.Sequence): for value in output: diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index 53ba983a1a..c1493b35fc 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -3,6 +3,12 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +from . import _utils +from . import _ortmodule_output_transformation as _ortmodule_io +from onnxruntime.training import register_custom_ops_pytorch_exporter +from onnxruntime.capi.onnxruntime_inference_collection import OrtValue +from onnxruntime.capi import _pybind_state as C + import functools import io import logging @@ -11,6 +17,7 @@ import onnxruntime import torch import inspect from inspect import signature +from enum import IntEnum from torch.utils.dlpack import from_dlpack, to_dlpack from torch.utils.cpp_extension import load_inline @@ -19,11 +26,6 @@ from torch.utils.cpp_extension import load_inline from typing import TypeVar T = TypeVar('T', bound='Module') -from onnxruntime.capi import _pybind_state as C -from onnxruntime.capi.onnxruntime_inference_collection import OrtValue -from onnxruntime.training import register_custom_ops_pytorch_exporter -from . import _utils, _ortmodule_output_transformation - ONNX_OPSET_VERSION = 12 @@ -31,13 +33,24 @@ ONNX_OPSET_VERSION = 12 def _ortvalue_to_dlpack(ortvalue): return ortvalue._ortvalue.to_dlpack() + def _ortvalue_from_dlpack(dlpack_tensor): return OrtValue(C.OrtValue.from_dlpack(dlpack_tensor)) + +class Verbosity(IntEnum): + VERBOSE = 0 + INFO = 1 + WARNING = 2 + ERROR = 3 + FATAL = 4 + + def _create_iobinding(io_binding, inputs, model, device): '''Creates IO binding for a `model` inputs and output''' for idx, value_info in enumerate(model.graph.input): - io_binding.bind_ortvalue_input(value_info.name, _ortvalue_from_dlpack(to_dlpack(inputs[idx]))) + io_binding.bind_ortvalue_input( + value_info.name, _ortvalue_from_dlpack(to_dlpack(inputs[idx]))) for value_info in model.graph.output: io_binding.bind_output(value_info.name, device.type, @@ -51,21 +64,24 @@ def _check_same_device(device, argument_str, *args): if arg is not None and isinstance(arg, torch.Tensor): arg_device = torch.device(arg.device) if arg_device != device: - raise RuntimeError(f"{argument_str} found on device {arg_device}, but expected it to be on module device {device}.") + raise RuntimeError( + f"{argument_str} found on device {arg_device}, but expected it to be on module device {device}.") + -# TODO: PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8, -# and convert the config to torch.uint8 tensor duing from_dlpack(). So a boolean tensor -# from forward graph outputs will be converted to torch.uint8 tensor. When this tensor -# is feeded to backward graph as input, it will cause data type mismatch issue during -# inference session running. We cannot change the from_dlpack() in PyTorch side, so we -# have to handle this specially, which will introduce a cast here and there is data copied. -# Always cast from torch.uint8 to torch.bool is not logically right, we need to check the -# real data type of the inputs in the backeard graph, and perform the cast only necessary. def _ort_output_to_torch_tensor(ort_output): + # TODO: PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8, + # and convert the config to torch.uint8 tensor duing from_dlpack(). So a boolean tensor + # from forward graph outputs will be converted to torch.uint8 tensor. When this tensor + # is feeded to backward graph as input, it will cause data type mismatch issue during + # inference session running. We cannot change the from_dlpack() in PyTorch side, so we + # have to handle this specially, which will introduce a cast here and there is data copied. + # Always cast from torch.uint8 to torch.bool is not logically right, we need to check the + # real data type of the inputs in the backeard graph, and perform the cast only necessary. tensor = from_dlpack(_ortvalue_to_dlpack(ort_output)) return tensor.to(torch.bool) if tensor.dtype == torch.uint8 else tensor -def _load_torch_allocator_cpp_extension(): + +def _load_torch_allocator_cpp_extension(verbosity): torch_cuda_allocator_addresses_cpp_source = """ #include #include @@ -78,13 +94,16 @@ def _load_torch_allocator_cpp_extension(): """ return load_inline(name='inline_extension', cpp_sources=[torch_cuda_allocator_addresses_cpp_source], - functions=['cuda_caching_allocator_raw_alloc_address', 'cuda_caching_allocator_raw_delete_address'], - verbose=True, with_cuda=True) + functions=['cuda_caching_allocator_raw_alloc_address', + 'cuda_caching_allocator_raw_delete_address'], + verbose=verbosity < Verbosity.WARNING, with_cuda=True) + class ORTModule(torch.nn.Module): def __init__(self, module): - assert isinstance(module, torch.nn.Module), "'module' must be a torch.nn.Module" + assert isinstance( + module, torch.nn.Module), "'module' must be a torch.nn.Module" # Create forward dynamically, so each ORTModule instance will have its own copy. # This is needed to be able to copy the forward signatures from the original PyTorch models @@ -102,17 +121,20 @@ class ORTModule(torch.nn.Module): # Exporting module to ONNX for the first time if not self._onnx_training: - device_from_module = _utils.get_device_from_module(self._original_module) + device_from_module = _utils.get_device_from_module( + self._original_module) if not self._device or self._device != device_from_module: self._device = device_from_module if not self._device: - raise RuntimeError('A device must be specified in the model or data!') - self._get_inference_graph_and_init_gradient_graph_builder(*inputs, **kwargs) + raise RuntimeError( + 'A device must be specified in the model or data!') + self._get_inference_graph_and_init_gradient_graph_builder( + *inputs, **kwargs) _, _, input_names_require_grad, new_input_shape = \ - _ortmodule_output_transformation.parse_inputs_for_onnx_export( + _ortmodule_io.parse_inputs_for_onnx_export( self._original_module_parameters, self._onnx_inference, *inputs, **kwargs) - # If inputs requiring gradient change from one call to forward to the next, the module_gradient_graph_builder + # If inputs requiring gradient change from forward to the next, the module_gradient_graph_builder # needs to be reinitialized so it can compute the backward output for the new inputs that require_grad if input_names_require_grad != self._input_names_require_grad: self._input_names_require_grad = input_names_require_grad @@ -123,15 +145,16 @@ class ORTModule(torch.nn.Module): self._build_training_graph() self._create_training_session() - module_device = _utils.get_device_from_module(self._original_module) + module_device = _utils.get_device_from_module( + self._original_module) if self._device != module_device: self._device = module_device self._create_training_session() - - # Use a custom torch.autograd.Function to associate self.backward_graph as the - # gradient implementation for self.forward_graph. class _ORTModuleFunction(torch.autograd.Function): + '''Use a custom torch.autograd.Function to associate self.backward_graph as the + gradient implementation for self.forward_graph.''' + @staticmethod def forward(ctx, *inputs, **kwargs): '''Performs forward pass based on user input and PyTorch initializer @@ -144,23 +167,31 @@ class ORTModule(torch.nn.Module): ''' # Assert that the input and model device match - _check_same_device(self._device, "Input argument to forward", *inputs) + _check_same_device( + self._device, "Input argument to forward", *inputs) # Use IO binding - _create_iobinding(self._training_io_binding, inputs, self._onnx_training, self._device) + _create_iobinding(self._training_io_binding, + inputs, self._onnx_training, self._device) # Run and return module outputs. - forward_outputs, run_id = self._training_session.run_forward(self._training_io_binding, self._run_options) - user_outputs = tuple(_ort_output_to_torch_tensor(forward_output) for forward_output in forward_outputs) + forward_outputs, run_id = self._training_session.run_forward( + self._training_io_binding, self._run_options) + user_outputs = tuple(_ort_output_to_torch_tensor( + forward_output) for forward_output in forward_outputs) ctx.run_id = run_id - # Disable materializing grads then None object will not be converted to a tensor filled with zeros prior to calling backward. - # Also save shape, device and type info to ctx for materializing tensor in backward if output grad is None. + # Disable materializing grads then None object will not be converted + # to a tensor filled with zeros prior to calling backward. + # Also save shape, device and type info to ctx for materializing + # tensor in backward if output grad is None. ctx.set_materialize_grads(False) - ctx.output_info = [(output.shape, output.device, output.dtype) for output in user_outputs] + ctx.output_info = [ + (output.shape, output.device, output.dtype) for output in user_outputs] # Assert that the outputs and model device match - _check_same_device(self._device, "Output argument from forward", *user_outputs) + _check_same_device( + self._device, "Output argument from forward", *user_outputs) return user_outputs @@ -170,7 +201,8 @@ class ORTModule(torch.nn.Module): ''' # Assert that the grad_outputs and model device match - _check_same_device(self._device, "Input argument to backward", *grad_outputs) + _check_same_device( + self._device, "Input argument to backward", *grad_outputs) # Use IO binding # Push user output grads to ONNX backend. @@ -179,17 +211,21 @@ class ORTModule(torch.nn.Module): if grad_output is None: shape, device, dtype = ctx.output_info[idx] if idx in self._onnx_graphs_info.output_grad_indices_require_full_shape: - grad_output = torch.zeros(shape, device=device, dtype=dtype) + grad_output = torch.zeros( + shape, device=device, dtype=dtype) else: - grad_output = torch.tensor(0., device=device, dtype=dtype) + grad_output = torch.tensor( + 0., device=device, dtype=dtype) elif not grad_output.is_contiguous(): grad_output = grad_output.contiguous() contiguous_grad_outputs.append(grad_output) - backward_grad_output_ortvalue = [_ortvalue_from_dlpack(to_dlpack(grad_output)) for grad_output in contiguous_grad_outputs] + backward_grad_output_ortvalue = [_ortvalue_from_dlpack( + to_dlpack(grad_output)) for grad_output in contiguous_grad_outputs] # Run and get results run_id = ctx.run_id - self._training_session.run_backward(backward_grad_output_ortvalue, run_id) + self._training_session.run_backward( + backward_grad_output_ortvalue, run_id) backward_outputs = self._training_io_binding.get_outputs() # Return input and initializer gradients @@ -206,26 +242,32 @@ class ORTModule(torch.nn.Module): # Append None to results for each input that did not require grad results.append(None) # Append gradients of initializer to results - results += [_ort_output_to_torch_tensor(backward_output) + results += [_ort_output_to_torch_tensor(backward_output) for backward_output in backward_outputs[num_user_input_grads:]] - # The OrtValue has a shared_ptr to the data. At this point there are two shared_ptrs to the data, one through the + # The OrtValue has a shared_ptr to the data. + # At this point there are two shared_ptrs to the data, one through the # OrtValue in the output iobinding, and the other through the copy in OrtDLManagedTensor. - # The following call clears the iobinding output, reducing the use_count to 1, so that once torch finishes computation - # on the DLpack tensors, the memory can be freed. + # The following call clears the iobinding output, reducing the use_count to 1, + # so that once torch finishes computation on the DLpack tensors, the memory can be freed. self._training_io_binding.clear_binding_outputs() return tuple(results) - return _ortmodule_output_transformation.populate_user_output_from_schema_and_outputs(self._original_module_output_schema, + return _ortmodule_io.populate_user_output_from_schema_and_outputs( + self._original_module_output_schema, self._onnx_graphs_info.user_output_names, _ORTModuleFunction.apply(*self._convert_training_graph_input_to_list(*inputs, **kwargs))) # Bind the forward method. self.forward = _forward.__get__(self) # Copy the forward signature from the PyTorch module. - functools.update_wrapper(self.forward.__func__, module.forward.__func__) + functools.update_wrapper( + self.forward.__func__, module.forward.__func__) super(ORTModule, self).__init__() + # Verbosity for logging + self._verbosity = Verbosity.WARNING + # Support contrib OPs register_custom_ops_pytorch_exporter.register_custom_op() @@ -236,13 +278,16 @@ class ORTModule(torch.nn.Module): self._original_module = module # Get the module that flattens the output from the original module into a tuple self._flattened_output_module = \ - _ortmodule_output_transformation.get_flattened_output_module(self._original_module) - self._original_module_parameters = signature(self._original_module.forward).parameters.values() + _ortmodule_io.get_flattened_output_module( + self._original_module) + self._original_module_parameters = signature( + self._original_module.forward).parameters.values() # TODO: remove after PyTorch ONNX exporter supports VAR_KEYWORD parameters. for input_parameter in self._original_module_parameters: if input_parameter.kind == inspect.Parameter.VAR_KEYWORD: - raise NotImplementedError("The model's forward method has **kwargs parameter which is currently not supported.") + raise NotImplementedError( + "The model's forward method has **kwargs parameter which is currently not supported.") self._onnx_inference = None self._is_training = True @@ -250,8 +295,8 @@ class ORTModule(torch.nn.Module): # 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 benefit of static shape, we can expose this flag to user - self._use_static_shape = False + # if we observe benefit of static shape, we can expose this flag to user + self._use_static_shape = False self._module_gradient_graph_builder = None self._input_names_require_grad = None self._original_module_output_schema = None @@ -270,34 +315,42 @@ class ORTModule(torch.nn.Module): self._save_onnx_prefix = '' from torch.utils.cpp_extension import ROCM_HOME - self.is_rocm_pytorch = (True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False) + self.is_rocm_pytorch = (True if ( + (torch.version.hip is not None) and (ROCM_HOME is not None)) else False) # CPP extension to get torch CUDA allocator's alloc and free function addresses # Disable external allocator for ROCM EP since external allocator is not supported yet. - self._use_external_cuda_allocator = (False if self.is_rocm_pytorch else True) + self._use_external_cuda_allocator = ( + False if self.is_rocm_pytorch else True) if self._use_external_cuda_allocator: - self._torch_cuda_allocator = _load_torch_allocator_cpp_extension() + self._torch_cuda_allocator = _load_torch_allocator_cpp_extension( + self._verbosity) self._torch_alloc = self._torch_cuda_allocator.cuda_caching_allocator_raw_alloc_address() self._torch_free = self._torch_cuda_allocator.cuda_caching_allocator_raw_delete_address() def _initialize_module_gradient_graph_builder(self): # TODO: PyTorch exporter bug: changes the initializer order in ONNX model - initializer_names = [p[0] for p in self._flattened_output_module.named_parameters()] - onnx_initializer_names = [p.name for p in self._onnx_inference.graph.initializer] - initializer_names = [p for p in initializer_names if p in onnx_initializer_names] + initializer_names = [p[0] + for p in self._flattened_output_module.named_parameters()] + onnx_initializer_names = { + p.name for p in self._onnx_inference.graph.initializer} + initializer_names = [ + p for p in initializer_names if p in onnx_initializer_names] # Build full training graph grad_builder_config = C.ModuleGradientGraphBuilderConfiguration() grad_builder_config.initializer_names_to_train = initializer_names grad_builder_config.input_names_require_grad = self._input_names_require_grad self._module_gradient_graph_builder = C.ModuleGradientGraphBuilder() - self._module_gradient_graph_builder.initialize(self._onnx_inference.SerializeToString(), grad_builder_config) + self._module_gradient_graph_builder.initialize( + self._onnx_inference.SerializeToString(), grad_builder_config) def _get_inference_graph_and_init_gradient_graph_builder(self, *inputs, **kwargs): self._onnx_inference = self._get_inference_graph(*inputs, **kwargs) if self._save_onnx: - onnx.save(self._onnx_inference, self._save_onnx_prefix + '_inference.onnx') + onnx.save(self._onnx_inference, + self._save_onnx_prefix + '_inference.onnx') self._initialize_module_gradient_graph_builder() @@ -306,10 +359,12 @@ class ORTModule(torch.nn.Module): provider_options = None if self._device.type == 'cuda': # Configure the InferenceSessions to use the specific GPU on which the model is placed. - providers = (["ROCMExecutionProvider"] if self.is_rocm_pytorch else ["CUDAExecutionProvider"]) + providers = (["ROCMExecutionProvider"] if self.is_rocm_pytorch else [ + "CUDAExecutionProvider"]) providers.append("CPUExecutionProvider") if self._use_external_cuda_allocator: - provider_options = [{"device_id": str(self._device.index), "cuda_external_alloc": str(self._torch_alloc), "cuda_external_free": str(self._torch_free)}, {}] + provider_options = [{"device_id": str(self._device.index), "cuda_external_alloc": str( + self._torch_alloc), "cuda_external_free": str(self._torch_free)}, {}] else: provider_options = [{"device_id": str(self._device.index)}, {}] elif self._device.type == 'cpu': @@ -322,7 +377,7 @@ class ORTModule(torch.nn.Module): # default to PRIORITY_BASED execution order session_options.execution_order = onnxruntime.ExecutionOrder.PRIORITY_BASED # 0:Verbose, 1:Info, 2:Warning. 3:Error, 4:Fatal. Default is 2. - session_options.log_severity_level = 2 + session_options.log_severity_level = int(self._verbosity) self._training_session = onnxruntime.training.TrainingAgent(self._onnx_training.SerializeToString(), session_options, providers, provider_options) @@ -331,19 +386,23 @@ class ORTModule(torch.nn.Module): self._run_options = C.RunOptions() # IO binding - # TODO: we should try to reuse the output buffers as some of the output tensors are same sizes, expecially the backward graph outputs. + # TODO: Reuse output buffers as some of output tensors have same shape, + # especially the backward graph outputs. self._training_io_binding = self._training_session.io_binding() def _build_training_graph(self, *inputs, **kwargs): if self._use_static_shape: - self._module_gradient_graph_builder.build(self._current_input_shape) + self._module_gradient_graph_builder.build( + self._current_input_shape) else: self._module_gradient_graph_builder.build() - self._onnx_training = onnx.load_model_from_string(self._module_gradient_graph_builder.get_training_model()) + self._onnx_training = onnx.load_model_from_string( + self._module_gradient_graph_builder.get_training_model()) self._onnx_graphs_info = self._module_gradient_graph_builder.get_training_graph_info() if self._save_onnx: - onnx.save(self._onnx_training, self._save_onnx_prefix + '_training.onnx') + onnx.save(self._onnx_training, + self._save_onnx_prefix + '_training.onnx') def eval(self: T) -> T: self._is_training = False @@ -375,7 +434,8 @@ class ORTModule(torch.nn.Module): result.append(inp) else: # TODO: Re-export ONNX if any input from _onnx_graphs_info.user_input_names is None. - raise RuntimeError(f'Input is present in ONNX graph but not provided: {name}.') + raise RuntimeError( + f'Input is present in ONNX graph but not provided: {name}.') # Initializers for param in self._flattened_output_module.named_parameters(): @@ -391,10 +451,10 @@ class ORTModule(torch.nn.Module): # Setup dynamic axes for onnx model input_names, dynamic_axes, self._input_names_require_grad, _ = \ - _ortmodule_output_transformation.parse_inputs_for_onnx_export( + _ortmodule_io.parse_inputs_for_onnx_export( self._original_module_parameters, None, *inputs, **kwargs) output_names, output_dynamic_axes, self._original_module_output_schema = \ - _ortmodule_output_transformation.parse_outputs_for_onnx_export_and_extract_output_schema( + _ortmodule_io.parse_outputs_for_onnx_export_and_extract_output_schema( self._original_module, inputs, kwargs) dynamic_axes.update(output_dynamic_axes) @@ -405,20 +465,23 @@ class ORTModule(torch.nn.Module): # NOTE: Inputs may contain tensors that have attributes preventing their deepcopy (example grad_fn). # Therefore, deepcopy only the data component of the input tensors for export. sample_inputs_copy, sample_kwargs_copy = \ - _ortmodule_output_transformation.deepcopy_model_input(*inputs, **kwargs) + _ortmodule_io.deepcopy_model_input( + *inputs, **kwargs) try: with torch.no_grad(): torch.onnx.export(self._flattened_output_module, - sample_inputs_copy + (sample_kwargs_copy, ), - f, - input_names=input_names, - output_names=output_names, - opset_version=ONNX_OPSET_VERSION, - do_constant_folding=False, - training=torch.onnx.TrainingMode.TRAINING, - dynamic_axes=dynamic_axes) + sample_inputs_copy + (sample_kwargs_copy, ), + f, + input_names=input_names, + output_names=output_names, + opset_version=ONNX_OPSET_VERSION, + do_constant_folding=False, + training=torch.onnx.TrainingMode.TRAINING, + dynamic_axes=dynamic_axes, + verbose=self._verbosity < Verbosity.WARNING) except RuntimeError as e: - raise RuntimeError('There was an error while exporting the PyTorch model to ONNX: {}'.format(e)) + raise RuntimeError( + 'There was an error while exporting the PyTorch model to ONNX: {}'.format(e)) return onnx.load_model_from_string(f.getvalue()) diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index ba1bc66a81..61e92d4a1a 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -1288,3 +1288,26 @@ def test_forward_data_and_model_on_different_devices(data_device, model_device): with pytest.raises(RuntimeError) as runtime_error: ort_model(x) assert f"Input argument to forward found on device {torch.device(x.device)}, but expected it to be on module device {ort_model._device}." in str(runtime_error.value) + +def test_forward_returns_none_type_as_output(): + class NeuralNetNoneTypeOutput(torch.nn.Module): + def __init__(self, input_size, num_classes): + super(NeuralNetNoneTypeOutput, self).__init__() + + self.fc1 = torch.nn.Linear(input_size, num_classes) + self.relu1 = torch.nn.ReLU() + + def forward(self, input1): + out1 = self.fc1(input1) + out1 = self.relu1(out1) + return {'out': out1, 'none_output': None} + + device = 'cuda' + N, D_in, H, D_out = 64, 784, 500, 10 + model = NeuralNetNoneTypeOutput(D_in, D_out).to(device) + model = ORTModule(model) + x = torch.randn(N, D_in, device=device) + output = model(x) + + assert output['out'] is not None + assert output['none_output'] is None diff --git a/orttraining/orttraining/test/python/perf_log/ort_module_perf_test_tools.py b/orttraining/orttraining/test/python/perf_log/ort_module_perf_test_tools.py new file mode 100644 index 0000000000..0bd5e70e3c --- /dev/null +++ b/orttraining/orttraining/test/python/perf_log/ort_module_perf_test_tools.py @@ -0,0 +1,206 @@ +# https://docs.microsoft.com/en-us/azure/mysql/connect-python + +import mysql.connector +from mysql.connector import errorcode +import git +import os + +import argparse +from datetime import datetime + +def get_repo_commit(repo_path): + repo = git.Repo(repo_path, search_parent_directories=True) + sha = repo.head.object.hexsha + short_sha = repo.git.rev_parse(sha, short=4) + return short_sha + +create_table_script = "CREATE TABLE perf_test_training_ort_module_data (\ + id int(11) NOT NULL AUTO_INCREMENT,\ + Model varchar(64) COLLATE utf8_bin DEFAULT NULL,\ + BatchId varchar(32) COLLATE utf8_bin DEFAULT NULL,\ + CommitId varchar(32) COLLATE utf8_bin DEFAULT NULL,\ + ModelName varchar(256) COLLATE utf8_bin DEFAULT NULL,\ + DisplayName varchar(512) COLLATE utf8_bin DEFAULT NULL,\ + UseMixedPrecision tinyint(1) DEFAULT NULL,\ + UseAutoCast tinyint(1) DEFAULT NULL,\ + UseDeepSpeed tinyint(1) DEFAULT NULL,\ + Optimizer varchar(32) COLLATE utf8_bin DEFAULT NULL,\ + BatchSize int(11) DEFAULT NULL,\ + SeqLen int(11) DEFAULT NULL,\ + PredictionsPerSeq int(11) DEFAULT NULL,\ + NumOfBatches int(11) DEFAULT NULL,\ + WeightUpdateSteps int(11) DEFAULT NULL,\ + Round int(11) DEFAULT NULL,\ + GradAccSteps int(11) DEFAULT NULL,\ + AvgTimePerBatch float DEFAULT NULL,\ + Throughput float DEFAULT NULL,\ + StabilizedThroughput float DEFAULT NULL,\ + EndToEndThroughput float DEFAULT NULL,\ + TotalTime float DEFAULT NULL,\ + AvgCPU int(11) DEFAULT NULL,\ + Memory int(11) DEFAULT NULL,\ + RunConfig varchar(2048) COLLATE utf8_bin DEFAULT NULL,\ + Time datetime DEFAULT NULL,\ + PRIMARY KEY (id),\ + UNIQUE KEY config_unique (Model,BatchId,CommitId,UseMixedPrecision,UseAutoCast,UseDeepSpeed,Optimizer,BatchSize,SeqLen,ModelName)\ +) ENGINE=InnoDB AUTO_INCREMENT=1696 DEFAULT CHARSET=utf8 COLLATE=utf8_bin;" + +insert_table_script = "INSERT INTO onnxruntime.perf_test_training_ort_module_data\ + (\ + Model,\ + BatchId,\ + CommitId,\ + ModelName,\ + DisplayName,\ + UseMixedPrecision,\ + UseAutoCast,\ + UseDeepSpeed,\ + Optimizer,\ + BatchSize,\ + SeqLen,\ + PredictionsPerSeq,\ + NumOfBatches,\ + WeightUpdateSteps,\ + Round,\ + GradAccSteps,\ + AvgTimePerBatch,\ + Throughput,\ + StabilizedThroughput,\ + EndToEndThroughput,\ + TotalTime,\ + AvgCPU,\ + Memory,\ + RunConfig,\ + Time)\ + VALUES\ + (\ + %(Model)s,\ + %(BatchId)s,\ + %(CommitId)s,\ + %(ModelName)s,\ + %(DisplayName)s,\ + %(UseMixedPrecision)s,\ + %(UseAutoCast)s,\ + %(UseDeepSpeed)s,\ + %(Optimizer)s,\ + %(BatchSize)s,\ + %(SeqLen)s,\ + %(PredictionsPerSeq)s,\ + %(NumOfBatches)s,\ + %(WeightUpdateSteps)s,\ + %(Round)s,\ + %(GradAccSteps)s,\ + %(AvgTimePerBatch)s,\ + %(Throughput)s,\ + %(StabilizedThroughput)s,\ + %(EndToEndThroughput)s,\ + %(TotalTime)s,\ + %(AvgCPU)s,\ + %(Memory)s,\ + %(RunConfig)s,\ + %(Time)s)" + +# Obtain connection string information from the portal +def connect_to_perf_dashboard_db(mysql_server_name, power_bi_user_name, password, database): + config = { + 'host': mysql_server_name, + 'user': power_bi_user_name, + 'password': password, + 'database': database, + } + + try: + conn = mysql.connector.connect(**config) + print("Connection established") + return conn + except mysql.connector.Error as err: + if err.errno == errorcode.ER_ACCESS_DENIED_ERROR: + print("Something is wrong with the user name or password") + elif err.errno == errorcode.ER_BAD_DB_ERROR: + print("Database does not exist") + else: + print(err) + +def log_perf_metrics(perf_metrics, + mysql_server_name, power_bi_user_name, power_bi_password, power_bi_database, perf_repo_path=None): + if perf_repo_path: + perf_metrics['CommitId'] = get_repo_commit(perf_repo_path) + else: + perf_metrics['CommitId'] = get_repo_commit(os.path.realpath(__file__)) + + connect_and_insert_perf_metrics( + mysql_server_name, + power_bi_user_name, + power_bi_password, + power_bi_database, + perf_metrics) + +required_attributes_for_perf_metrics = ['model_name', 'optimizer', 'batch_size', 'epochs', 'train_steps', + 'sequence_length'] + +def calculate_and_log_perf_metrics(args, start_time, + mysql_server_name, power_bi_user_name, power_bi_password, power_bi_database, ort_repo_path=None): + completion_time = datetime.datetime.now() + perf_metrics_duration = completion_time - start_time + + for attribute in required_attributes_for_perf_metrics: + if not hasattr(args, attribute): + raise ValueError('args does not contain all attributes needed to calculate perf metrics. \ + Please prepare perf_metrics and call log_perf_metrics instead') + + perf_metrics = {} + perf_metrics['Model'] = args.model_name + perf_metrics['BatchId'] = 'NA' + perf_metrics['ModelName'] = args.model_name + perf_metrics['DisplayName'] = args.model_name + perf_metrics['UseMixedPrecision'] = args.fp16 if hasattr(args, 'fp16') else False + perf_metrics['UseAutoCast'] = args.use_auto_cast if hasattr(args, 'use_auto_cast') else False + perf_metrics['UseDeepSpeed'] = args.use_deep_speed if hasattr(args, 'use_deep_speed') else False + perf_metrics['Optimizer'] = args.optimizer + perf_metrics['BatchSize'] = args.batch_size + perf_metrics['SeqLen'] = args.sequence_length + perf_metrics['PredictionsPerSeq'] = args.prediction_per_seq if hasattr(args, 'prediction_per_seq') else 0 + perf_metrics['NumOfBatches'] = args.epochs * args.train_steps + perf_metrics['WeightUpdateSteps'] = args.epochs * args.train_steps + perf_metrics['Round'] = 0 # NA + perf_metrics['GradAccSteps'] = args.gradient_accumulation_steps + + perf_metrics['AvgTimePerBatch'] = \ + perf_metrics_duration.microseconds / args.train_steps + + perf_metrics['Throughput'] = \ + args.batch_size * args.train_steps / perf_metrics_duration.seconds + + perf_metrics['StabilizedThroughput'] = 0 # TODO + perf_metrics['EndToEndThroughput'] = 0 # TODO + perf_metrics['TotalTime'] = perf_metrics_duration.seconds + + perf_metrics['AvgCPU'] = 0 # TODO + perf_metrics['Memory'] = 0 # TODO + perf_metrics['RunConfig'] = 'na' + perf_metrics['Time'] = completion_time.strftime("%Y-%m-%d %H:%M:%S") + + log_perf_metrics(perf_metrics, mysql_server_name, power_bi_user_name, power_bi_password, power_bi_database, + ort_repo_path) + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--mysql_server_name', help='Perf dashboard mysql server name') + parser.add_argument('--power_bi_user_name', help='Power BI user name') + parser.add_argument('--password', help='password', default=None) + parser.add_argument('--database', help='The dashboard database') + return parser.parse_args() + +def connect_and_insert_perf_metrics(mysql_server_name, power_bi_user_name, password, database, perf_metrics): + conn = connect_to_perf_dashboard_db(mysql_server_name, power_bi_user_name, password, database) + # https://dev.mysql.com/doc/connector-python/en/connector-python-api-mysqlcursor-execute.html + conn.cursor().execute(insert_table_script, perf_metrics) + conn.commit() + conn.cursor().close() + conn.close() + print("perf_metrics logged into power-bi database.") + +if __name__ == '__main__': + args = parse_arguments() + conn = connect_to_perf_dashboard_db(args.mysql_server_name, args.power_bi_user_name, args.password, args.database) + conn.cursor().execute(create_table_script)