From e05b15175db366200deef19e2cba00b489949bd8 Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Tue, 18 May 2021 12:57:05 -0700 Subject: [PATCH] Add cpp ext lock file check during ORTModule init (#7740) * Add cpp ext lock file check during ORTModule init * Address comments --- .../python/training/ortmodule/__init__.py | 18 +- .../training/ortmodule/_cpp_extensions.py | 185 ++++++++++++++++++ .../ortmodule/_graph_execution_manager.py | 8 +- .../python/training/ortmodule/_utils.py | 165 +--------------- 4 files changed, 205 insertions(+), 171 deletions(-) create mode 100644 orttraining/orttraining/python/training/ortmodule/_cpp_extensions.py diff --git a/orttraining/orttraining/python/training/ortmodule/__init__.py b/orttraining/orttraining/python/training/ortmodule/__init__.py index a9f373d7a2..5bdcd39840 100644 --- a/orttraining/orttraining/python/training/ortmodule/__init__.py +++ b/orttraining/orttraining/python/training/ortmodule/__init__.py @@ -3,14 +3,23 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +import os from packaging import version -# All global constant goes here, before ORTModule is imported +################################################################################ +# All global constant goes here, before ORTModule is imported ################## +################################################################################ ONNX_OPSET_VERSION = 12 MINIMUM_TORCH_VERSION_STR = '1.8.1' +TORCH_CPP_BUILD_DIR = os.path.join(os.path.dirname(__file__),'torch_inline_extensions') -from .ortmodule import ORTModule - +# Check whether Torch C++ extension compilation was aborted in previous runs +if not os.path.exists(TORCH_CPP_BUILD_DIR): + os.makedirs(TORCH_CPP_BUILD_DIR, exist_ok = True) +elif os.path.exists(os.path.join(TORCH_CPP_BUILD_DIR,'lock')): + print("WARNING: ORTModule detected PyTorch CPP extension's lock file during initialization, " + "which can cause unexpected hangs. " + f"Delete {os.path.join(TORCH_CPP_BUILD_DIR,'lock')} to prevent unexpected behavior.") # Verify proper PyTorch is installed before proceding to ONNX Runtime initializetion try: @@ -23,3 +32,6 @@ try: f'but version {torch.__version__} was found instead.') except: raise(f'PyTorch {MINIMUM_TORCH_VERSION_STR} must be installed in order to run ONNXRuntime ORTModule frontend!') + +# ORTModule must be loaded only after all validation passes +from .ortmodule import ORTModule diff --git a/orttraining/orttraining/python/training/ortmodule/_cpp_extensions.py b/orttraining/orttraining/python/training/ortmodule/_cpp_extensions.py new file mode 100644 index 0000000000..ad775412cd --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/_cpp_extensions.py @@ -0,0 +1,185 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Support for PyTorch C++ extensions within ORTModule + +IMPORTANT: All extensions must explicitly use TORCH_CPP_BUILD_DIR as `build_directory` + to allow ORTModule to monitor TORCH_CPP_BUILD_DIR/lock and warn the user + when abnormal initialization occurs + +TODO: Implement mechanism to register extensions and prevent issues with incorrect/missing flags + for each :meth:`torch.utils.cpp_extension.load_inline` call +""" + +import threading +from functools import wraps +from torch.utils.cpp_extension import load_inline + +from onnxruntime.capi import _pybind_state as C +from onnxruntime.training.ortmodule import TORCH_CPP_BUILD_DIR + + +def _load_torch_gpu_allocator_cpp_extension(verbosity, is_rocm_pytorch): + gpu_identifier = "hip" if is_rocm_pytorch else "cuda" + gpu_allocator_header = "HIPCachingAllocator" if is_rocm_pytorch else "CUDACachingAllocator" + torch_gpu_allocator_addresses_cpp_source = f''' + #include + #include + + size_t gpu_caching_allocator_raw_alloc_address() {{ + return reinterpret_cast(&c10::{gpu_identifier}::{gpu_allocator_header}::raw_alloc); + }} + + size_t gpu_caching_allocator_raw_delete_address() {{ + return reinterpret_cast(&c10::{gpu_identifier}::{gpu_allocator_header}::raw_delete); + }} + ''' + + return load_inline(name='torch_allocator', + cpp_sources=[torch_gpu_allocator_addresses_cpp_source], + extra_cflags=['-D__HIP_PLATFORM_HCC__=1' if is_rocm_pytorch else ''], + functions=['gpu_caching_allocator_raw_alloc_address', + 'gpu_caching_allocator_raw_delete_address'], + verbose=verbosity, + with_cuda=True, + build_directory=TORCH_CPP_BUILD_DIR) + +def run_once_aten_op_executor(f): + """ + Decorator to run a function only once. + :param f: function to be run only once during execution time despite the number of calls + :return: The original function with the params passed to it if it hasn't already been run before + """ + @wraps(f) + def aten_op_executor_wrapper(*args, **kwargs): + if not aten_op_executor_wrapper.has_run: + with aten_op_executor_wrapper.lock: + if not aten_op_executor_wrapper.has_run: + aten_op_executor_wrapper.has_run = True + return f(*args, **kwargs) + + aten_op_executor_wrapper.lock = threading.Lock() + aten_op_executor_wrapper.has_run = False + return aten_op_executor_wrapper + +@run_once_aten_op_executor +def _load_aten_op_executor_cpp_extension(verbosity, is_rocm_pytorch): + aten_op_executor_cpp_source = """ +#include +#include +#include +#include +#include + +class ATenOperatorCache { + public: + static ATenOperatorCache& Instance() { + static ATenOperatorCache instance; + return instance; + } + + std::shared_ptr GetOperator(const std::string& op_name) { + if (ops_.find(op_name) == ops_.end()) { + auto& ops = torch::jit::getAllOperatorsFor(torch::jit::Symbol::fromQualString(op_name)); + TORCH_INTERNAL_ASSERT(ops.size() == 1); + ops_[op_name] = ops.front(); + } + + return ops_.at(op_name); + } + + private: + ATenOperatorCache() = default; + std::unordered_map> ops_; +}; + +// Some arguments of backward operator are not from forward operator's input or output, +// but need some processing. Since we cannot build such processing to ONNX graph for now, +// we are putting such processing code here if needed. +// Take embedding_backward as example: +// weight: embedding_backward(grad, indices, weight.size(0), padding_idx, scale_grad_by_freq, sparse) +// the 3rd argument (index 2) is weight.size(0), we add this processing here. +using TensorTransformFunc = std::function; +static const TensorTransformFunc embedding_num_weights = [](const at::Tensor& tensor) { + return c10::IValue(tensor.size(0)); +}; + +static const std::unordered_map> TENSOR_TRANSFORM_FUNCS = { + {"aten::embedding_backward", {{2, embedding_num_weights}}}, +}; + +template +void SetIValueArguments(const std::vector>& raw_arguments, + std::vector& ivalue_arguments) { + for (size_t i = 0; i < raw_arguments.size(); i++) { + size_t index = std::get<0>(raw_arguments[i]); + TORCH_INTERNAL_ASSERT(index < ivalue_arguments.size()); + ivalue_arguments[index] = c10::IValue(std::get<1>(raw_arguments[i])); + } +} + +// TODO: Add more argument types, such as list type. +std::vector ExecuteATenOperator( + const char* op_name, const std::vector>& tensor_arguments, + const std::vector>& int_arguments, + const std::vector>& float_arguments, + const std::vector>& bool_arguments) { + std::string op_name_str(op_name); + std::shared_ptr op = ATenOperatorCache::Instance().GetOperator(op_name_str); + + // TODO: need to handle optional argument and arguments with default values. + std::vector arguments; + arguments.resize(op->schema().arguments().size()); + for (size_t i = 0; i < tensor_arguments.size(); i++) { + size_t index = std::get<0>(tensor_arguments[i]); + at::Tensor tensor = at::fromDLPack(std::get<1>(tensor_arguments[i])); + bool has_transform_func = false; + if (TENSOR_TRANSFORM_FUNCS.find(op_name_str) != TENSOR_TRANSFORM_FUNCS.end()) { + const auto& transform_funcs = TENSOR_TRANSFORM_FUNCS.at(op_name_str); + if (transform_funcs.find(index) != transform_funcs.end()) { + arguments[index] = transform_funcs.at(index)(tensor); + has_transform_func = true; + } + } + + if (!has_transform_func) { + arguments[index] = c10::IValue(tensor); + } + } + + SetIValueArguments(int_arguments, arguments); + SetIValueArguments(float_arguments, arguments); + SetIValueArguments(bool_arguments, arguments); + + torch::jit::Stack stack; + for (size_t i = 0; i < arguments.size(); i++) { + torch::jit::push(stack, arguments[i]); + } + + op->getOperation()(&stack); + // TODO: need to handle multiple-tensor outputs. + at::Tensor output; + torch::jit::pop(stack, output); + std::vector result; + result.emplace_back(at::toDLPack(output)); + return result; +} + +size_t execute_aten_operator_address() { return reinterpret_cast(&ExecuteATenOperator); } + """ + + aten_op_executor_cpp_extension = load_inline(name='aten_op_executor', cpp_sources=[aten_op_executor_cpp_source], + extra_cflags=['-D__HIP_PLATFORM_HCC__=1' if is_rocm_pytorch else ''], + functions=['execute_aten_operator_address'], + verbose=verbosity, with_cuda=True, + build_directory=TORCH_CPP_BUILD_DIR) + + C.register_aten_op_executor(str(aten_op_executor_cpp_extension.execute_aten_operator_address())) + +def _load_aten_op_executor_cpp_extension_if_needed(onnx_model, verbosity, is_rocm_pytorch): + for node in onnx_model.graph.node: + if node.op_type == 'ATenOp' and node.domain == 'com.microsoft': + _load_aten_op_executor_cpp_extension(verbosity, is_rocm_pytorch) + break diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index e7b7b60613..e726a62105 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -3,7 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from . import _utils, _io, _logger +from . import _utils, _io, _logger, _cpp_extensions as _cpp_ext from onnxruntime.training.ortmodule import ONNX_OPSET_VERSION from onnxruntime.capi import _pybind_state as C @@ -115,8 +115,8 @@ class GraphExecutionManager(ABC): self._use_external_gpu_allocator = True if self._use_external_gpu_allocator: # CPP extension to get torch GPU allocator's alloc and free function addresses - self._torch_gpu_allocator = _utils._load_torch_gpu_allocator_cpp_extension(self._loglevel < _logger.LogLevel.WARNING, - self.is_rocm_pytorch) + self._torch_gpu_allocator = _cpp_ext._load_torch_gpu_allocator_cpp_extension(self._loglevel < _logger.LogLevel.WARNING, + self.is_rocm_pytorch) self._torch_alloc = self._torch_gpu_allocator.gpu_caching_allocator_raw_alloc_address() self._torch_free = self._torch_gpu_allocator.gpu_caching_allocator_raw_delete_address() @@ -209,7 +209,7 @@ class GraphExecutionManager(ABC): self._set_device_from_module(inputs, kwargs) self._onnx_model = self._get_exported_model(*inputs, **kwargs) - _utils._load_aten_op_executor_cpp_extension_if_needed(self._onnx_model, self._loglevel < _logger.LogLevel.WARNING, self.is_rocm_pytorch) + _cpp_ext._load_aten_op_executor_cpp_extension_if_needed(self._onnx_model, self._loglevel < _logger.LogLevel.WARNING, self.is_rocm_pytorch) if self._save_onnx: onnx.save(self._onnx_model, self._save_onnx_prefix + '_torch_exporter.onnx') diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py index 67ae17f403..6c28942069 100644 --- a/orttraining/orttraining/python/training/ortmodule/_utils.py +++ b/orttraining/orttraining/python/training/ortmodule/_utils.py @@ -6,11 +6,8 @@ from onnxruntime.capi.onnxruntime_inference_collection import OrtValue from onnxruntime.capi import _pybind_state as C -import threading import torch from torch.utils.dlpack import from_dlpack, to_dlpack -from torch.utils.cpp_extension import load_inline -from functools import wraps def _ortvalue_to_torch_tensor(ortvalue): @@ -20,33 +17,10 @@ def _ortvalue_to_torch_tensor(ortvalue): torch_tensor = from_dlpack(ortvalue.to_dlpack()) return torch_tensor.to(torch.bool) if ortvalue.data_type() == 'tensor(bool)' else torch_tensor + def _ortvalue_from_torch_tensor(torch_tensor): return C.OrtValue.from_dlpack(to_dlpack(torch_tensor), torch_tensor.dtype == torch.bool) -def _load_torch_gpu_allocator_cpp_extension(verbosity, is_rocm_pytorch): - gpu_identifier = "hip" if is_rocm_pytorch else "cuda" - gpu_allocator_header = "HIPCachingAllocator" if is_rocm_pytorch else "CUDACachingAllocator" - torch_gpu_allocator_addresses_cpp_source = f''' - #include - #include - - size_t gpu_caching_allocator_raw_alloc_address() {{ - return reinterpret_cast(&c10::{gpu_identifier}::{gpu_allocator_header}::raw_alloc); - }} - - size_t gpu_caching_allocator_raw_delete_address() {{ - return reinterpret_cast(&c10::{gpu_identifier}::{gpu_allocator_header}::raw_delete); - }} - ''' - - return load_inline(name='inline_extension', - cpp_sources=[torch_gpu_allocator_addresses_cpp_source], - extra_cflags=['-D__HIP_PLATFORM_HCC__=1' if is_rocm_pytorch else ''], - functions=['gpu_caching_allocator_raw_alloc_address', - 'gpu_caching_allocator_raw_delete_address'], - verbose=verbosity, - with_cuda=True) - def _check_same_device(device, argument_str, *args): '''Check that all tensor arguments in *args reside on the same device as the input device''' @@ -118,140 +92,3 @@ def _create_iobinding(io_binding, inputs, model, device): for value_info in model.graph.output: io_binding.bind_output(value_info.name, device.type, device_id=get_device_index(device)) - -def run_once_aten_op_executor(f): - """ - Decorator to run a function only once. - :param f: function to be run only once during execution time despite the number of calls - :return: The original function with the params passed to it if it hasn't already been run before - """ - @wraps(f) - def aten_op_executor_wrapper(*args, **kwargs): - if not aten_op_executor_wrapper.has_run: - with aten_op_executor_wrapper.lock: - if not aten_op_executor_wrapper.has_run: - aten_op_executor_wrapper.has_run = True - return f(*args, **kwargs) - - aten_op_executor_wrapper.lock = threading.Lock() - aten_op_executor_wrapper.has_run = False - return aten_op_executor_wrapper - -@run_once_aten_op_executor -def _load_aten_op_executor_cpp_extension(verbosity, is_rocm_pytorch): - aten_op_executor_cpp_source = """ -#include -#include -#include -#include -#include - -class ATenOperatorCache { - public: - static ATenOperatorCache& Instance() { - static ATenOperatorCache instance; - return instance; - } - - std::shared_ptr GetOperator(const std::string& op_name) { - if (ops_.find(op_name) == ops_.end()) { - auto& ops = torch::jit::getAllOperatorsFor(torch::jit::Symbol::fromQualString(op_name)); - TORCH_INTERNAL_ASSERT(ops.size() == 1); - ops_[op_name] = ops.front(); - } - - return ops_.at(op_name); - } - - private: - ATenOperatorCache() = default; - std::unordered_map> ops_; -}; - -// Some arguments of backward operator are not from forward operator's input or output, -// but need some processing. Since we cannot build such processing to ONNX graph for now, -// we are putting such processing code here if needed. -// Take embedding_backward as example: -// weight: embedding_backward(grad, indices, weight.size(0), padding_idx, scale_grad_by_freq, sparse) -// the 3rd argument (index 2) is weight.size(0), we add this processing here. -using TensorTransformFunc = std::function; -static const TensorTransformFunc embedding_num_weights = [](const at::Tensor& tensor) { - return c10::IValue(tensor.size(0)); -}; - -static const std::unordered_map> TENSOR_TRANSFORM_FUNCS = { - {"aten::embedding_backward", {{2, embedding_num_weights}}}, -}; - -template -void SetIValueArguments(const std::vector>& raw_arguments, - std::vector& ivalue_arguments) { - for (size_t i = 0; i < raw_arguments.size(); i++) { - size_t index = std::get<0>(raw_arguments[i]); - TORCH_INTERNAL_ASSERT(index < ivalue_arguments.size()); - ivalue_arguments[index] = c10::IValue(std::get<1>(raw_arguments[i])); - } -} - -// TODO: Add more argument types, such as list type. -std::vector ExecuteATenOperator( - const char* op_name, const std::vector>& tensor_arguments, - const std::vector>& int_arguments, - const std::vector>& float_arguments, - const std::vector>& bool_arguments) { - std::string op_name_str(op_name); - std::shared_ptr op = ATenOperatorCache::Instance().GetOperator(op_name_str); - - // TODO: need to handle optional argument and arguments with default values. - std::vector arguments; - arguments.resize(op->schema().arguments().size()); - for (size_t i = 0; i < tensor_arguments.size(); i++) { - size_t index = std::get<0>(tensor_arguments[i]); - at::Tensor tensor = at::fromDLPack(std::get<1>(tensor_arguments[i])); - bool has_transform_func = false; - if (TENSOR_TRANSFORM_FUNCS.find(op_name_str) != TENSOR_TRANSFORM_FUNCS.end()) { - const auto& transform_funcs = TENSOR_TRANSFORM_FUNCS.at(op_name_str); - if (transform_funcs.find(index) != transform_funcs.end()) { - arguments[index] = transform_funcs.at(index)(tensor); - has_transform_func = true; - } - } - - if (!has_transform_func) { - arguments[index] = c10::IValue(tensor); - } - } - - SetIValueArguments(int_arguments, arguments); - SetIValueArguments(float_arguments, arguments); - SetIValueArguments(bool_arguments, arguments); - - torch::jit::Stack stack; - for (size_t i = 0; i < arguments.size(); i++) { - torch::jit::push(stack, arguments[i]); - } - - op->getOperation()(&stack); - // TODO: need to handle multiple-tensor outputs. - at::Tensor output; - torch::jit::pop(stack, output); - std::vector result; - result.emplace_back(at::toDLPack(output)); - return result; -} - -size_t execute_aten_operator_address() { return reinterpret_cast(&ExecuteATenOperator); } - """ - - aten_op_executor_cpp_extension = load_inline(name='inline_extension_aten_op_executor', cpp_sources=[aten_op_executor_cpp_source], - extra_cflags=['-D__HIP_PLATFORM_HCC__=1' if is_rocm_pytorch else ''], - functions=['execute_aten_operator_address'], - verbose=verbosity, with_cuda=True) - - C.register_aten_op_executor(str(aten_op_executor_cpp_extension.execute_aten_operator_address())) - -def _load_aten_op_executor_cpp_extension_if_needed(onnx_model, verbosity, is_rocm_pytorch): - for node in onnx_model.graph.node: - if node.op_type == 'ATenOp' and node.domain == 'com.microsoft': - _load_aten_op_executor_cpp_extension(verbosity, is_rocm_pytorch) - break