saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-06 04:28:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add cpp ext lock file check during ORTModule init (#7740)

Browse source 
* Add cpp ext lock file check during ORTModule init

* Address comments
This commit is contained in:
Thiago Crepaldi 2021-05-18 12:57:05 -07:00 committed by GitHub
parent 224a664811
commit e05b15175d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 205 additions and 171 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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

185
orttraining/orttraining/python/training/ortmodule/_cpp_extensions.py Normal file
View file

@ -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 <torch/extension.h>
#include <c10/{gpu_identifier}/{gpu_allocator_header}.h>
size_t gpu_caching_allocator_raw_alloc_address() {{
return reinterpret_cast<size_t>(&c10::{gpu_identifier}::{gpu_allocator_header}::raw_alloc);
}}
size_t gpu_caching_allocator_raw_delete_address() {{
return reinterpret_cast<size_t>(&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 <torch/torch.h>
#include <ATen/DLConvertor.h>
#include <unordered_map>
#include <tuple>
#include <vector>
class ATenOperatorCache {
public:
static ATenOperatorCache& Instance() {
static ATenOperatorCache instance;
return instance;
}
std::shared_ptr<torch::jit::Operator> 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<std::string, std::shared_ptr<torch::jit::Operator>> 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<c10::IValue(const at::Tensor&)>;
static const TensorTransformFunc embedding_num_weights = [](const at::Tensor& tensor) {
return c10::IValue(tensor.size(0));
};
static const std::unordered_map<std::string, std::unordered_map<size_t, TensorTransformFunc>> TENSOR_TRANSFORM_FUNCS = {
{"aten::embedding_backward", {{2, embedding_num_weights}}},
};
template <typename T>
void SetIValueArguments(const std::vector<std::tuple<size_t, T>>& raw_arguments,
std::vector<c10::IValue>& 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<DLManagedTensor*> ExecuteATenOperator(
const char* op_name, const std::vector<std::tuple<size_t, DLManagedTensor*>>& tensor_arguments,
const std::vector<std::tuple<size_t, int64_t>>& int_arguments,
const std::vector<std::tuple<size_t, float>>& float_arguments,
const std::vector<std::tuple<size_t, bool>>& bool_arguments) {
std::string op_name_str(op_name);
std::shared_ptr<torch::jit::Operator> op = ATenOperatorCache::Instance().GetOperator(op_name_str);
// TODO: need to handle optional argument and arguments with default values.
std::vector<c10::IValue> 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<int64_t>(int_arguments, arguments);
SetIValueArguments<float>(float_arguments, arguments);
SetIValueArguments<bool>(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<DLManagedTensor*> result;
result.emplace_back(at::toDLPack(output));
return result;
}
size_t execute_aten_operator_address() { return reinterpret_cast<size_t>(&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

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

@ -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')

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

@ -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 <torch/extension.h>
#include <c10/{gpu_identifier}/{gpu_allocator_header}.h>
size_t gpu_caching_allocator_raw_alloc_address() {{
return reinterpret_cast<size_t>(&c10::{gpu_identifier}::{gpu_allocator_header}::raw_alloc);
}}
size_t gpu_caching_allocator_raw_delete_address() {{
return reinterpret_cast<size_t>(&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 <torch/torch.h>
#include <ATen/DLConvertor.h>
#include <unordered_map>
#include <tuple>
#include <vector>
class ATenOperatorCache {
public:
static ATenOperatorCache& Instance() {
static ATenOperatorCache instance;
return instance;
}
std::shared_ptr<torch::jit::Operator> 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<std::string, std::shared_ptr<torch::jit::Operator>> 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<c10::IValue(const at::Tensor&)>;
static const TensorTransformFunc embedding_num_weights = [](const at::Tensor& tensor) {
return c10::IValue(tensor.size(0));
};
static const std::unordered_map<std::string, std::unordered_map<size_t, TensorTransformFunc>> TENSOR_TRANSFORM_FUNCS = {
{"aten::embedding_backward", {{2, embedding_num_weights}}},
};
template <typename T>
void SetIValueArguments(const std::vector<std::tuple<size_t, T>>& raw_arguments,
std::vector<c10::IValue>& 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<DLManagedTensor*> ExecuteATenOperator(
const char* op_name, const std::vector<std::tuple<size_t, DLManagedTensor*>>& tensor_arguments,
const std::vector<std::tuple<size_t, int64_t>>& int_arguments,
const std::vector<std::tuple<size_t, float>>& float_arguments,
const std::vector<std::tuple<size_t, bool>>& bool_arguments) {
std::string op_name_str(op_name);
std::shared_ptr<torch::jit::Operator> op = ATenOperatorCache::Instance().GetOperator(op_name_str);
// TODO: need to handle optional argument and arguments with default values.
std::vector<c10::IValue> 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<int64_t>(int_arguments, arguments);
SetIValueArguments<float>(float_arguments, arguments);
SetIValueArguments<bool>(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<DLManagedTensor*> result;
result.emplace_back(at::toDLPack(output));
return result;
}
size_t execute_aten_operator_address() { return reinterpret_cast<size_t>(&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