saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix ortmodule for the pytorch model with ort device (#9927)

Browse source 
* add ortmodule and eager mode test

* add ortmodule dependency

* convert between aten ort tensor and ortvalue

* register the EP to ortmodule using ort device information

* remove duplicated test

* remove useless dependency

* handle half precision type for ortmodule outputs

* adjust the tensor conversion python code

Co-authored-by: Cheng Tang <chenta@microsoft.com@orttrainingdev9.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
This commit is contained in:
Tang, Cheng 2021-12-08 13:04:29 -08:00 committed by GitHub
parent fb30e9fdae
commit 0adeb86bfd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 120 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
orttraining/orttraining/eager/ort_backends.cpp
View file

@ -56,6 +56,9 @@ onnxruntime::Status ORTBackendsManager::set_device(size_t device_index, const st
custom_op_schema_);
backends_[device_index] = std::move(invoker);
ProviderInfoMap provider_info;
provider_info[provider_type] = provider_options;
device_ep_info_[device_index] = provider_info;
return onnxruntime::Status::OK();
}
@ -66,6 +69,22 @@ OrtDevice ORTBackendsManager::GetOrtDeviceInfo(size_t torch_device_index){
return allocator->Info().device;
}
size_t ORTBackendsManager::GetOrtDeviceIndex(const OrtMemoryInfo& ort_memory_info){
for (auto it = backends_.begin(); it != backends_.end(); ++it){
//eager mode currently only operate on EP's default memory type
auto allocator = it->second->GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault);
if (allocator->Info() == ort_memory_info)
return it->first;
}
ORT_THROW("Can't find the eager mode ORT device index for target ort tensor");
}
const ProviderInfoMap& ORTBackendsManager::GetOrtDeviceProviderInfo(size_t torch_device_index) const {
auto lookup = device_ep_info_.find(torch_device_index);
ORT_ENFORCE(lookup != device_ep_info_.end());
return lookup->second;
}
onnxruntime::ORTInvoker& ORTBackendsManager::GetInvoker(const at::Device device) {
ORT_LOG_FN(device);

9
orttraining/orttraining/eager/ort_backends.h
View file

@ -15,6 +15,8 @@
namespace torch_ort {
namespace eager {
using ProviderInfoMap = std::unordered_map<std::string, onnxruntime::ProviderOptions >;
class ORTBackendsManager {
public:
ORTBackendsManager(const onnxruntime::logging::Logger& logger);
@ -26,12 +28,19 @@ public:
OrtDevice GetOrtDeviceInfo(size_t torch_device_index);
size_t GetOrtDeviceIndex(const OrtMemoryInfo& ort_memory_info);
const ProviderInfoMap& GetOrtDeviceProviderInfo(size_t torch_device_index) const;
private:
std::map<at::DeviceIndex, std::unique_ptr<onnxruntime::ORTInvoker>> backends_;
const onnxruntime::logging::Logger& logger_;
//custom op schema registry
//TODO: we might want to support load custom op schema on the fly
onnxruntime::IOnnxRuntimeOpSchemaRegistryList custom_op_schema_ = {};
// record the device associated provider information, so ortmodule can restore the ep
std::unordered_map<at::DeviceIndex, ProviderInfoMap> device_ep_info_;
};
ORTBackendsManager& GetORTBackendsManager();

51
orttraining/orttraining/eager/ort_eager.cpp
View file

@ -11,6 +11,7 @@
#include "orttraining/core/framework/ortmodule_graph_builder.h"
#include "ort_customops.h"
#include "torch/csrc/autograd/python_variable.h"
#include "core/framework/tensor.h"
namespace onnxruntime{
namespace python{
@ -18,19 +19,42 @@ namespace python{
using namespace onnxruntime::training;
using namespace torch_ort::eager;
py::object ORTTensor_toDLPack(const at::Tensor& data)
{
OrtValue ort_value = torch_ort::eager::create_ort_value(data);
return py::reinterpret_steal<py::object>(onnxruntime::training::framework::torch::ToDlpack(ort_value));
static at::ScalarType aten_scalar_type_from_ort(
onnxruntime::MLDataType dtype) {
if (dtype == onnxruntime::DataTypeImpl::GetType<float>())
return at::kFloat;
else if (dtype == onnxruntime::DataTypeImpl::GetType<double>())
return at::kDouble;
else if (dtype == onnxruntime::DataTypeImpl::GetType<onnxruntime::MLFloat16>())
return at::kHalf;
else if (dtype == onnxruntime::DataTypeImpl::GetType<onnxruntime::BFloat16>())
return at::kBFloat16;
else if (dtype == onnxruntime::DataTypeImpl::GetType<int>())
return at::kInt;
else if (dtype == onnxruntime::DataTypeImpl::GetType<int16_t>())
return at::kShort;
else if (dtype == onnxruntime::DataTypeImpl::GetType<int64_t>())
return at::kLong;
else if (dtype == onnxruntime::DataTypeImpl::GetType<bool>())
return at::kBool;
else
ORT_THROW("Unsupport aten scalar type: ", dtype);
}
at::Tensor ORTTensor_FromDLPack(const py::object& dlpack_tensor)
OrtValue ORTTensor_toORTValue(const at::Tensor& data)
{
OrtValue ort_value = onnxruntime::training::framework::torch::FromDlpack(dlpack_tensor.ptr(), false);
return torch_ort::eager::create_ort_value(data);
}
at::Tensor OrtValue_To_ATen_Tensor(OrtValue& ortvalue)
{
auto& ort_tensor = ortvalue.Get<Tensor>();
size_t ort_device_idx = torch_ort::eager::GetORTBackendsManager().GetOrtDeviceIndex(ort_tensor.Location());
return torch_ort::eager::aten_tensor_from_ort(
std::move(ort_value),
std::move(ortvalue),
at::TensorOptions()
.device(at::Device(at::DeviceType::ORT, 0)));
.device(at::Device(at::DeviceType::ORT, ort_device_idx))
.dtype(onnxruntime::python::aten_scalar_type_from_ort(ort_tensor.DataType())));
}
void addObjectMethodsForEager(py::module& m){
@ -44,11 +68,11 @@ void addObjectMethodsForEager(py::module& m){
},
py::arg("device_index") = 0);
m.def("ort_to_dlpack", [](at::Tensor data) {
return ORTTensor_toDLPack(data);
m.def("aten_ort_tensor_to_ort_value", [](at::Tensor data) {
return ORTTensor_toORTValue(data);
});
m.def("ort_from_dlpack", [](py::object dlpack_tensor) {
return ORTTensor_FromDLPack(dlpack_tensor);
m.def("to_aten_ort_device_tensor", [](OrtValue& ortvalue) {
return OrtValue_To_ATen_Tensor(ortvalue);
});
m.def("set_device", [](size_t device_index,
@ -61,6 +85,9 @@ void addObjectMethodsForEager(py::module& m){
m.def("get_ort_device", [](size_t torch_device_index){
return torch_ort::eager::GetORTBackendsManager().GetOrtDeviceInfo(torch_device_index);
});
m.def("get_ort_device_provider_info", [](size_t torch_device_index){
return torch_ort::eager::GetORTBackendsManager().GetOrtDeviceProviderInfo(torch_device_index);
});
auto customop_module = m.def_submodule("custom_ops");
torch_ort::eager::GenerateCustomOpsBindings(customop_module);

20
orttraining/orttraining/eager/test/linux_only_ortmodule_eager_test.py
View file

@ -25,7 +25,25 @@ class NeuralNet(nn.Module):
out = self.fc2(out)
return out
class NoOpNet(torch.nn.Module):
def __init__(self):
super(NoOpNet, self).__init__()
self.dummy_weight = torch.nn.Parameter(
torch.ones(128, 128, dtype=torch.float16))
def forward(self, input):
return input
class OrtModuleEagerTest(unittest.TestCase):
def test_half_type(self):
model = NoOpNet()
device = torch.device('ort')
model.to(device)
model = ORTModule(model)
input = torch.ones(2,2).to(torch.float16)
y = model(input.to(device))
assert(y.dtype == torch.float16)
def test_ort_module_and_eager_mode(self):
input_size = 784
hidden_size = 500
@ -67,4 +85,4 @@ class OrtModuleEagerTest(unittest.TestCase):
assert torch.allclose(cpu_updated_state[state_tensor], ort_updated_state[state_tensor].cpu(), atol=1e-3)
if __name__ == '__main__':
unittest.main()
unittest.main()

5
orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
View file

@ -274,6 +274,11 @@ class GraphExecutionManager(GraphExecutionInterface):
elif self._device.type == 'cpu':
providers = ["CPUExecutionProvider"]
provider_options = [{}]
elif self._device.type == 'ort':
provider_info = C.get_ort_device_provider_info(self._device.index)
assert len(provider_info.keys()) == 1
providers = list(provider_info.keys())
provider_options = [provider_info[providers[0]]]
session_options = onnxruntime.SessionOptions()
session_options.enable_mem_pattern = False

18
orttraining/orttraining/python/training/ortmodule/_training_manager.py
View file

@ -47,8 +47,11 @@ class TrainingManager(GraphExecutionManager):
forward_inputs = C.OrtValueVector()
forward_inputs.reserve(len(inputs))
for input in inputs:
dlp = _utils._torch_tensor_to_dlpack(input)
forward_inputs.push_back(dlp, input.dtype == torch.bool)
if input.device.type == 'ort':
forward_inputs.push_back(C.aten_ort_tensor_to_ort_value(input))
else:
valid_ort_tensor = _utils._torch_tensor_to_dlpack(input)
forward_inputs.push_back(valid_ort_tensor, input.dtype == torch.bool)
forward_outputs = C.OrtValueVector()
# Run and return module outputs.
@ -145,7 +148,10 @@ class TrainingManager(GraphExecutionManager):
grad_output = torch.tensor(0., device=device, dtype=dtype)
elif not grad_output.is_contiguous():
grad_output = grad_output.contiguous()
backward_inputs.push_back(_utils._torch_tensor_to_dlpack(grad_output),
if grad_output.device.type == 'ort':
backward_inputs.push_back(C.aten_ort_tensor_to_ort_value(grad_output))
else:
backward_inputs.push_back(_utils._torch_tensor_to_dlpack(grad_output),
grad_output.dtype is torch.bool)
backward_inputs.shrink_to_fit()
@ -163,8 +169,7 @@ class TrainingManager(GraphExecutionManager):
for input_name in self._graph_info.user_input_names:
# Append to the results the backward output for each input that required grad
if input_name in require_grad_names_set:
results.append(_utils._torch_tensor_from_dl_pack(
backward_outputs.dlpack_at(require_grad_names_index),
results.append(_utils._ortvalue_to_torch_tensor(
backward_outputs[require_grad_names_index], self._device))
require_grad_names_index += 1
else:
@ -177,8 +182,7 @@ class TrainingManager(GraphExecutionManager):
initializer_index = num_user_input_grads
for initializer_name in self._graph_info.initializer_names:
if initializer_name in self._graph_initializer_names_to_train:
results.append(_utils._torch_tensor_from_dl_pack(
backward_outputs.dlpack_at(initializer_index),
results.append(_utils._ortvalue_to_torch_tensor(
backward_outputs[initializer_index], self._device))
initializer_index += 1
else:

34
orttraining/orttraining/python/training/ortmodule/_utils.py
View file

@ -38,33 +38,33 @@ def _ortvalue_from_torch_tensor(torch_tensor):
return C.OrtValue.from_dlpack(to_dlpack(torch_tensor), is_bool_tensor)
def _torch_tensor_from_dl_pack(dlpack, ortvalue, device):
torch_tensor = from_dlpack(dlpack) if device.type != 'ort' else C.ort_from_dlpack(dlpack)
def _torch_tensor_from_dl_pack(dlpack, ortvalue):
torch_tensor = from_dlpack(dlpack)
return torch_tensor.to(torch.bool) if ortvalue.data_type() == 'tensor(bool)' else torch_tensor
def _ortvalue_to_torch_tensor(ortvalue, device):
if device.type == 'ort':
return C.to_aten_ort_device_tensor(ortvalue)
# 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 we need to convert the torch tensor to torch.bool type if OrtValue is bool tensor.
dlpack_tensor = ortvalue.to_dlpack()
return _torch_tensor_from_dl_pack(dlpack_tensor, ortvalue, device)
return _torch_tensor_from_dl_pack(dlpack_tensor, ortvalue)
# convert torch tensor to ort accept tensors, could be dlpack or ortvalue
def _torch_tensor_to_dlpack(tensor):
if tensor.device.type == 'ort':
return C.ort_to_dlpack(tensor)
else:
# 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.
# DLPack is discussing how to support bool type, we can remove this workaround once both DLPack
# and PyTorch support bool type.
if not tensor.is_contiguous():
raise ORTModuleIOError(
"Only contiguous tensors are supported.")
if tensor.dtype == torch.bool and LooseVersion(torch.__version__) >= LooseVersion('1.10.0'):
tensor = tensor.to(torch.uint8)
return to_dlpack(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.
# DLPack is discussing how to support bool type, we can remove this workaround once both DLPack
# and PyTorch support bool type.
if not tensor.is_contiguous():
raise ORTModuleIOError(
"Only contiguous tensors are supported.")
if tensor.dtype == torch.bool and LooseVersion(torch.__version__) >= LooseVersion('1.10.0'):
tensor = tensor.to(torch.uint8)
return to_dlpack(tensor)
def _check_same_device(device, argument_str, *args):

1
tools/ci_build/github/windows/eager/requirements.txt
View file

@ -3,3 +3,4 @@ wheel
numpy
typing_extensions
torch