From 0adeb86bfda43b42b3351519a648a1824cbdbfff Mon Sep 17 00:00:00 2001 From: "Tang, Cheng" Date: Wed, 8 Dec 2021 13:04:29 -0800 Subject: [PATCH] Fix ortmodule for the pytorch model with ort device (#9927) * 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 --- .../orttraining/eager/ort_backends.cpp | 19 +++++++ orttraining/orttraining/eager/ort_backends.h | 9 ++++ orttraining/orttraining/eager/ort_eager.cpp | 51 ++++++++++++++----- .../test/linux_only_ortmodule_eager_test.py | 20 +++++++- .../ortmodule/_graph_execution_manager.py | 5 ++ .../training/ortmodule/_training_manager.py | 18 ++++--- .../python/training/ortmodule/_utils.py | 34 ++++++------- .../github/windows/eager/requirements.txt | 1 + 8 files changed, 120 insertions(+), 37 deletions(-) diff --git a/orttraining/orttraining/eager/ort_backends.cpp b/orttraining/orttraining/eager/ort_backends.cpp index d552393514..e14ba9d2a1 100644 --- a/orttraining/orttraining/eager/ort_backends.cpp +++ b/orttraining/orttraining/eager/ort_backends.cpp @@ -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); diff --git a/orttraining/orttraining/eager/ort_backends.h b/orttraining/orttraining/eager/ort_backends.h index 26cc495a61..4f3d885a5e 100644 --- a/orttraining/orttraining/eager/ort_backends.h +++ b/orttraining/orttraining/eager/ort_backends.h @@ -15,6 +15,8 @@ namespace torch_ort { namespace eager { +using ProviderInfoMap = std::unordered_map; + 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> 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 device_ep_info_; }; ORTBackendsManager& GetORTBackendsManager(); diff --git a/orttraining/orttraining/eager/ort_eager.cpp b/orttraining/orttraining/eager/ort_eager.cpp index e3fdc332f1..f679335258 100644 --- a/orttraining/orttraining/eager/ort_eager.cpp +++ b/orttraining/orttraining/eager/ort_eager.cpp @@ -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(onnxruntime::training::framework::torch::ToDlpack(ort_value)); +static at::ScalarType aten_scalar_type_from_ort( + onnxruntime::MLDataType dtype) { + if (dtype == onnxruntime::DataTypeImpl::GetType()) + return at::kFloat; + else if (dtype == onnxruntime::DataTypeImpl::GetType()) + return at::kDouble; + else if (dtype == onnxruntime::DataTypeImpl::GetType()) + return at::kHalf; + else if (dtype == onnxruntime::DataTypeImpl::GetType()) + return at::kBFloat16; + else if (dtype == onnxruntime::DataTypeImpl::GetType()) + return at::kInt; + else if (dtype == onnxruntime::DataTypeImpl::GetType()) + return at::kShort; + else if (dtype == onnxruntime::DataTypeImpl::GetType()) + return at::kLong; + else if (dtype == onnxruntime::DataTypeImpl::GetType()) + 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(); + 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); diff --git a/orttraining/orttraining/eager/test/linux_only_ortmodule_eager_test.py b/orttraining/orttraining/eager/test/linux_only_ortmodule_eager_test.py index d6415f1353..8134a116bf 100644 --- a/orttraining/orttraining/eager/test/linux_only_ortmodule_eager_test.py +++ b/orttraining/orttraining/eager/test/linux_only_ortmodule_eager_test.py @@ -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() \ No newline at end of file + unittest.main() diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 874b0c605e..ede9ffb2ab 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -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 diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index a66ed955d4..35f8208114 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -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: diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py index d5bb43951f..131c8fa80c 100644 --- a/orttraining/orttraining/python/training/ortmodule/_utils.py +++ b/orttraining/orttraining/python/training/ortmodule/_utils.py @@ -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): diff --git a/tools/ci_build/github/windows/eager/requirements.txt b/tools/ci_build/github/windows/eager/requirements.txt index 3dbc855e49..557d5efd0b 100644 --- a/tools/ci_build/github/windows/eager/requirements.txt +++ b/tools/ci_build/github/windows/eager/requirements.txt @@ -3,3 +3,4 @@ wheel numpy typing_extensions torch +