mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-12 17:57:38 +00:00
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 <liqun@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net> * 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 <liqfu@microsoft.com> Co-authored-by: baijumeswani <bmeswani@microsoft.com>
This commit is contained in:
Thiago Crepaldi
GitHub
parent
4fd9fef9ee
commit
c60ef62190
10 changed files with 394 additions and 94 deletions
|
|
@ -132,7 +132,7 @@ std::unordered_set<NodeIndex> 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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -115,7 +115,7 @@ common::Status TrainingAgent::RunBackward(int64_t run_id, const std::vector<OrtV
|
|||
}
|
||||
|
||||
void TrainingAgent::CancelPendingBackwardRun(int64_t run_id) {
|
||||
LOGS(*inference_session_->GetLogger(), 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);
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -506,9 +506,9 @@ py::class_<TrainingAgent>(m, "TrainingAgent", R"pbdoc(This is the main class use
|
|||
.def_readwrite("use_invertible_layernorm_grad",
|
||||
&ModuleGradientGraphBuilderConfiguration::use_invertible_layernorm_grad);
|
||||
|
||||
py::class_<TrainingGraphInfo> split_graphs_info(m, "TrainingGraphInfo",
|
||||
py::class_<TrainingGraphInfo> 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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 <torch/extension.h>
|
||||
#include <c10/cuda/CUDACachingAllocator.h>
|
||||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
Loading…
Reference in a new issue